use compressed buffer

Author: angus-langchain

cmd/collector/main.go

@@ -11,16 +11,16 @@ import (
 
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/aggregator"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/config"
-	"github.com/langchain-ai/langsmith-collector-proxy/internal/queue"
+	"github.com/langchain-ai/langsmith-collector-proxy/internal/model"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/server"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/uploader"
 )
 
 func main() {
 	cfg := config.Load()
 
-	qm := queue.New(cfg.QueueSize)
-	trSrv := server.NewRouter(cfg, qm)
+	ch := make(chan *model.Run, 1024)
+	trSrv := server.NewRouter(cfg, ch)
 
 	up := uploader.New(uploader.Config{
 		BaseURL:        cfg.LangsmithHost,
@@ -30,13 +30,12 @@ func main() {
 		InFlight:       100,
 		APIKey:         cfg.DefaultAPIKey,
 	})
-	up.Start()
-	defer up.Stop()
 
-	agg := aggregator.New(qm, up, aggregator.Config{
-		BatchSize:     cfg.BatchSize,
-		FlushInterval: cfg.FlushInterval,
-	})
+	agg := aggregator.New(up, aggregator.Config{
+		BatchSize:      cfg.BatchSize,
+		FlushInterval:  cfg.FlushInterval,
+		MaxBufferBytes: cfg.MaxBufferBytes,
+	}, ch)
 	agg.Start()
 	defer agg.Stop()
 

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/go-chi/chi/v5 v5.2.1
 	github.com/google/uuid v1.6.0
 	go.opentelemetry.io/proto/otlp v1.6.0
+	golang.org/x/sync v0.14.0
 	google.golang.org/protobuf v1.36.6
 )
 

go.sum

@@ -30,6 +30,8 @@ go.opentelemetry.io/proto/otlp v1.6.0 h1:jQjP+AQyTf+Fe7OKj/MfkDrmK4MNVtw2NpXsf9f
 go.opentelemetry.io/proto/otlp v1.6.0/go.mod h1:cicgGehlFuNdgZkcALOCh3VE6K/u2tAjzlRhDwmVpZc=
 golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY=
 golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E=
+golang.org/x/sync v0.14.0 h1:woo0S4Yywslg6hp4eUFjTVOyKt0RookbpAHG4c1HmhQ=
+golang.org/x/sync v0.14.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
 golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
 golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0=

internal/aggregator/aggregator.go

@@ -7,59 +7,44 @@ import (
 	"time"
 
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/model"
-	"github.com/langchain-ai/langsmith-collector-proxy/internal/queue"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/serializer"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/uploader"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/util"
 )
 
 type Config struct {
-	BatchSize     int
-	FlushInterval time.Duration
-
-	GCInterval time.Duration
-	EntryTTL   time.Duration
+	BatchSize      int
+	FlushInterval  time.Duration
+	MaxBufferBytes int
+	GCInterval     time.Duration
+	EntryTTL       time.Duration
 }
 
 type Aggregator struct {
-	qm     *queue.QueueManager
+	ch     chan *model.Run
 	cfg    Config
 	up     *uploader.Uploader
 	cancel context.CancelFunc
 }
 
-func New(qm *queue.QueueManager, up *uploader.Uploader, cfg Config) *Aggregator {
+func New(up *uploader.Uploader, cfg Config, ch chan *model.Run) *Aggregator {
 	if cfg.GCInterval == 0 {
 		cfg.GCInterval = 2 * time.Minute
 	}
 	if cfg.EntryTTL == 0 {
 		cfg.EntryTTL = 5 * time.Minute
 	}
-	return &Aggregator{qm: qm, cfg: cfg, up: up}
+	if cfg.MaxBufferBytes == 0 {
+		cfg.MaxBufferBytes = 10 * 1024 * 1024 // 10MB
+	}
+	return &Aggregator{up: up, cfg: cfg, ch: ch}
 }
 
 func (a *Aggregator) Start() {
 	ctx, cancel := context.WithCancel(context.Background())
 	a.cancel = cancel
 
-	started := map[string]struct{}{}
-	go func() {
-		tick := time.NewTicker(1 * time.Second)
-		defer tick.Stop()
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-tick.C:
-				for project, ch := range a.qm.All() {
-					if _, ok := started[project]; !ok {
-						started[project] = struct{}{}
-						go a.worker(ctx, project, ch)
-					}
-				}
-			}
-		}
-	}()
+	go a.worker(ctx, a.ch)
 }
 
 func (a *Aggregator) Stop() {
@@ -68,10 +53,9 @@ func (a *Aggregator) Stop() {
 	}
 }
 
-func (a *Aggregator) worker(ctx context.Context, project string, ch <-chan *model.Run) {
-	batch := make([]*model.Run, 0, a.cfg.BatchSize)
-	flushTicker := time.NewTicker(a.cfg.FlushInterval)
-	defer flushTicker.Stop()
+func (a *Aggregator) worker(ctx context.Context, ch <-chan *model.Run) {
+	sc := serializer.NewStreamingCompressor()
+	var scMu sync.Mutex
 
 	type entry struct {
 		dotted string
@@ -82,28 +66,36 @@ func (a *Aggregator) worker(ctx context.Context, project string, ch <-chan *mode
 	waitingChildren := make(map[string][]*model.Run) // parentID → parked kids
 	waitingSince := make(map[string]time.Time)       // parentID → first‑park time
 
-	var mu sync.Mutex
-
 	flush := func() {
-		if len(batch) == 0 {
+		scMu.Lock()
+		defer scMu.Unlock()
+		if sc == nil || sc.RunCount() == 0 {
 			return
 		}
-		comp, boundary, _, err := serializer.SerialiseAndCompress(batch)
+		comp, boundary, _, err := sc.Close()
 		if err != nil {
 			slog.Error("serialize", "err", err)
 			return
 		}
-		a.up.Enqueue(uploader.Batch{Project: project, Data: comp, Boundary: boundary})
-		batch = batch[:0]
+		go a.up.Send(ctx, uploader.Batch{Data: comp, Boundary: boundary})
+		sc = serializer.NewStreamingCompressor()
 	}
 
 	add := func(r *model.Run) {
-		batch = append(batch, r)
-		if len(batch) >= a.cfg.BatchSize {
-			flush()
+		scMu.Lock()
+		defer scMu.Unlock()
+		if err := sc.AddRun(r); err != nil {
+			slog.Error("stream add", "err", err)
+			return
+		}
+		if sc.RunCount() >= a.cfg.BatchSize || sc.Uncompressed() >= a.cfg.MaxBufferBytes {
+			go flush()
 		}
 	}
 
+	flushTicker := time.NewTicker(a.cfg.FlushInterval)
+	defer flushTicker.Stop()
+
 	var cascade func(parentID, parentDotted string)
 	cascade = func(parentID, parentDotted string) {
 		kids := waitingChildren[parentID]
@@ -157,16 +149,13 @@ func (a *Aggregator) worker(ctx context.Context, project string, ch <-chan *mode
 	for {
 		select {
 		case <-ctx.Done():
-			mu.Lock()
 			flush()
-			mu.Unlock()
 			return
 
 		case r := <-ch:
 			if r == nil {
 				continue
 			}
-			mu.Lock()
 			switch {
 			case r.ParentRunID == nil || *r.ParentRunID == "":
 				d := util.NewDottedOrder(*r.ID)
@@ -190,17 +179,12 @@ func (a *Aggregator) worker(ctx context.Context, project string, ch <-chan *mode
 					}
 				}
 			}
-			mu.Unlock()
 
 		case <-flushTicker.C:
-			mu.Lock()
 			flush()
-			mu.Unlock()
 
 		case <-gcTicker.C:
-			mu.Lock()
 			gc()
-			mu.Unlock()
 		}
 	}
 }

internal/config/config.go

@@ -9,12 +9,12 @@ import (
 type Config struct {
 	Port           string
 	MaxBodyBytes   int64
-	QueueSize      int
 	LangsmithHost  string
 	DefaultAPIKey  string
 	DefaultProject string
 	BatchSize      int
 	FlushInterval  time.Duration
+	MaxBufferBytes int
 	MaxRetries     int
 	BackoffInitial time.Duration
 }
@@ -44,10 +44,16 @@ func Load() Config {
 		LangsmithHost:  env("LANGSMITH_ENDPOINT", "https://api.smith.langchain.com"),
 		DefaultAPIKey:  env("LANGSMITH_API_KEY", ""),
 		DefaultProject: env("LANGSMITH_PROJECT", ""),
-		// Collector Config
-		QueueSize:     int(envInt64("QUEUE_SIZE", 1000)),
-		BatchSize:     int(envInt64("BATCH_SIZE", 100)), // default 100 runs
-		FlushInterval: time.Duration(envInt64("FLUSH_INTERVAL_MS", 5000)) * time.Millisecond,
+		// Collector Config. The following values control how frequently and how many runs are sent to LangSmith.
+		// The collector will buffer zstd compressed runs in memory until one of the following conditions is met:
+		// 1. The buffer size exceeds MaxBufferBytes.
+		// 2. The flush interval is reached.
+		// 3. The batch size is reached.
+		// Because zstd compression can be extremely effective, we use the uncompressed size of the buffer to determine how much data
+		// will be sent to LangSmith.
+		BatchSize:      int(envInt64("BATCH_SIZE", 100)), // default 100 runs
+		FlushInterval:  time.Duration(envInt64("FLUSH_INTERVAL_MS", 5000)) * time.Millisecond,
+		MaxBufferBytes: int(envInt64("MAX_BUFFER_BYTES", 10*1024*1024)), // default 10MB
 		// Uploader Config
 		MaxRetries:     int(envInt64("MAX_RETRIES", 3)),
 		BackoffInitial: time.Duration(envInt64("RETRY_BACKOFF_MS", 100)) * time.Millisecond,

internal/handler/traces.go

@@ -9,8 +9,7 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/langchain-ai/langsmith-collector-proxy/internal/contextkey"
-	"github.com/langchain-ai/langsmith-collector-proxy/internal/queue"
+	"github.com/langchain-ai/langsmith-collector-proxy/internal/model"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/translator"
 	collectortracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1"
 	"google.golang.org/protobuf/encoding/protojson"
@@ -41,14 +40,14 @@ func handleError(w http.ResponseWriter, r *http.Request, status int, err error)
 }
 
 // TracesHandler ingests OTLP/HTTP requests.
-func TracesHandler(maxBody int64, tr *translator.Translator, qm *queue.QueueManager) http.HandlerFunc {
+func TracesHandler(maxBody int64, tr *translator.Translator, ch chan *model.Run) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		r.Body = http.MaxBytesReader(w, r.Body, maxBody)
 		defer r.Body.Close()
 		var reader io.Reader = r.Body
 		switch strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Encoding"))) {
 		case "", "identity":
-		case "gzip":
+		case "gzip": // todo add comment about this ( add comments in genral)
 			gzr, err := gzip.NewReader(reader)
 			if err != nil {
 				handleError(w, r, http.StatusBadRequest, err)
@@ -82,24 +81,12 @@ func TracesHandler(maxBody int64, tr *translator.Translator, qm *queue.QueueMana
 
 		runs := tr.Translate(&req)
 
-		project := r.Header.Get("Langsmith-Project")
-		if project == "" {
-			if v := r.Context().Value(contextkey.ProjectIDKey); v != nil {
-				project, _ = v.(string)
-			}
-		}
+		project := r.Header.Get("Langsmith-Project") // todo: remove queues by project
 		for _, run := range runs {
 			if project != "" {
 				run.SessionName = &project
 			}
-			if err := qm.Enqueue(r.Context(), project, run); err != nil {
-				if errors.Is(err, queue.ErrContextDone) {
-					handleError(w, r, 499, err)
-				} else {
-					handleError(w, r, http.StatusServiceUnavailable, err)
-				}
-				return
-			}
+			ch <- run
 		}
 		w.WriteHeader(http.StatusAccepted)
 	}