feat(llmobs): add internal FlushSync function (#4695)

rarguelloF · web-flow · commit e2961fddeddc · 2026-05-05T14:30:53.000+02:00
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
@@ -2335,3 +2335,89 @@ func TestSpanEventsSizeBasedFlushing(t *testing.T) {
 				"all spans accumulate in a single batch that is too large to send", size)
 	}
 }
+
+func TestFlushSync(t *testing.T) {
+	t.Run("does-not-hang-with-empty-buffer", func(t *testing.T) {
+		// FlushSync must return promptly even when there is nothing to flush.
+		_, ll := testTracer(t)
+
+		done := make(chan struct{})
+		go func() {
+			ll.FlushSync()
+			close(done)
+		}()
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			t.Fatal("FlushSync hung with empty buffer")
+		}
+	})
+	t.Run("waits-for-slow-transport", func(t *testing.T) {
+		// FlushSync must block the caller until the HTTP send completes.
+		// Verify by timing: if FlushSync returned before the request delay,
+		// the blocking guarantee would be violated.
+		const delay = 200 * time.Millisecond
+
+		tt, ll := testTracer(t, testtracer.WithRequestDelay(delay))
+
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "slow-span", llmobs.StartSpanConfig{})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		start := time.Now()
+		ll.FlushSync()
+		elapsed := time.Since(start)
+
+		// Must have waited at least the request delay.
+		assert.GreaterOrEqual(t, elapsed, delay, "FlushSync should block until the slow HTTP send completes")
+
+		spans := tt.WaitForLLMObsSpans(t, 1)
+		require.Len(t, spans, 1)
+		assert.Equal(t, "slow-span", spans[0].Name)
+	})
+	t.Run("no-panic-without-active-llmobs", func(t *testing.T) {
+		// Package-level FlushSync should be safe when no active LLMObs.
+		assert.NotPanics(t, func() {
+			llmobs.FlushSync()
+		})
+	})
+	t.Run("waits-for-data-queued-by-flush", func(t *testing.T) {
+		// Flush() triggers batchSend asynchronously. If FlushSync() is called
+		// right after, it should still block until that batchSend completes.
+		// Without a fix, FlushSync gets an already-empty buffer (Flush cleared
+		// it) and returns immediately while the in-flight batchSend is still running.
+		const delay = 200 * time.Millisecond
+		_, ll := testTracer(t, testtracer.WithRequestDelay(delay))
+
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "test-span", llmobs.StartSpanConfig{})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		ll.Flush()
+		start := time.Now()
+		ll.FlushSync()
+		elapsed := time.Since(start)
+
+		assert.GreaterOrEqual(t, elapsed, delay-20*time.Millisecond,
+			"FlushSync should block until the batchSend triggered by the preceding Flush completes")
+	})
+	t.Run("does-not-hang-after-stop", func(t *testing.T) {
+		_, ll := testTracer(t)
+		ll.Stop()
+
+		// After Stop, stopCh is closed and the worker is gone. FlushSync's select
+		// has two ready cases: <-stopCh and flushNowCh<-done (buffered, empty).
+		// Go picks randomly, so loop to reliably hit the flushNowCh branch —
+		// without the fix that branch blocks forever on <-done.
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			for range 20 {
+				ll.FlushSync()
+			}
+		}()
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			t.Fatal("FlushSync hung after Stop")
+		}
+	})
+}
