perf improvements

Author: angus-langchain

internal/aggregator/aggregator.go

@@ -97,18 +97,22 @@ func (a *Aggregator) worker(ctx context.Context, ch <-chan *model.Run) {
 		if len(comp) > 0 {
 			go a.up.Send(ctx, uploader.Batch{Data: comp, Boundary: boundary})
 		}
-		sc = serializer.NewStreamingCompressor()
 	}
 
 	add := func(r *model.Run) {
+		needFlush := false
+
 		scMu.Lock()
-		defer scMu.Unlock()
 		if err := sc.AddRun(r); err != nil {
-			slog.Error("stream add", "err", err)
+			scMu.Unlock()
 			return
 		}
-		if sc.RunCount() >= a.cfg.BatchSize || sc.Uncompressed() >= a.cfg.MaxBufferBytes {
-			go flush()
+		needFlush = sc.RunCount() >= a.cfg.BatchSize ||
+			sc.Uncompressed() >= a.cfg.MaxBufferBytes
+		scMu.Unlock()
+
+		if needFlush {
+			flush()
 		}
 	}
 

internal/serializer/stream.go

@@ -13,9 +13,9 @@ import (
 
 // StreamingCompressor uses zstd + multipart/form-data to format and compress traces.
 //
-// Treat StreamingCompressor as a short-lived object for a single upload. Call Close()
-// when done to get the final bytes and cleanup resources. Then create a new StreamingCompressor
-// for the next upload.
+// Treat StreamingCompressor as a long-lived object for a multiple uploads. Call Close()
+// when done with a single upload to get the final bytes and cleanup resources. You should
+// reuse the StreamingCompressor object for multiple uploads.
 //
 // When interacting with the underlying buffer, the StreamingCompressor is not thread-safe.
 // It is the responsibility of the caller to ensure thread-safety by using a mutex.
@@ -89,7 +89,13 @@ func (sc *StreamingCompressor) Close() ([]byte, string, int, error) {
 	if err := sc.w.Close(); err != nil {
 		return nil, "", 0, err
 	}
-	return sc.buf.Bytes(), sc.boundary, sc.uncompressed, nil
+	out := sc.buf.Bytes()
+	uncompressed := sc.uncompressed
+	// reset the buffer and writer for the next upload
+	sc.buf = bytes.NewBuffer(nil)
+	sc.w = zstd.NewWriter(sc.buf)
+	sc.runCount, sc.uncompressed = 0, 0
+	return out, sc.boundary, uncompressed, nil
 }
 
 func (sc *StreamingCompressor) Uncompressed() int {

internal/translator/translator.go

@@ -15,6 +15,7 @@ type spanEntry struct {
 }
 
 type Translator struct {
+	converter  *GenAiConverter
 	mu         sync.RWMutex
 	span2trace map[string]spanEntry
 	ttl        time.Duration
@@ -50,7 +51,6 @@ func (t *Translator) Translate(req *collectortracepb.ExportTraceServiceRequest)
 	}
 
 	runs := make([]*model.Run, 0, total)
-	converter := &GenAiConverter{}
 	for _, rs := range req.ResourceSpans {
 		for _, ss := range rs.ScopeSpans {
 			for _, span := range ss.Spans {
@@ -63,7 +63,7 @@ func (t *Translator) Translate(req *collectortracepb.ExportTraceServiceRequest)
 						}
 					}
 				}
-				run, err := converter.ConvertSpan(span, false)
+				run, err := t.converter.ConvertSpan(span, false)
 				if err != nil || run == nil {
 					continue
 				}