Merge branch 'main' into feat/go-ci-visibility-failnow-teardown

tonyredondo · web-flow · commit 2c6f3bb080f8 · 2026-05-05T15:07:12.000+02:00
diff --git a/contrib/mark3labs/mcp-go/go.mod b/contrib/mark3labs/mcp-go/go.mod
@@ -31,7 +31,7 @@ require (
 	github.com/DataDog/sketches-go v1.4.8 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
-	github.com/buger/jsonparser v1.1.1 // indirect
+	github.com/buger/jsonparser v1.1.2 // indirect
 	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cihub/seelog v0.0.0-20170130134532-f561c5e57575 // indirect
diff --git a/contrib/mark3labs/mcp-go/go.sum b/contrib/mark3labs/mcp-go/go.sum
diff --git a/ddtrace/tracer/textmap.go b/ddtrace/tracer/textmap.go
@@ -1441,27 +1441,34 @@ func (*propagatorBaggage) extractTextMap(reader TextMapReader) (*SpanContext, er
 		return &ctx, nil
 	}
 
-	parts := strings.Split(baggageHeader, ",")
-
-	// 1) validation & single-trim pass
-	for i, kv := range parts {
+	// Single pass: enforce baggageMaxItems and baggageMaxBytes, validate, and apply.
+	ctr := 0
+	byteCount := 0
+	for kv := range strings.SplitSeq(baggageHeader, ",") {
+		itemBytes := len(kv)
+		if ctr > 0 {
+			itemBytes++ // comma separator
+		}
+		if ctr >= baggageMaxItems {
+			log.Warn("baggage item count exceeded limit (%d), dropping remaining items", baggageMaxItems)
+			break
+		}
+		if byteCount+itemBytes > baggageMaxBytes {
+			log.Warn("baggage byte limit exceeded (%d), dropping remaining items", baggageMaxBytes)
+			break
+		}
 		k, v, ok := strings.Cut(kv, "=")
 		trimmedK := strings.TrimSpace(k)
 		trimmedV := strings.TrimSpace(v)
 		if !ok || trimmedK == "" || trimmedV == "" {
 			log.Warn("invalid baggage item: %q, dropping entire header", kv)
-			return &ctx, nil
+			return &SpanContext{}, nil
 		}
-		// store back the trimmed pair so we don't re-trim below
-		parts[i] = trimmedK + "=" + trimmedV
-	}
-
-	// 2) safe to URL-decode & apply
-	for _, kv := range parts {
-		rawK, rawV, _ := strings.Cut(kv, "=")
-		key, _ := url.QueryUnescape(rawK)
-		val, _ := url.QueryUnescape(rawV)
+		key, _ := url.QueryUnescape(trimmedK)
+		val, _ := url.QueryUnescape(trimmedV)
 		ctx.setBaggageItem(key, val)
+		byteCount += itemBytes
+		ctr++
 	}
 
 	return &ctx, nil
diff --git a/ddtrace/tracer/textmap_test.go b/ddtrace/tracer/textmap_test.go
@@ -2892,6 +2892,127 @@ func TestExtractBaggagePropagatorMalformedHeader(t *testing.T) {
 	})
 }
 
+func TestExtractBaggagePropagatorMaxItems(t *testing.T) {
+	tracer, err := newTracer()
+	assert.NoError(t, err)
+	defer tracer.Stop()
+
+	var b strings.Builder
+	for i := range baggageMaxItems + 5 {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		iStr := strconv.Itoa(i)
+		b.WriteString("key" + iStr + "=val" + iStr)
+	}
+
+	headers := TextMapCarrier{
+		DefaultTraceIDHeader:  "4",
+		DefaultParentIDHeader: "1",
+		DefaultBaggageHeader:  b.String(),
+	}
+	s, err := tracer.Extract(headers)
+	assert.NoError(t, err)
+
+	got := make(map[string]string)
+	s.ForeachBaggageItem(func(k, v string) bool {
+		got[k] = v
+		return true
+	})
+	assert.Len(t, got, baggageMaxItems)
+	for i := range baggageMaxItems {
+		iStr := strconv.Itoa(i)
+		assert.Equal(t, "val"+iStr, got["key"+iStr])
+	}
+	for i := baggageMaxItems; i < baggageMaxItems+5; i++ {
+		iStr := strconv.Itoa(i)
+		_, present := got["key"+iStr]
+		assert.False(t, present, "key%s should not be present", iStr)
+	}
+}
+
+func TestExtractBaggagePropagatorMaxBytes(t *testing.T) {
+	tracer, err := newTracer()
+	assert.NoError(t, err)
+	defer tracer.Stop()
+
+	// 12 items, each "keyN=" + 1000 'a's = 1005 wire bytes. Including comma
+	// separators, the first 8 fit under baggageMaxBytes (8192); the 9th would
+	// push the running total to 9053 > 8192, so items 8..11 are dropped.
+	const itemValLen = 1000
+	const numItems = 12
+	const expectedKept = 8
+	val := strings.Repeat("a", itemValLen)
+
+	var b strings.Builder
+	for i := range numItems {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		b.WriteString("key" + strconv.Itoa(i) + "=" + val)
+	}
+
+	headers := TextMapCarrier{
+		DefaultTraceIDHeader:  "4",
+		DefaultParentIDHeader: "1",
+		DefaultBaggageHeader:  b.String(),
+	}
+	s, err := tracer.Extract(headers)
+	assert.NoError(t, err)
+
+	got := make(map[string]string)
+	s.ForeachBaggageItem(func(k, v string) bool {
+		got[k] = v
+		return true
+	})
+	assert.Len(t, got, expectedKept)
+	for i := range expectedKept {
+		assert.Equal(t, val, got["key"+strconv.Itoa(i)])
+	}
+	for i := expectedKept; i < numItems; i++ {
+		_, present := got["key"+strconv.Itoa(i)]
+		assert.False(t, present, "key%d should not be present", i)
+	}
+}
+
+func TestExtractBaggagePropagatorMalformedPastLimit(t *testing.T) {
+	tracer, err := newTracer()
+	assert.NoError(t, err)
+	defer tracer.Stop()
+
+	// baggageMaxItems valid entries followed by a malformed entry. Because
+	// the malformed entry sits past the items limit it is never inspected,
+	// so the valid prefix is kept (regression check on the single-pass design).
+	var b strings.Builder
+	for i := range baggageMaxItems {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		iStr := strconv.Itoa(i)
+		b.WriteString("key" + iStr + "=val" + iStr)
+	}
+	b.WriteString(",malformed_no_equals_sign")
+
+	headers := TextMapCarrier{
+		DefaultTraceIDHeader:  "4",
+		DefaultParentIDHeader: "1",
+		DefaultBaggageHeader:  b.String(),
+	}
+	s, err := tracer.Extract(headers)
+	assert.NoError(t, err)
+
+	got := make(map[string]string)
+	s.ForeachBaggageItem(func(k, v string) bool {
+		got[k] = v
+		return true
+	})
+	assert.Len(t, got, baggageMaxItems)
+	for i := range baggageMaxItems {
+		iStr := strconv.Itoa(i)
+		assert.Equal(t, "val"+iStr, got["key"+iStr])
+	}
+}
+
 func TestExtractOnlyBaggage(t *testing.T) {
 	t.Setenv("DD_TRACE_PROPAGATION_STYLE", "baggage")
 	headers := TextMapCarrier(map[string]string{
diff --git a/internal/llmobs/llmobs.go b/internal/llmobs/llmobs.go
@@ -146,9 +146,10 @@ type LLMObs struct {
 	// lifecycle
 	mu            sync.Mutex
 	running       bool
-	wg            sync.WaitGroup
-	stopCh        chan struct{} // signal stop
-	flushNowCh    chan struct{}
+	sendWg        sync.WaitGroup // tracks in-flight batchSend goroutines
+	workerDone    chan struct{}  // closed when the worker loop exits
+	stopCh        chan struct{}  // signal stop
+	flushNowCh    chan chan struct{}
 	flushInterval time.Duration
 }
 
@@ -187,8 +188,9 @@ func newLLMObs(cfg *config.Config, tracer Tracer) (*LLMObs, error) {
 		Tracer:        tracer,
 		spanEventsCh:  make(chan *transport.LLMObsSpanEvent),
 		evalMetricsCh: make(chan *transport.LLMObsMetric),
+		workerDone:    make(chan struct{}),
 		stopCh:        make(chan struct{}),
-		flushNowCh:    make(chan struct{}, 1),
+		flushNowCh:    make(chan chan struct{}, 1),
 		flushInterval: defaultFlushInterval,
 	}, nil
 }
@@ -245,6 +247,13 @@ func Flush() {
 	}
 }
 
+// FlushSync forces a flush of all buffered LLMObs data and blocks until the flush completes.
+func FlushSync() {
+	if activeLLMObs != nil {
+		activeLLMObs.FlushSync()
+	}
+}
+
 // Run starts the worker loop that processes span events and metrics.
 func (l *LLMObs) Run() {
 	l.mu.Lock()
@@ -255,7 +264,8 @@ func (l *LLMObs) Run() {
 	l.running = true
 	l.mu.Unlock()
 
-	l.wg.Go(func() {
+	go func() {
+		defer close(l.workerDone)
 		// this goroutine should be the only one writing to the internal buffers
 
 		ticker := time.NewTicker(l.flushInterval)
@@ -268,9 +278,7 @@ func (l *LLMObs) Run() {
 				if l.bufSpanEventsSize+evSize > sizeLimitEVPEvent {
 					log.Debug("llmobs: span events buffer size limit reached, flushing before adding new event")
 					params := l.clearBuffersNonLocked()
-					l.wg.Go(func() {
-						l.batchSend(params)
-					})
+					l.sendWg.Go(func() { l.batchSend(params) })
 				}
 				l.bufSpanEvents = append(l.bufSpanEvents, ev)
 				l.bufSpanEventsSize += evSize
@@ -280,16 +288,22 @@ func (l *LLMObs) Run() {
 
 			case <-ticker.C:
 				params := l.clearBuffersNonLocked()
-				l.wg.Go(func() {
-					l.batchSend(params)
-				})
+				l.sendWg.Go(func() { l.batchSend(params) })
 
-			case <-l.flushNowCh:
+			case done := <-l.flushNowCh:
 				log.Debug("llmobs: on-demand flush signal")
 				params := l.clearBuffersNonLocked()
-				l.wg.Go(func() {
+				l.sendWg.Add(1)
+				go func() {
+					defer func() {
+						l.sendWg.Done()
+						if done != nil {
+							l.sendWg.Wait()
+							close(done)
+						}
+					}()
 					l.batchSend(params)
-				})
+				}()
 
 			case <-l.stopCh:
 				log.Debug("llmobs: stop signal")
@@ -299,7 +313,7 @@ func (l *LLMObs) Run() {
 				return
 			}
 		}
-	})
+	}()
 }
 
 // clearBuffersNonLocked clears the internal buffers and returns the corresponding batchSendParams to send to the backend.
@@ -320,11 +334,24 @@ func (l *LLMObs) clearBuffersNonLocked() batchSendParams {
 func (l *LLMObs) Flush() {
 	// non-blocking edge trigger so multiple calls coalesce
 	select {
-	case l.flushNowCh <- struct{}{}:
+	case l.flushNowCh <- nil:
 	default:
 	}
 }
 
+// FlushSync forces an immediate flush and blocks until the flush completes.
+func (l *LLMObs) FlushSync() {
+	done := make(chan struct{})
+	select {
+	case l.flushNowCh <- done:
+		select {
+		case <-done:
+		case <-l.stopCh:
+		}
+	case <-l.stopCh:
+	}
+}
+
 // Stop requests shutdown, drains what’s already in the channels, flushes, and waits.
 func (l *LLMObs) Stop() {
 	l.mu.Lock()
@@ -342,8 +369,10 @@ func (l *LLMObs) Stop() {
 		close(l.stopCh)
 	}
 
-	// Wait for the main worker to exit (it will do a final flush)
-	l.wg.Wait()
+	// Wait for the worker loop to exit (it does a final synchronous flush),
+	// then wait for any async batchSend goroutines still in flight.
+	<-l.workerDone
+	l.sendWg.Wait()
 }
 
 // drainChannels pulls everything currently buffered in the channels into our in-memory buffers.
diff --git a/internal/llmobs/llmobs_test.go b/internal/llmobs/llmobs_test.go
