fix: metrics flush timeout canceled shutdown

router/core/graph_server.go

@@ -151,7 +151,7 @@ type reusedGraphMux struct {
 }
 
 // newGraphServer creates a new server instance.
-func newGraphServer(routerCtx context.Context, r *Router, response *routerconfig.Response, proxy ProxyFunc) (*graphServer, error) {
+func newGraphServer(routerCtx context.Context, r *Router, response *routerconfig.Response, proxy ProxyFunc) (_ *graphServer, resErr error) {
 	/* Older versions of composition will not populate a compatibility version.
 	 * Currently, all "old" router execution configurations are compatible as there have been no breaking
 	 * changes.
@@ -231,6 +231,16 @@ func newGraphServer(routerCtx context.Context, r *Router, response *routerconfig
 		headerPropagation: r.headerPropagation,
 	}
 
+	defer func() {
+		if resErr == nil {
+			return
+		}
+		// Shutdown the graph server to clean up resources
+		if err := s.Shutdown(routerCtx); err != nil {
+			s.logger.Error("Failed to shut down graph server during failed-build cleanup", zap.Error(err))
+		}
+	}()
+
 	baseOtelAttributes := []attribute.KeyValue{
 		otel.WgRouterVersion.String(Version),
 		otel.WgRouterClusterName.String(r.clusterName),
@@ -678,8 +688,9 @@ type graphMux struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	mux    *chi.Mux
-	reused atomic.Bool
+	mux       *chi.Mux
+	reused    atomic.Bool
+	finalized atomic.Bool
 
 	planCache                   *ristretto.Cache[uint64, *planWithMetaData]
 	planFallbackCache           *slowplancache.Cache[*planWithMetaData]
@@ -958,6 +969,11 @@ func (s *graphMux) configureCacheMetrics(srv *graphServer, baseOtelAttributes []
 }
 
 func (s *graphMux) Shutdown(ctx context.Context) error {
+	// Make sure we do not shutdown the mux multiple times
+	if !s.finalized.CompareAndSwap(false, true) {
+		return nil
+	}
+
 	// cancel the graph muxes context to close its resources like websocket connections, resolvers, etc.
 	s.cancel()
 
@@ -997,12 +1013,6 @@ func (s *graphMux) Shutdown(ctx context.Context) error {
 		}
 	}
 
-	if s.streamMetricStore != nil {
-		if aErr := s.streamMetricStore.Shutdown(ctx); aErr != nil {
-			err = errors.Join(err, aErr)
-		}
-	}
-
 	if s.prometheusMetricsExporter != nil {
 		if aErr := s.prometheusMetricsExporter.Shutdown(ctx); aErr != nil {
 			err = errors.Join(err, aErr)
@@ -1021,7 +1031,7 @@ func (s *graphMux) Shutdown(ctx context.Context) error {
 // The mux is appended internally to the graph server's list of muxes to clean up later when the server is swapped.
 func (s *graphServer) buildGraphMux(
 	opts BuildGraphMuxOptions,
-) (*graphMux, error) {
+) (_ *graphMux, resErr error) {
 	graphMuxCtx, graphMuxCancel := context.WithCancel(s.routerCtx)
 
 	gm := &graphMux{
@@ -1031,6 +1041,21 @@ func (s *graphServer) buildGraphMux(
 		streamMetricStore: rmetric.NewNoopStreamMetricStore(),
 	}
 
+	// A failed mux isn't in s.graphMuxList yet (added on success below), so the graph
+	// server's Shutdown won't see it. Clean it up here to avoid leaking its callbacks,
+	// context and caches.
+	defer func() {
+		if resErr == nil {
+			return
+		}
+
+		cleanupCtx, cancel := context.WithTimeout(s.routerCtx, time.Second*30)
+		defer cancel()
+		if cErr := gm.Shutdown(cleanupCtx); cErr != nil {
+			s.logger.Error("Failed to clean up partially-built graph mux after build error", zap.Error(cErr))
+		}
+	}()
+
 	httpRouter := chi.NewRouter()
 
 	// we only enable the attribute mapper if we are not using the default cloud exporter
@@ -2100,6 +2125,38 @@ func (s *graphServer) wait(ctx context.Context) error {
 	}
 }
 
+// metricsFlushTimeout bounds the single, central flush of the shared meter
+// providers during graph server shutdown.
+const metricsFlushTimeout = 30 * time.Second
+
+// flushMeterProviders flushes the OTLP and Prometheus meter providers once. These
+// providers are shared by every metric store (request, connection, stream,
+// engine, runtime), so a single flush drains all of their metrics.
+func (s *graphServer) flushMeterProviders(ctx context.Context) error {
+	wg := &sync.WaitGroup{}
+
+	var otlpErr error
+	var promErr error
+
+	if s.otlpMeterProvider != nil {
+		wg.Go(func() {
+			if err := s.otlpMeterProvider.ForceFlush(ctx); err != nil {
+				otlpErr = errors.Join(otlpErr, fmt.Errorf("failed to flush otlp metrics: %w", err))
+			}
+		})
+	}
+	if s.promMeterProvider != nil {
+		wg.Go(func() {
+			if err := s.promMeterProvider.ForceFlush(ctx); err != nil {
+				promErr = errors.Join(promErr, fmt.Errorf("failed to flush prometheus metrics: %w", err))
+			}
+		})
+	}
+	wg.Wait()
+
+	return errors.Join(otlpErr, promErr)
+}
+
 // Shutdown gracefully shutdown the server and waits for all in-flight requests to finish.
 // After all requests are done, it will shut down the metric store and runtime metrics.
 // Shutdown does cancel the context after all non-hijacked requests such as WebSockets has been handled.
@@ -2125,6 +2182,16 @@ func (s *graphServer) Shutdown(ctx context.Context) error {
 		zap.String("config_version", s.baseRouterConfigVersion),
 	)
 
+	// Flush the meter providers exactly once, with their own timeout,
+	// before tearing down the individual metric stores.
+	// As all the stores share the same meter providers, we only need to flush once
+	// before initiating the shutdown of the individual stores.
+	flushCtx, flushCancel := context.WithTimeout(ctx, metricsFlushTimeout)
+	if err := s.flushMeterProviders(flushCtx); err != nil {
+		finalErr = errors.Join(finalErr, fmt.Errorf("failed to flush metrics: %w", err))
+	}
+	flushCancel()
+
 	// Ensure that we don't wait indefinitely for shutdown
 	if s.routerGracePeriod > 0 {
 		newCtx, cancel := context.WithTimeout(ctx, s.routerGracePeriod)

router/core/graphql_prehandler.go

@@ -1244,21 +1244,18 @@ func (h *PreHandler) flushMetrics(ctx context.Context, requestLogger *zap.Logger
 	now := time.Now()
 
 	wg := &sync.WaitGroup{}
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
+
+	wg.Go(func() {
 		if err := h.metrics.MetricStore().Flush(ctx); err != nil {
 			requestLogger.Error("Failed to flush OTEL metrics", zap.Error(err))
 		}
-	}()
+	})
 
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
+	wg.Go(func() {
 		if err := h.tracerProvider.ForceFlush(ctx); err != nil {
 			requestLogger.Error("Failed to flush OTEL tracer", zap.Error(err))
 		}
-	}()
+	})
 
 	wg.Wait()
 

router/internal/exporter/exporter.go

@@ -9,6 +9,7 @@ import (
 	"github.com/cloudflare/backoff"
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
+	"golang.org/x/sync/semaphore"
 )
 
 // Exporter is a generic, thread-safe batch exporter that queues items and sends them
@@ -23,7 +24,8 @@ type Exporter[T any] struct {
 	acceptTrafficSema chan struct{}
 	queue             chan T
 	inflightBatches   *atomic.Int64
-	batchBufferPool   *sync.Pool // Pool for batch slice buffers to reduce allocations
+	batchBufferPool   *sync.Pool          // Pool for batch slice buffers to reduce allocations
+	exportSema        *semaphore.Weighted // Bounds the number of concurrent in-flight export goroutines
 
 	// exportRequestContext is used to cancel all requests that started before the shutdown
 	exportRequestContext context.Context
@@ -47,6 +49,7 @@ const (
 	defaultMaxBatchItems          = 1024
 	defaultMaxQueueSize           = 1024 * 10
 	defaultBatchInterval          = time.Duration(10) * time.Second
+	defaultMaxConcurrentExports   = 100
 )
 
 type ExporterSettings struct {
@@ -60,14 +63,19 @@ type ExporterSettings struct {
 	RetryOptions RetryOptions
 	// ExportTimeout is the timeout for the export request.
 	ExportTimeout time.Duration
+	// MaxConcurrentExports caps the number of in-flight export goroutines.
+	// When <= 0, defaultMaxConcurrentExports is used. Bounding this prevents a slow
+	// or failing sink from accumulating unbounded retrying batches in memory.
+	MaxConcurrentExports int
 }
 
 func NewDefaultExporterSettings() *ExporterSettings {
 	return &ExporterSettings{
-		BatchSize:     defaultMaxBatchItems,
-		QueueSize:     defaultMaxQueueSize,
-		Interval:      defaultBatchInterval,
-		ExportTimeout: defaultExportTimeout,
+		BatchSize:            defaultMaxBatchItems,
+		QueueSize:            defaultMaxQueueSize,
+		Interval:             defaultBatchInterval,
+		ExportTimeout:        defaultExportTimeout,
+		MaxConcurrentExports: defaultMaxConcurrentExports,
 		RetryOptions: RetryOptions{
 			Enabled:     true,
 			MaxRetry:    defaultExportMaxRetryAttempts,
@@ -93,6 +101,13 @@ func NewExporter[T any](logger *zap.Logger, sink Sink[T], isRetryableError SinkE
 		isRetryableError = func(err error) bool { return true }
 	}
 
+	// Bound concurrent exports; fall back to the default when unset so callers
+	// that construct ExporterSettings directly don't get an unbounded pipeline.
+	maxConcurrentExports := settings.MaxConcurrentExports
+	if maxConcurrentExports <= 0 {
+		maxConcurrentExports = defaultMaxConcurrentExports
+	}
+
 	e := &Exporter[T]{
 		logger:                  logger.With(zap.String("component", "exporter")),
 		settings:                settings,
@@ -102,6 +117,7 @@ func NewExporter[T any](logger *zap.Logger, sink Sink[T], isRetryableError SinkE
 		shutdownSignal:          make(chan struct{}),
 		acceptTrafficSema:       make(chan struct{}),
 		inflightBatches:         atomic.NewInt64(0),
+		exportSema:              semaphore.NewWeighted(int64(maxConcurrentExports)),
 		exportRequestContext:    ctx,
 		cancelAllExportRequests: cancel,
 		batchBufferPool: &sync.Pool{
@@ -161,9 +177,18 @@ func (e *Exporter[T]) getBatchBuffer() []T {
 }
 
 // putBatchBuffer returns a batch buffer to the pool for reuse.
-// The buffer is reset to zero length before being pooled.
+// Element references are cleared and the slice is reset to zero length before pooling.
 func (e *Exporter[T]) putBatchBuffer(buffer []T) {
-	// Reset the slice to zero length while keeping capacity
+	// Drop buffers that grew beyond the configured batch size instead of pooling
+	// them, so the pool doesn't retain oversized backing arrays indefinitely.
+	if cap(buffer) > e.settings.BatchSize {
+		return
+	}
+	// Clear element references before pooling. Reslicing to [:0] only changes the
+	// length: the backing array still holds the (already-exported) item pointers,
+	// and the GC scans the full capacity, so without this the pooled buffer would
+	// pin every item from the last batch until those slots are overwritten.
+	clear(buffer)
 	buffer = buffer[:0]
 	e.batchBufferPool.Put(&buffer)
 }
@@ -232,8 +257,20 @@ func (e *Exporter[T]) exportBatch(batch []T) error {
 func (e *Exporter[T]) prepareAndSendBatch(batch []T) {
 	e.logger.Debug("Preparing to send batch", zap.Int("batch_size", len(batch)))
 	e.inflightBatches.Inc()
+	// Acquire a slot before spawning the goroutine. When all slots are taken this
+	// blocks the caller (the start/drain loop), applying backpressure all the way
+	// up to Record's queue instead of accumulating retrying batches in memory.
+	// The context is only cancelled on forced shutdown, in which case exports would
+	// fail anyway: return the buffer to the pool and drop the batch.
+	if err := e.exportSema.Acquire(e.exportRequestContext, 1); err != nil {
+		e.logger.Debug("Skipping batch export, exporter is shutting down", zap.Error(err))
+		e.putBatchBuffer(batch)
+		e.inflightBatches.Dec()
+		return
+	}
 	go func() {
 		defer e.inflightBatches.Dec()
+		defer e.exportSema.Release(1) // Release the slot once the export (and any retries) completes
 		defer e.putBatchBuffer(batch) // Return buffer to pool after export completes
 		e.exportBatchWithRetry(batch)
 	}()