fix: accumulate RetrieveResultCache size instead of recomputing it per merge

Project: milvus-io/milvus | Language: Go | Concepts: algorithmic-optimization, data-structures, memory-management | Difficulty: mid
Link: https://github.com/milvus-io/milvus/pull/52274

The problem

The RetrieveResultCache.merge method was recalculating the entire size of the accumulated results (proto.Size(c.result)) every time a new result was merged. Since the size calculation walks the whole message, this made merging a stream of N results O(N²) in the number of IDs—very inefficient for large N. This unnecessary recomputation happened under a lock, creating a performance bottleneck when handling many results, especially for delete-by-expression operations.

Why it's non-trivial

A naive fix might just try to optimize or cache the size calculation, but would miss that only a subset of the incoming data is actually retained after merging (specifically, just the IDs and counters, not all the fields). Over-accounting for discarded data would artificially limit batch sizes, fragmenting downstream processing and reducing throughput. Careful accounting is needed to avoid both underestimating (which could cause oversized messages) and overestimating (which would shrink batches unnecessarily).

The approach

1. Add a fixed headroom for message growth that isn't per-merge:

+// Those add up to well under 128 bytes for an entire accumulation; 256 is taken
+// once, at Put, so that the per-merge charge below can stay tight.
+const sizeGrowthHeadroom = 256

Annotation: This constant covers the small, bounded sources of growth in the protobuf message (like widening length prefixes or counters) that don't scale with each merge, so we don't have to overcharge per merge.

2. Only charge for the retained part of each result (IDs), not the whole envelope:

+// mergedSize reports what result will contribute to the accumulated message
+// once merge folds it in.
+func mergedSize(result *internalpb.RetrieveResults) int {
+	return proto.Size(result.GetIds())
+}

Annotation: This function calculates the size of just the IDs part of the result, which is the only part retained after merging. This avoids overestimating the size by excluding dropped fields.

3. Accumulate size instead of recomputing it:

-func (c *RetrieveResultCache) merge(result *internalpb.RetrieveResults) {
-	c.result.Ids = mergeIDs(c.result.GetIds(), result.GetIds())
-	c.result.AllRetrieveCount = c.result.GetAllRetrieveCount() + result.GetAllRetrieveCount()
-	c.result.CostAggregation = mergeCostAggregation(c.result.GetCostAggregation(), result.GetCostAggregation())
-	c.result.ScannedRemoteBytes = c.result.GetScannedRemoteBytes() + result.GetScannedRemoteBytes()
-	c.result.ScannedTotalBytes = c.result.GetScannedTotalBytes() + result.GetScannedTotalBytes()
-	c.size = proto.Size(c.result)
+func (c *RetrieveResultCache) merge(result *internalpb.RetrieveResults) {
+	c.result.Ids = mergeIDs(c.result.GetIds(), result.GetIds())
+	c.result.AllRetrieveCount = c.result.GetAllRetrieveCount() + result.GetAllRetrieveCount()
+	c.result.CostAggregation = mergeCostAggregation(c.result.GetCostAggregation(), result.GetCostAggregation())
+	c.result.ScannedRemoteBytes = c.result.GetScannedRemoteBytes() + result.GetScannedRemoteBytes()
+	c.result.ScannedTotalBytes = c.result.GetScannedTotalBytes() + result.GetScannedTotalBytes()
+	// Accumulate rather than recompute. `c.size = proto.Size(c.result)` walked
+	// the whole accumulated message on every merge, making a stream of N
+	// results O(N^2) in the total number of IDs.
+	c.size += mergedSize(result)
 }

Annotation: Instead of recalculating the size of the entire result after each merge, the code now just adds the size of the new IDs being merged. This makes the operation O(N) instead of O(N²).

4. Ensure allocation and flush logic use the same size accounting:

-func (c *RetrieveResultCache) Alloc(result *internalpb.RetrieveResults) bool {
-	return proto.Size(result)+c.size <= c.cap
+func (c *RetrieveResultCache) Alloc(result *internalpb.RetrieveResults) bool {
+	// Charge what Put will actually charge, so the flush decision here and the
+	// accounting in merge cannot disagree about the same message.
+	if c.result == nil {
+		return proto.Size(result)+sizeGrowthHeadroom+c.size <= c.cap
+	}
+	return mergedSize(result)+c.size <= c.cap
 }

Annotation: This ensures that the logic deciding whether a new result can be added to the batch matches the actual size accounting used during merging, preventing inconsistencies.

5. Add tests to check both correctness and utilization:

+func TestRetrieveResultCacheSizeNeverUnderestimates(t *testing.T) {
+	...
+	assert.GreaterOrEqual(t, cache.size, proto.Size(cache.result),
+		"reported size must not be below the real encoded size after %d merges", i+1)
+	...
+	assert.Less(t, cache.size-real, real/100,
+		"over-estimate should stay under 1%% of the real size, got %d of %d",
+		cache.size-real, real)
+}

Annotation: This test ensures that the accumulated size never underestimates the true serialized size (which would risk oversized messages), and that the overestimate stays minimal (less than 1%).

From the review thread

No notable review discussion on this PR.

Takeaways

• Always charge for exactly what is retained, not what is received, when accumulating or batching data. Over-accounting can fragment batches and reduce throughput; under-accounting can cause correctness or safety issues.
• Avoid O(N²) algorithms in hot code paths, especially under locks. Incremental accumulation is often possible and much more efficient.
• Keep size/accounting logic consistent between allocation checks and actual updates. Mismatches can lead to subtle bugs or violations of invariants.
• Test both safety (never underestimating) and utilization (not overestimating) when optimizing memory or size accounting.