preserving the capacity of the Symbols slice during resets

Signed-off-by: SungJin1212 <tjdwls1201@gmail.com>

Author: SungJin1212

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Changelog
 
 ## master / unreleased
+* [CHANGE] Querier: Make query time range configurations per-tenant: `query_ingesters_within`, `query_store_after`, and `shuffle_sharding_ingesters_lookback_period`. Uses `model.Duration` instead of `time.Duration` to support serialization but has minimum unit of 1ms (nanoseconds/microseconds not supported). #7160
 * [FEATURE] Distributor: Add experimental `-distributor.enable-start-timestamp` flag for Prometheus Remote Write 2.0. When enabled, `StartTimestamp (ST)` is ingested. #7371
 * [FEATURE] Memberlist: Add `-memberlist.cluster-label` and `-memberlist.cluster-label-verification-disabled` to prevent accidental cross-cluster gossip joins and support rolling label rollout. #7385
 * [FEATURE] Querier: Add timeout classification to classify query timeouts as 4XX (user error) or 5XX (system error) based on phase timing. When enabled, queries that spend most of their time in PromQL evaluation return `422 Unprocessable Entity` instead of `503 Service Unavailable`. #7374
@@ -9,11 +10,11 @@
 * [ENHANCEMENT] Metrics Helper: Add native histogram support for aggregating and merging, including dual-format histogram handling that exposes both native and classic bucket formats. #7359
 * [ENHANCEMENT] Cache: Add per-tenant TTL configuration for query results cache to control cache expiration on a per-tenant basis with separate TTLs for regular and out-of-order data. #7357
 * [ENHANCEMENT] Update build image and Go version to 1.26. #7437
-* [CHANGE] Querier: Make query time range configurations per-tenant: `query_ingesters_within`, `query_store_after`, and `shuffle_sharding_ingesters_lookback_period`. Uses `model.Duration` instead of `time.Duration` to support serialization but has minimum unit of 1ms (nanoseconds/microseconds not supported). #7160
 * [ENHANCEMENT] Tenant Federation: Add a local cache to regex resolver. #7363
 * [ENHANCEMENT] Query Scheduler: Add `cortex_query_scheduler_tracked_requests` metric to track the current number of requests held by the scheduler. #7355
 * [ENHANCEMENT] Compactor: Prevent partition compaction to compact any blocks marked for deletion. #7391
 * [ENHANCEMENT] Distributor: Optimize memory allocations by reusing the existing capacity of these pooled slices in the Prometheus Remote Write 2.0 path. #7392
+* [ENHANCEMENT] Distributor: Optimize memory allocations by pooling PreallocWriteRequestV2 and preserving the capacity of the Symbols slice during resets. #7404
 * [BUGFIX] Alertmanager: Fix disappearing user config and state when ring is temporarily unreachable. #7372
 * [BUGFIX] Fix nil when ingester_query_max_attempts > 1. #7369
 * [BUGFIX] Querier: Fix queryWithRetry and labelsWithRetry returning (nil, nil) on cancelled context by propagating ctx.Err(). #7370

pkg/cortexpb/timeseriesv2.go

@@ -130,6 +130,14 @@ func PreallocWriteRequestV2FromPool() *PreallocWriteRequestV2 {
 	return writeRequestPoolV2.Get().(*PreallocWriteRequestV2)
 }
 
+// Reset implements proto.Message and preserves the capacity of the Symbols slice.
+func (p *PreallocWriteRequestV2) Reset() {
+	savedSymbols := p.Symbols
+	p.WriteRequestV2.Reset()
+	p.Symbols = savedSymbols[:0]
+	p.data = nil
+}
+
 // PreallocTimeseriesV2SliceFromPool retrieves a slice of PreallocTimeseriesV2 from a sync.Pool.
 // ReuseSliceV2 should be called once done.
 func PreallocTimeseriesV2SliceFromPool() []PreallocTimeseriesV2 {

pkg/cortexpb/timeseriesv2_test.go

@@ -122,6 +122,71 @@ func TestReuseWriteRequestV2(t *testing.T) {
 	})
 }
 
+func TestPreallocWriteRequestV2Reset(t *testing.T) {
+	t.Run("preserves Symbols capacity", func(t *testing.T) {
+		const symbolsCap = 100
+		req := &PreallocWriteRequestV2{
+			WriteRequestV2: WriteRequestV2{
+				Symbols: make([]string, 0, symbolsCap),
+			},
+		}
+		req.Symbols = append(req.Symbols, "a", "b", "c")
+
+		ptrBefore := &req.Symbols[:cap(req.Symbols)][0]
+
+		req.Reset()
+
+		assert.Equal(t, 0, len(req.Symbols), "Symbols length should be 0 after Reset")
+		assert.Equal(t, symbolsCap, cap(req.Symbols), "Symbols capacity should be preserved after Reset")
+		assert.Same(t, ptrBefore, &req.Symbols[:cap(req.Symbols)][0], "Symbols backing array should be reused after Reset")
+	})
+
+	t.Run("clears non-Symbols WriteRequestV2 fields", func(t *testing.T) {
+		b := []byte{1, 2, 3}
+		req := &PreallocWriteRequestV2{
+			WriteRequestV2: WriteRequestV2{
+				Source:                  RULE,
+				SkipLabelNameValidation: true,
+				Timeseries:              []PreallocTimeseriesV2{{TimeSeriesV2: &TimeSeriesV2{}}},
+			},
+			data: &b,
+		}
+
+		req.Reset()
+
+		assert.Equal(t, SourceEnum(0), req.Source)
+		assert.False(t, req.SkipLabelNameValidation)
+		assert.Nil(t, req.Timeseries)
+		assert.Nil(t, req.data)
+	})
+
+	t.Run("Unmarshal after Reset reuses Symbols backing array", func(t *testing.T) {
+		const symbolsCount = 50
+		symbols := make([]string, symbolsCount)
+		for i := range symbols {
+			symbols[i] = fmt.Sprintf("symbol_%04d", i)
+		}
+		data, err := (&WriteRequestV2{Symbols: symbols}).Marshal()
+		require.NoError(t, err)
+
+		req := &PreallocWriteRequestV2{
+			WriteRequestV2: WriteRequestV2{
+				Symbols: make([]string, 0, symbolsCount*2),
+			},
+		}
+
+		// Simulate Reset in util.ParseProtoReader()
+		req.Reset()
+		ptrAfterReset := &req.Symbols[:cap(req.Symbols)][0]
+		capAfterReset := cap(req.Symbols)
+
+		require.NoError(t, req.WriteRequestV2.Unmarshal(data))
+		assert.Equal(t, symbolsCount, len(req.Symbols))
+		assert.Equal(t, capAfterReset, cap(req.Symbols), "capacity should not change: Unmarshal reused the existing backing array")
+		assert.Same(t, ptrAfterReset, &req.Symbols[:cap(req.Symbols)][0], "backing array pointer should be identical: no new allocation occurred")
+	})
+}
+
 func BenchmarkMarshallWriteRequestV2(b *testing.B) {
 	ts := PreallocTimeseriesV2SliceFromPool()
 

pkg/util/push/push.go

@@ -94,14 +94,14 @@ func Handler(remoteWrite2Enabled bool, acceptUnknownRemoteWriteContentType bool,
 				return
 			}
 
-			var req cortexpb.PreallocWriteRequestV2
+			req := cortexpb.PreallocWriteRequestV2FromPool()
 			// v1 request is put back into the pool by the Distributor.
 			defer func() {
-				cortexpb.ReuseWriteRequestV2(&req)
+				cortexpb.ReuseWriteRequestV2(req)
 				req.Free()
 			}()
 
-			err = util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), maxRecvMsgSize, &req, util.RawSnappy)
+			err = util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), maxRecvMsgSize, req, util.RawSnappy)
 			if err != nil {
 				level.Error(logger).Log("err", err.Error())
 				http.Error(w, err.Error(), http.StatusBadRequest)
@@ -113,7 +113,7 @@ func Handler(remoteWrite2Enabled bool, acceptUnknownRemoteWriteContentType bool,
 				req.Source = cortexpb.API
 			}
 
-			v1Req, err := convertV2RequestToV1(&req, overrides.EnableTypeAndUnitLabels(userID), overrides.EnableStartTimestamp(userID))
+			v1Req, err := convertV2RequestToV1(req, overrides.EnableTypeAndUnitLabels(userID), overrides.EnableStartTimestamp(userID))
 			if err != nil {
 				level.Error(logger).Log("err", err.Error())
 				http.Error(w, err.Error(), http.StatusBadRequest)