From f95c3e0b3db2269b1a416ea19641526ac0ab917b Mon Sep 17 00:00:00 2001 From: mkusumdb Date: Tue, 3 Feb 2026 16:50:44 -0800 Subject: [PATCH 1/4] receive: Add buffer pooling to reduce heap allocations and GC pressure (#283) --- pkg/receive/handler.go | 189 +++++++++++++++++++++++++++++++-------- test/e2e/receive_test.go | 6 +- 2 files changed, 157 insertions(+), 38 deletions(-) diff --git a/pkg/receive/handler.go b/pkg/receive/handler.go index b2fdd382178..624f2b983f8 100644 --- a/pkg/receive/handler.go +++ b/pkg/receive/handler.go @@ -87,8 +87,77 @@ var ( errNotReady = errors.New("target not ready") errUnavailable = errors.New("target not available") errInternal = errors.New("internal error") + + // Used internally to abort reads when the limiter is exceeded mid-stream. + errRequestTooLarge = errors.New("write request too large") + + // Default / max capacities for pooled buffers. These caps prevent "pool ballooning" + // where a single large request permanently inflates process RSS. + defaultCompressedBufCap = 32 * 1024 + defaultDecompressedBufCap = 128 * 1024 + maxPooledCompressedCap = 1 << 20 // 1MB + maxPooledDecompressedCap = 4 << 20 // 4MB + copyBufSize = 32 * 1024 + + // Buffer/message pools to reduce allocations in receive hot path. + compressedBufPool = sync.Pool{ + New: func() interface{} { + return bytes.NewBuffer(make([]byte, 0, defaultCompressedBufCap)) + }, + } + decompressedBufPool = sync.Pool{ + New: func() interface{} { + b := make([]byte, 0, defaultDecompressedBufCap) + return &b + }, + } + writeRequestPool = sync.Pool{ + New: func() interface{} { + return &prompb.WriteRequest{} + }, + } + copyBufPool = sync.Pool{ + New: func() interface{} { + b := make([]byte, copyBufSize) + return &b + }, + } ) +type sizeLimiter interface { + AllowSizeBytes(string, int64) bool +} + +// limitedBufferWriter writes to a buffer but aborts if the tenant exceeds the limiter. +// This protects the server when Content-Length is missing or incorrect. +type limitedBufferWriter struct { + b *bytes.Buffer + limiter sizeLimiter + tenant string + seen int64 +} + +func (w *limitedBufferWriter) Write(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + w.seen += int64(len(p)) + if !w.limiter.AllowSizeBytes(w.tenant, w.seen) { + return 0, errRequestTooLarge + } + return w.b.Write(p) +} + +// zlabelsGet avoids ZLabels -> PromLabels conversion in hot paths. +func zlabelsGet(lbls []labelpb.ZLabel, name string) (string, bool) { + for _, l := range lbls { + if l.Name == name { + return l.Value, true + } + } + return "", false +} + type WriteableStoreAsyncClient interface { storepb.WriteableStoreClient RemoteWriteAsync(context.Context, *storepb.WriteRequest, endpointReplica, []int, chan writeResponse, func(error)) @@ -442,6 +511,18 @@ func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusap return h.options.TSDBStats.TenantStats(statsLimit, statsByLabelName, tenantID), nil } +// tenantKeyForDistribution matches distributeTimeseriesToReplicas semantics exactly. +func (h *Handler) tenantKeyForDistribution(tenantHTTP string, ts prompb.TimeSeries) string { + tenant := tenantHTTP + if h.splitTenantLabelName == "" { + return tenant + } + if v, ok := zlabelsGet(ts.Labels, h.splitTenantLabelName); ok && v != "" { + return h.splitTenantLabelName + ":" + v + } + return h.options.DefaultTenantID +} + // Close stops the Handler. func (h *Handler) Close() { _ = h.peers.Close() @@ -562,40 +643,75 @@ func (h *Handler) receiveHTTP(w http.ResponseWriter, r *http.Request) { } requestLimiter := h.Limiter.RequestLimiter() - // io.ReadAll dynamically adjust the byte slice for read data, starting from 512B. + // io.ReadAll dynamically adjusts the byte slice for read data, starting from 512B. // Since this is receive hot path, grow upfront saving allocations and CPU time. - compressed := bytes.Buffer{} + compressed := compressedBufPool.Get().(*bytes.Buffer) + defer func() { + // Avoid pooling huge buffers forever. + if compressed.Cap() <= maxPooledCompressedCap { + compressed.Reset() + compressedBufPool.Put(compressed) + } + }() + if r.ContentLength >= 0 { if !requestLimiter.AllowSizeBytes(tenantHTTP, r.ContentLength) { - http.Error(w, "write request too large", http.StatusRequestEntityTooLarge) + http.Error(w, errRequestTooLarge.Error(), http.StatusRequestEntityTooLarge) return } compressed.Grow(int(r.ContentLength)) } else { compressed.Grow(512) } - _, err = io.Copy(&compressed, r.Body) + + // Enforce size limits even when Content-Length is missing or wrong. + lw := &limitedBufferWriter{ + b: compressed, + limiter: requestLimiter, + tenant: tenantHTTP, + } + copyBuf := copyBufPool.Get().(*[]byte) + defer copyBufPool.Put(copyBuf) + _, err = io.CopyBuffer(lw, r.Body, *copyBuf) if err != nil { + if err == errRequestTooLarge { + http.Error(w, errRequestTooLarge.Error(), http.StatusRequestEntityTooLarge) + return + } http.Error(w, errors.Wrap(err, "read compressed request body").Error(), http.StatusInternalServerError) return } - reqBuf, err := s2.Decode(nil, compressed.Bytes()) + + // Decode into a pooled buffer to avoid allocs. (cap-guarded on return) + reqBuf := decompressedBufPool.Get().(*[]byte) + defer func() { + if cap(*reqBuf) <= maxPooledDecompressedCap { + *reqBuf = (*reqBuf)[:0] + decompressedBufPool.Put(reqBuf) + } + }() + *reqBuf, err = s2.Decode((*reqBuf)[:0], compressed.Bytes()) if err != nil { level.Error(tLogger).Log("msg", "snappy decode error", "err", err) http.Error(w, errors.Wrap(err, "snappy decode error").Error(), http.StatusBadRequest) return } - if !requestLimiter.AllowSizeBytes(tenantHTTP, int64(len(reqBuf))) { - http.Error(w, "write request too large", http.StatusRequestEntityTooLarge) + if !requestLimiter.AllowSizeBytes(tenantHTTP, int64(len(*reqBuf))) { + http.Error(w, errRequestTooLarge.Error(), http.StatusRequestEntityTooLarge) return } // NOTE: Due to zero copy ZLabels, Labels used from WriteRequests keeps memory // from the whole request. Ensure that we always copy those when we want to // store them for longer time. - var wreq prompb.WriteRequest - if err := proto.Unmarshal(reqBuf, &wreq); err != nil { + wreq := writeRequestPool.Get().(*prompb.WriteRequest) + wreq.Reset() + defer func() { + wreq.Reset() + writeRequestPool.Put(wreq) + }() + if err := proto.Unmarshal(*reqBuf, wreq); err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } @@ -637,14 +753,21 @@ func (h *Handler) receiveHTTP(w http.ResponseWriter, r *http.Request) { } // Apply relabeling configs. - h.relabel(&wreq) + h.relabel(wreq) if len(wreq.Timeseries) == 0 { level.Debug(tLogger).Log("msg", "remote write request dropped due to relabeling.") return } + // Deep copy all label strings to detach them from pooled decode buffer. + // Required for correctness when pooled buffers are reused and for preventing + // retention of the whole request buffer via zero-copy label references. + for i := range wreq.Timeseries { + labelpb.ReAllocZLabelsStrings(&wreq.Timeseries[i].Labels, h.writer.opts.Intern) + } + responseStatusCode := http.StatusOK - tenantStats, err := h.handleRequest(ctx, rep, tenantHTTP, &wreq) + tenantStats, err := h.handleRequest(ctx, rep, tenantHTTP, wreq) if err != nil { level.Debug(tLogger).Log("msg", "failed to handle request", "err", err.Error()) switch errors.Cause(err) { @@ -839,6 +962,13 @@ func (h *Handler) fanoutForward(ctx context.Context, params remoteWriteParams) ( stats = h.gatherWriteStats(len(params.replicas), localWrites, remoteWrites) + // Precompute seriesID -> tenantKey used by distributeTimeseriesToReplicas so we can + // attribute errorSeries correctly even when the responses channel closes. + seriesTenantKey := make([]string, len(params.writeRequest.Timeseries)) + for i, ts := range params.writeRequest.Timeseries { + seriesTenantKey[i] = h.tenantKeyForDistribution(params.tenant, ts) + } + // Prepare a buffered channel to receive the responses from the local and remote writes. Remote writes will all go // asynchronously and with this capacity we will never block on writing to the channel. maxBufferedResponses := len(localWrites) @@ -880,12 +1010,17 @@ func (h *Handler) fanoutForward(ctx context.Context, params remoteWriteParams) ( return stats, ctx.Err() case resp, hasMore := <-responses: if !hasMore { - for _, seriesErr := range seriesErrs { + for i, seriesErr := range seriesErrs { + // Count only the series that actually saw at least one error. + if len(seriesErr.errs) > 0 && i < len(seriesTenantKey) { + tk := seriesTenantKey[i] + if st, ok := stats[tk]; ok { + st.errorSeries++ + stats[tk] = st + } + } writeErrors.Add(seriesErr) } - if stat, ok := stats[resp.tenant]; ok { - stat.errorSeries += len(seriesErrs) - } return stats, writeErrors.ErrOrNil() } @@ -922,18 +1057,7 @@ func (h *Handler) distributeTimeseriesToReplicas( remoteWrites := make(map[endpointReplica]map[string]trackedSeries) localWrites := make(map[endpointReplica]map[string]trackedSeries) for tsIndex, ts := range timeseries { - var tenant = tenantHTTP - - if h.splitTenantLabelName != "" { - lbls := labelpb.ZLabelsToPromLabels(ts.Labels) - - tenantLabel := lbls.Get(h.splitTenantLabelName) - if tenantLabel != "" { - tenant = h.splitTenantLabelName + ":" + tenantLabel - } else { - tenant = h.options.DefaultTenantID - } - } + tenant := h.tenantKeyForDistribution(tenantHTTP, ts) for _, rn := range replicas { endpoint, err := h.hashring.GetN(tenant, &ts, rn) @@ -947,12 +1071,8 @@ func (h *Handler) distributeTimeseriesToReplicas( } writeableSeries, ok := writeDestination[endpointReplica] if !ok { - writeDestination[endpointReplica] = map[string]trackedSeries{ - tenant: { - seriesIDs: make([]int, 0), - timeSeries: make([]prompb.TimeSeries, 0), - }, - } + writeableSeries = make(map[string]trackedSeries, 1) + writeDestination[endpointReplica] = writeableSeries } tenantSeries := writeableSeries[tenant] @@ -1018,8 +1138,7 @@ func (h *Handler) sendLocalWrite( for _, ts := range trackedSeries.timeSeries { var tenant = tenantHTTP if h.splitTenantLabelName != "" { - lbls := labelpb.ZLabelsToPromLabels(ts.Labels) - if tnt := lbls.Get(h.splitTenantLabelName); tnt != "" { + if tnt, ok := zlabelsGet(ts.Labels, h.splitTenantLabelName); ok && tnt != "" { tenant = tnt } } diff --git a/test/e2e/receive_test.go b/test/e2e/receive_test.go index c938a4f0407..efe04bb9b71 100644 --- a/test/e2e/receive_test.go +++ b/test/e2e/receive_test.go @@ -1052,8 +1052,8 @@ func TestReceiveExtractsTenant(t *testing.T) { Timeseries: []prompb.TimeSeries{ { Labels: []prompb.Label{ - {Name: tenantLabelName, Value: "tenant-1"}, {Name: "aa", Value: "bb"}, + {Name: tenantLabelName, Value: "tenant-1"}, }, Samples: []prompb.Sample{ {Value: 1, Timestamp: time.Now().UnixMilli()}, @@ -1072,8 +1072,8 @@ func TestReceiveExtractsTenant(t *testing.T) { Timeseries: []prompb.TimeSeries{ { Labels: []prompb.Label{ - {Name: tenantLabelName, Value: "tenant-2"}, {Name: "aa", Value: "bb"}, + {Name: tenantLabelName, Value: "tenant-2"}, }, Samples: []prompb.Sample{ {Value: 1, Timestamp: time.Now().UnixMilli()}, @@ -1101,8 +1101,8 @@ func TestReceiveExtractsTenant(t *testing.T) { Timeseries: []prompb.TimeSeries{ { Labels: []prompb.Label{ - {Name: tenantLabelName, Value: "tenant-3"}, {Name: "aa", Value: "bb"}, + {Name: tenantLabelName, Value: "tenant-3"}, }, Samples: []prompb.Sample{ {Value: 1, Timestamp: time.Now().UnixMilli()}, From c1359fe5ff7be1cf5bc26805686b138abdb76a34 Mon Sep 17 00:00:00 2001 From: mkusumdb Date: Tue, 3 Feb 2026 17:27:08 -0800 Subject: [PATCH 2/4] receive: Detach exemplar labels from pooled buffer (#289) --- pkg/receive/handler.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pkg/receive/handler.go b/pkg/receive/handler.go index 624f2b983f8..4a1655edee6 100644 --- a/pkg/receive/handler.go +++ b/pkg/receive/handler.go @@ -764,6 +764,10 @@ func (h *Handler) receiveHTTP(w http.ResponseWriter, r *http.Request) { // retention of the whole request buffer via zero-copy label references. for i := range wreq.Timeseries { labelpb.ReAllocZLabelsStrings(&wreq.Timeseries[i].Labels, h.writer.opts.Intern) + // Also detach exemplar labels from the pooled buffer. + for j := range wreq.Timeseries[i].Exemplars { + labelpb.ReAllocZLabelsStrings(&wreq.Timeseries[i].Exemplars[j].Labels, h.writer.opts.Intern) + } } responseStatusCode := http.StatusOK From e96d18821b63161266aca14e310c2a009c3ea556 Mon Sep 17 00:00:00 2001 From: Yuchen Wang <162491048+yuchen-db@users.noreply.github.com> Date: Tue, 3 Feb 2026 23:37:08 -0800 Subject: [PATCH 3/4] receive: stagger head compaction across tenants (#288) --- cmd/thanos/receive.go | 13 ++++++++++++ pkg/receive/multitsdb.go | 37 ++++++++++++++++++++++++++++++++++- pkg/receive/multitsdb_test.go | 28 ++++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/cmd/thanos/receive.go b/cmd/thanos/receive.go index 1df000f4436..3e92649c85a 100644 --- a/cmd/thanos/receive.go +++ b/cmd/thanos/receive.go @@ -172,6 +172,11 @@ func runReceive( level.Info(logger).Log("msg", "tenant path segments before tenant feature enabled", "segments", path.Join(conf.tsdbPathSegmentsBeforeTenant...)) } + if *conf.compactionDelayInterval > 0 { + multiTSDBOptions = append(multiTSDBOptions, receive.WithCompactionDelayInterval(time.Duration(*conf.compactionDelayInterval))) + level.Info(logger).Log("msg", "deterministic compaction delay enabled", "interval", conf.compactionDelayInterval.String()) + } + rwTLSConfig, err := tls.NewServerConfig(log.With(logger, "protocol", "HTTP"), conf.rwServerCert, conf.rwServerKey, conf.rwServerClientCA, conf.rwServerTlsMinVersion) if err != nil { return err @@ -1000,6 +1005,8 @@ type receiveConfig struct { featureList *[]string noUploadTenants *[]string + + compactionDelayInterval *model.Duration } func (rc *receiveConfig) registerFlag(cmd extkingpin.FlagClause) { @@ -1112,6 +1119,12 @@ func (rc *receiveConfig) registerFlag(cmd extkingpin.FlagClause) { cmd.Flag("tsdb.no-lockfile", "Do not create lockfile in TSDB data directory. In any case, the lockfiles will be deleted on next startup.").Default("false").BoolVar(&rc.noLockFile) + rc.compactionDelayInterval = extkingpin.ModelDuration(cmd.Flag("receive.compaction-delay-interval", + "Interval for staggering head compaction across tenants. "+ + "Tenant N gets delay of N*interval (mod block duration). "+ + "0 uses random delay (default)."). + Default("0s")) + cmd.Flag("tsdb.max-exemplars", "Enables support for ingesting exemplars and sets the maximum number of exemplars that will be stored per tenant."+ " In case the exemplar storage becomes full (number of stored exemplars becomes equal to max-exemplars),"+ diff --git a/pkg/receive/multitsdb.go b/pkg/receive/multitsdb.go index 0eaa96bc2d1..a8749441f58 100644 --- a/pkg/receive/multitsdb.go +++ b/pkg/receive/multitsdb.go @@ -18,6 +18,7 @@ import ( "github.com/oklog/ulid" "github.com/pkg/errors" "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" "github.com/prometheus/prometheus/model/labels" "github.com/prometheus/prometheus/storage" "github.com/prometheus/prometheus/tsdb" @@ -70,6 +71,10 @@ type MultiTSDB struct { noUploadTenants []string // Support both exact matches and prefix patterns (e.g., "tenant1", "prod-*") enableTenantPathPrefix bool pathSegmentsBeforeTenant []string + + compactionDelayInterval time.Duration + tenantCounter atomic.Uint64 + compactionDelayGauge *prometheus.GaugeVec } // MultiTSDBOption is a functional option for MultiTSDB. @@ -110,6 +115,14 @@ func WithMatchersCache(cache storecache.MatchersCache) MultiTSDBOption { } } +// WithCompactionDelayInterval sets the interval for staggering head compaction across tenants. +// Tenant N gets delay of N*interval (mod block duration). +func WithCompactionDelayInterval(interval time.Duration) MultiTSDBOption { + return func(s *MultiTSDB) { + s.compactionDelayInterval = interval + } +} + // NewMultiTSDB creates new MultiTSDB. // NOTE: Passed labels must be sorted lexicographically (alphabetically). func NewMultiTSDB( @@ -149,6 +162,16 @@ func NewMultiTSDB( option(mt) } + if mt.compactionDelayInterval > 0 { + mt.compactionDelayGauge = promauto.With(reg).NewGaugeVec( + prometheus.GaugeOpts{ + Name: "thanos_receive_tsdb_compaction_delay_seconds", + Help: "Configured compaction delay for each tenant's TSDB in seconds.", + }, + []string{"tenant"}, + ) + } + return mt } @@ -792,7 +815,19 @@ func (t *MultiTSDB) startTSDB(logger log.Logger, tenantID string, tenant *tenant level.Info(logger).Log("msg", "opening TSDB") opts := *t.tsdbOpts opts.BlocksToDelete = tenant.blocksToDelete - opts.EnableDelayedCompaction = true + if t.compactionDelayInterval > 0 { + index := t.tenantCounter.Add(1) - 1 + chunkRange := time.Duration(t.tsdbOpts.MinBlockDuration) * time.Millisecond + delay := time.Duration(index) * t.compactionDelayInterval + if chunkRange > 0 { + delay = delay % chunkRange + } + opts.CompactionDelay = delay + level.Info(logger).Log("msg", "using deterministic compaction delay", "tenant", tenantID, "delay", delay) + t.compactionDelayGauge.WithLabelValues(tenantID).Set(delay.Seconds()) + } else { + opts.EnableDelayedCompaction = true + } tenant.blocksToDeleteFn = tsdb.DefaultBlocksToDelete // NOTE(GiedriusS): always set to false to properly handle OOO samples - OOO samples are written into the WBL diff --git a/pkg/receive/multitsdb_test.go b/pkg/receive/multitsdb_test.go index 65554bd7734..2c46607b6dd 100644 --- a/pkg/receive/multitsdb_test.go +++ b/pkg/receive/multitsdb_test.go @@ -1278,3 +1278,31 @@ func TestMultiTSDBSkipsLostAndFound(t *testing.T) { m.mtx.RUnlock() testutil.Assert(t, !exists, "lost+found should not be loaded as a tenant") } + +func TestMultiTSDBCompactionDelayInterval(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + interval := time.Minute + + m := NewMultiTSDB(dir, log.NewNopLogger(), prometheus.NewRegistry(), &tsdb.Options{ + MinBlockDuration: (2 * time.Hour).Milliseconds(), + MaxBlockDuration: (2 * time.Hour).Milliseconds(), + RetentionDuration: (6 * time.Hour).Milliseconds(), + }, labels.FromStrings("replica", "test"), "tenant_id", nil, false, metadata.NoneFunc, + WithCompactionDelayInterval(interval)) + t.Cleanup(func() { + testutil.Ok(t, m.Close()) + }) + + testutil.Ok(t, m.Open()) + + // Create multiple tenants and verify counter increments. + tenants := []string{"tenant-a", "tenant-b", "tenant-c"} + for _, tenant := range tenants { + testutil.Ok(t, appendSample(m, tenant, time.Now())) + } + + // Verify all tenants were created and counter was incremented. + testutil.Equals(t, uint64(len(tenants)), m.tenantCounter.Load()) +} From 034dae84641cbd48901fef569a90881c352b8b27 Mon Sep 17 00:00:00 2001 From: Yuchen Wang Date: Thu, 5 Feb 2026 17:05:14 -0800 Subject: [PATCH 4/4] ingestor only mode fast path --- pkg/receive/handler.go | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/pkg/receive/handler.go b/pkg/receive/handler.go index 4a1655edee6..3816d4bfaab 100644 --- a/pkg/receive/handler.go +++ b/pkg/receive/handler.go @@ -1248,6 +1248,28 @@ func (h *Handler) RemoteWrite(ctx context.Context, r *storepb.WriteRequest) (*st span, ctx := tracing.StartSpan(ctx, "receive_grpc") defer span.Finish() + // Fast path for IngestorOnly mode: write directly to local TSDB. + // This skips distributeTimeseriesToReplicas and sendLocalWrite since + // the Router already determined this data belongs to this node. + if h.receiverMode == IngestorOnly { + err := h.writer.Write(ctx, r.Tenant, r.Timeseries) + if err != nil { + level.Debug(h.logger).Log("msg", "failed to write to local TSDB", "err", err) + } + switch cause := errors.Cause(err); cause { + case nil: + return &storepb.WriteResponse{}, nil + default: + if isNotReady(cause) { + return nil, status.Error(codes.Unavailable, err.Error()) + } + if isConflict(cause) { + return nil, status.Error(codes.AlreadyExists, err.Error()) + } + return nil, status.Error(codes.Internal, err.Error()) + } + } + _, err := h.handleRequest(ctx, uint64(r.Replica), r.Tenant, &prompb.WriteRequest{Timeseries: r.Timeseries}) if err != nil { level.Debug(h.logger).Log("msg", "failed to handle request", "err", err)