receive: Add map pooling and optimize sendLocalWrite to reduce alloca…

pkg/receive/handler.go

@@ -122,6 +122,29 @@ var (
 			return &b
 		},
 	}
+
+	// Distribution map pools to reduce allocations in distributeTimeseriesToReplicas.
+	// These prevent the ~357GB of allocations observed in production profiling.
+	maxPooledDistributionMapSize = 100 // Max endpoints * replicas we'll pool
+
+	endpointReplicaMapPool = sync.Pool{
+		New: func() interface{} {
+			m := make(map[endpointReplica]map[string]trackedSeries)
+			return &m
+		},
+	}
+	tenantSeriesMapPool = sync.Pool{
+		New: func() interface{} {
+			m := make(map[string]trackedSeries, 1)
+			return &m
+		},
+	}
+	tenantTimeseriesMappingPool = sync.Pool{
+		New: func() interface{} {
+			m := make(map[string][]prompb.TimeSeries)
+			return &m
+		},
+	}
 )
 
 type sizeLimiter interface {
@@ -511,6 +534,73 @@ func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusap
 	return h.options.TSDBStats.TenantStats(statsLimit, statsByLabelName, tenantID), nil
 }
 
+// clearEndpointReplicaMap clears a distribution map and returns whether it should be pooled.
+// Maps that are too large are not pooled to prevent pool ballooning.
+func clearEndpointReplicaMap(m *map[endpointReplica]map[string]trackedSeries) bool {
+	if len(*m) > maxPooledDistributionMapSize {
+		return false
+	}
+	// Clear the outer map by deleting all entries
+	for k := range *m {
+		delete(*m, k)
+	}
+	return true
+}
+
+// clearTenantSeriesMap clears a tenant series map and returns whether it should be pooled.
+func clearTenantSeriesMap(m *map[string]trackedSeries) bool {
+	if len(*m) > maxPooledDistributionMapSize {
+		return false
+	}
+	// Clear the map by deleting all entries
+	for k := range *m {
+		delete(*m, k)
+	}
+	return true
+}
+
+// clearTenantTimeseriesMapping clears a tenant timeseries mapping and returns whether it should be pooled.
+func clearTenantTimeseriesMapping(m *map[string][]prompb.TimeSeries) bool {
+	if len(*m) > maxPooledDistributionMapSize {
+		return false
+	}
+	// Clear the map and the slices to avoid holding references
+	for k := range *m {
+		// Clear the slice to release references to TimeSeries
+		(*m)[k] = nil
+		delete(*m, k)
+	}
+	return true
+}
+
+// returnDistributionMapsToPool returns the distribution maps to their respective pools after clearing them.
+// This should be called after the maps are no longer needed to reduce memory allocations.
+func returnDistributionMapsToPool(localWrites, remoteWrites map[endpointReplica]map[string]trackedSeries) {
+	// Return nested maps first
+	for er := range localWrites {
+		if nestedMap := localWrites[er]; nestedMap != nil {
+			if clearTenantSeriesMap(&nestedMap) {
+				tenantSeriesMapPool.Put(&nestedMap)
+			}
+		}
+	}
+	for er := range remoteWrites {
+		if nestedMap := remoteWrites[er]; nestedMap != nil {
+			if clearTenantSeriesMap(&nestedMap) {
+				tenantSeriesMapPool.Put(&nestedMap)
+			}
+		}
+	}
+
+	// Return outer maps
+	if clearEndpointReplicaMap(&localWrites) {
+		endpointReplicaMapPool.Put(&localWrites)
+	}
+	if clearEndpointReplicaMap(&remoteWrites) {
+		endpointReplicaMapPool.Put(&remoteWrites)
+	}
+}
+
 // tenantKeyForDistribution matches distributeTimeseriesToReplicas semantics exactly.
 func (h *Handler) tenantKeyForDistribution(tenantHTTP string, ts prompb.TimeSeries) string {
 	tenant := tenantHTTP
@@ -963,6 +1053,8 @@ func (h *Handler) fanoutForward(ctx context.Context, params remoteWriteParams) (
 		level.Error(requestLogger).Log("msg", "failed to distribute timeseries to replicas", "err", err)
 		return stats, err
 	}
+	// Return distribution maps to pool after they're no longer needed to reduce allocations
+	defer returnDistributionMapsToPool(localWrites, remoteWrites)
 
 	stats = h.gatherWriteStats(len(params.replicas), localWrites, remoteWrites)
 
@@ -1058,14 +1150,21 @@ func (h *Handler) distributeTimeseriesToReplicas(
 ) (map[endpointReplica]map[string]trackedSeries, map[endpointReplica]map[string]trackedSeries, error) {
 	h.mtx.RLock()
 	defer h.mtx.RUnlock()
-	remoteWrites := make(map[endpointReplica]map[string]trackedSeries)
-	localWrites := make(map[endpointReplica]map[string]trackedSeries)
+
+	// Get outer maps from pool to reduce allocations (357GB in production profiling)
+	remoteWritesPtr := endpointReplicaMapPool.Get().(*map[endpointReplica]map[string]trackedSeries)
+	localWritesPtr := endpointReplicaMapPool.Get().(*map[endpointReplica]map[string]trackedSeries)
+	remoteWrites := *remoteWritesPtr
+	localWrites := *localWritesPtr
+
 	for tsIndex, ts := range timeseries {
 		tenant := h.tenantKeyForDistribution(tenantHTTP, ts)
 
 		for _, rn := range replicas {
 			endpoint, err := h.hashring.GetN(tenant, &ts, rn)
 			if err != nil {
+				// Clean up pooled maps on error
+				returnDistributionMapsToPool(localWrites, remoteWrites)
 				return nil, nil, err
 			}
 			endpointReplica := endpointReplica{endpoint: endpoint, replica: rn}
@@ -1075,7 +1174,9 @@ func (h *Handler) distributeTimeseriesToReplicas(
 			}
 			writeableSeries, ok := writeDestination[endpointReplica]
 			if !ok {
-				writeableSeries = make(map[string]trackedSeries, 1)
+				// Get nested map from pool
+				writeableSeriesPtr := tenantSeriesMapPool.Get().(*map[string]trackedSeries)
+				writeableSeries = *writeableSeriesPtr
 				writeDestination[endpointReplica] = writeableSeries
 			}
 			tenantSeries := writeableSeries[tenant]
@@ -1138,13 +1239,41 @@ func (h *Handler) sendLocalWrite(
 	span.SetTag("endpoint", writeDestination.endpoint)
 	span.SetTag("replica", writeDestination.replica)
 
-	tenantSeriesMapping := map[string][]prompb.TimeSeries{}
-	for _, ts := range trackedSeries.timeSeries {
+	// Get tenantSeriesMapping from pool to reduce allocations (10.30GB in production profiling)
+	tenantSeriesMappingPtr := tenantTimeseriesMappingPool.Get().(*map[string][]prompb.TimeSeries)
+	tenantSeriesMapping := *tenantSeriesMappingPtr
+	defer func() {
+		if clearTenantTimeseriesMapping(&tenantSeriesMapping) {
+			tenantTimeseriesMappingPool.Put(&tenantSeriesMapping)
+		}
+	}()
+
+	// Pre-allocate slices to reduce reallocations. Most requests have a single tenant,
+	// so we allocate the full capacity for the default tenant.
+	numSeries := len(trackedSeries.timeSeries)
+	if h.splitTenantLabelName == "" {
+		// Single tenant case - pre-allocate full capacity
+		tenantSeriesMapping[tenantHTTP] = make([]prompb.TimeSeries, 0, numSeries)
+	}
+
+	// Iterate by index to avoid copying large TimeSeries structs (contains 4 slice headers)
+	// This reduces allocations attributed to zlabelsGet (244.65GB in production profiling)
+	for i := range trackedSeries.timeSeries {
+		ts := trackedSeries.timeSeries[i]
 		var tenant = tenantHTTP
 		if h.splitTenantLabelName != "" {
 			if tnt, ok := zlabelsGet(ts.Labels, h.splitTenantLabelName); ok && tnt != "" {
 				tenant = tnt
 			}
+			// For multi-tenant case, lazily allocate with estimated capacity
+			if _, exists := tenantSeriesMapping[tenant]; !exists {
+				// Estimate: assume even distribution across tenants (heuristic)
+				estimatedCap := numSeries / 4
+				if estimatedCap < 1 {
+					estimatedCap = 1
+				}
+				tenantSeriesMapping[tenant] = make([]prompb.TimeSeries, 0, estimatedCap)
+			}
 		}
 		tenantSeriesMapping[tenant] = append(tenantSeriesMapping[tenant], ts)
 	}

pkg/receive/handler_test.go

@@ -2055,3 +2055,94 @@ func startIngestor(logger log.Logger, serverAddress string, delay time.Duration)
 	}()
 	return srv
 }
+
+// TestDistributionMapPooling verifies that distribution maps are properly
+// pooled and reused to reduce allocations.
+func TestDistributionMapPooling(t *testing.T) {
+	// Get initial pool counts by draining and refilling
+	initialRemote := endpointReplicaMapPool.Get().(*map[endpointReplica]map[string]trackedSeries)
+	initialLocal := endpointReplicaMapPool.Get().(*map[endpointReplica]map[string]trackedSeries)
+	endpointReplicaMapPool.Put(initialRemote)
+	endpointReplicaMapPool.Put(initialLocal)
+
+	// Create test data
+	localWrites := make(map[endpointReplica]map[string]trackedSeries)
+	remoteWrites := make(map[endpointReplica]map[string]trackedSeries)
+
+	// Add some data
+	ep1 := endpointReplica{endpoint: Endpoint{Address: "test1"}, replica: 0}
+	ep2 := endpointReplica{endpoint: Endpoint{Address: "test2"}, replica: 1}
+
+	tenantMap1 := make(map[string]trackedSeries)
+	tenantMap1["tenant1"] = trackedSeries{
+		seriesIDs:  []int{0, 1},
+		timeSeries: []prompb.TimeSeries{{}, {}},
+	}
+	localWrites[ep1] = tenantMap1
+
+	tenantMap2 := make(map[string]trackedSeries)
+	tenantMap2["tenant2"] = trackedSeries{
+		seriesIDs:  []int{2, 3},
+		timeSeries: []prompb.TimeSeries{{}, {}},
+	}
+	remoteWrites[ep2] = tenantMap2
+
+	// Verify maps have data
+	if len(localWrites) != 1 || len(remoteWrites) != 1 {
+		t.Fatalf("Expected maps to have data, got local=%d, remote=%d", len(localWrites), len(remoteWrites))
+	}
+
+	// Return maps to pool
+	returnDistributionMapsToPool(localWrites, remoteWrites)
+
+	// Verify maps were cleared
+	if len(localWrites) != 0 || len(remoteWrites) != 0 {
+		t.Errorf("Expected maps to be cleared after returning to pool, got local=%d, remote=%d", len(localWrites), len(remoteWrites))
+	}
+
+	// Get maps from pool again and verify they're empty
+	reusedRemote := endpointReplicaMapPool.Get().(*map[endpointReplica]map[string]trackedSeries)
+	reusedLocal := endpointReplicaMapPool.Get().(*map[endpointReplica]map[string]trackedSeries)
+
+	if len(*reusedRemote) != 0 || len(*reusedLocal) != 0 {
+		t.Errorf("Expected reused maps to be empty, got remote=%d, local=%d", len(*reusedRemote), len(*reusedLocal))
+	}
+
+	// Clean up
+	endpointReplicaMapPool.Put(reusedRemote)
+	endpointReplicaMapPool.Put(reusedLocal)
+}
+
+// TestClearMapFunctions verifies the map clearing helpers work correctly
+func TestClearMapFunctions(t *testing.T) {
+	// Test clearEndpointReplicaMap
+	outerMap := make(map[endpointReplica]map[string]trackedSeries)
+	outerMap[endpointReplica{endpoint: Endpoint{Address: "test"}, replica: 0}] = make(map[string]trackedSeries)
+
+	if !clearEndpointReplicaMap(&outerMap) {
+		t.Error("Expected small map to be poolable")
+	}
+	if len(outerMap) != 0 {
+		t.Errorf("Expected map to be cleared, got len=%d", len(outerMap))
+	}
+
+	// Test with oversized map
+	largeMap := make(map[endpointReplica]map[string]trackedSeries)
+	for i := 0; i <= maxPooledDistributionMapSize; i++ {
+		largeMap[endpointReplica{endpoint: Endpoint{Address: "test"}, replica: uint64(i)}] = make(map[string]trackedSeries)
+	}
+	if clearEndpointReplicaMap(&largeMap) {
+		t.Error("Expected large map to not be poolable")
+	}
+
+	// Test clearTenantSeriesMap
+	innerMap := make(map[string]trackedSeries)
+	innerMap["tenant1"] = trackedSeries{}
+
+	if !clearTenantSeriesMap(&innerMap) {
+		t.Error("Expected small map to be poolable")
+	}
+	if len(innerMap) != 0 {
+		t.Errorf("Expected map to be cleared, got len=%d", len(innerMap))
+	}
+}