fix(span-buffer): Do not flush forever after rebalancing (#88784)

When rebalancing happens, segments get sharded into different queues
than before. Meaning that done_flush_segments tries to delete the "done"
key from the wrong queue. Pull through the queue_key to
done_flush_segments to fix that.

Author: untitaker

src/sentry/spans/buffer.py

@@ -82,6 +82,8 @@
 # The segment ID in the Kafka protocol is only the span ID.
 SegmentKey = bytes
 
+QueueKey = bytes
+
 
 def _segment_key_to_span_id(segment_key: SegmentKey) -> bytes:
     return parse_segment_key(segment_key)[2]
@@ -127,6 +129,11 @@ class OutputSpan(NamedTuple):
     payload: dict[str, Any]
 
 
+class FlushedSegment(NamedTuple):
+    queue_key: QueueKey
+    spans: list[OutputSpan]
+
+
 class SpansBuffer:
     def __init__(
         self,
@@ -204,13 +211,13 @@ def process_spans(self, spans: Sequence[Span], now: int):
                         shard = self.assigned_shards[
                             int(span.trace_id, 16) % len(self.assigned_shards)
                         ]
-                        queue_keys.append(f"span-buf:q:{shard}")
+                        queue_keys.append(self._get_queue_key(shard))
 
                 results = p.execute()
 
         with metrics.timer("spans.buffer.process_spans.update_queue"):
-            queue_deletes: dict[str, set[bytes]] = {}
-            queue_adds: dict[str, MutableMapping[str | bytes, int]] = {}
+            queue_deletes: dict[bytes, set[bytes]] = {}
+            queue_adds: dict[bytes, MutableMapping[str | bytes, int]] = {}
 
             assert len(queue_keys) == len(results)
 
@@ -266,6 +273,9 @@ def _ensure_script(self):
         self.add_buffer_sha = self.client.script_load(add_buffer_script.script)
         return self.add_buffer_sha
 
+    def _get_queue_key(self, shard: int) -> bytes:
+        return f"span-buf:q:{shard}".encode("ascii")
+
     def _group_by_parent(self, spans: Sequence[Span]) -> dict[tuple[str, str], list[Span]]:
         """
         Groups partial trees of spans by their top-most parent span ID in the
@@ -296,32 +306,33 @@ def _group_by_parent(self, spans: Sequence[Span]) -> dict[tuple[str, str], list[
 
         return trees
 
-    def flush_segments(
-        self, now: int, max_segments: int = 0
-    ) -> tuple[int, dict[SegmentKey, list[OutputSpan]]]:
+    def flush_segments(self, now: int, max_segments: int = 0) -> dict[SegmentKey, FlushedSegment]:
         cutoff = now
 
+        queue_keys = []
+
         with metrics.timer("spans.buffer.flush_segments.load_segment_ids"):
             with self.client.pipeline(transaction=False) as p:
                 for shard in self.assigned_shards:
-                    key = f"span-buf:q:{shard}"
+                    key = self._get_queue_key(shard)
                     p.zrangebyscore(
                         key, 0, cutoff, start=0 if max_segments else None, num=max_segments or None
                     )
                     p.zcard(key)
+                    queue_keys.append(key)
 
                 result = iter(p.execute())
 
-        segment_keys = []
+        segment_keys: list[tuple[QueueKey, SegmentKey]] = []
         queue_sizes = []
 
         with metrics.timer("spans.buffer.flush_segments.load_segment_data"):
             with self.client.pipeline(transaction=False) as p:
                 # ZRANGEBYSCORE output
-                for segment_span_ids in result:
+                for queue_key, segment_span_ids in zip(queue_keys, result):
                     # process return value of zrevrangebyscore
                     for segment_key in segment_span_ids:
-                        segment_keys.append(segment_key)
+                        segment_keys.append((queue_key, segment_key))
                         p.smembers(segment_key)
 
                     # ZCARD output
@@ -340,10 +351,10 @@ def flush_segments(
 
         num_has_root_spans = 0
 
-        for segment_key, segment in zip(segment_keys, segments):
+        for (queue_key, segment_key), segment in zip(segment_keys, segments):
             segment_span_id = _segment_key_to_span_id(segment_key).decode("ascii")
 
-            return_segment = []
+            output_spans = []
             has_root_span = False
             metrics.timing("spans.buffer.flush_segments.num_spans_per_segment", len(segment))
             for payload in segment:
@@ -369,30 +380,30 @@ def flush_segments(
                     },
                 )
 
-                return_segment.append(OutputSpan(payload=val))
+                output_spans.append(OutputSpan(payload=val))
 
-            return_segments[segment_key] = return_segment
+            return_segments[segment_key] = FlushedSegment(queue_key=queue_key, spans=output_spans)
             num_has_root_spans += int(has_root_span)
+
         metrics.timing("spans.buffer.flush_segments.num_segments", len(return_segments))
         metrics.timing("spans.buffer.flush_segments.has_root_span", num_has_root_spans)
 
-        return sum(queue_sizes), return_segments
+        return return_segments
 
-    def done_flush_segments(self, segment_keys: dict[SegmentKey, list[OutputSpan]]):
+    def done_flush_segments(self, segment_keys: dict[SegmentKey, FlushedSegment]):
         metrics.timing("spans.buffer.done_flush_segments.num_segments", len(segment_keys))
         with metrics.timer("spans.buffer.done_flush_segments"):
             with self.client.pipeline(transaction=False) as p:
-                for segment_key, output_spans in segment_keys.items():
+                for segment_key, flushed_segment in segment_keys.items():
                     hrs_key = b"span-buf:hrs:" + segment_key
                     p.delete(hrs_key)
                     p.unlink(segment_key)
 
                     project_id, trace_id, _ = parse_segment_key(segment_key)
                     redirect_map_key = b"span-buf:sr:{%s:%s}" % (project_id, trace_id)
-                    shard = self.assigned_shards[int(trace_id, 16) % len(self.assigned_shards)]
-                    p.zrem(f"span-buf:q:{shard}".encode("ascii"), segment_key)
+                    p.zrem(flushed_segment.queue_key, segment_key)
 
-                    for span_batch in itertools.batched(output_spans, 100):
+                    for span_batch in itertools.batched(flushed_segment.spans, 100):
                         p.hdel(
                             redirect_map_key,
                             *[output_span.payload["span_id"] for output_span in span_batch],

src/sentry/spans/consumers/process/flusher.py

@@ -107,26 +107,22 @@ def produce(payload: KafkaPayload) -> None:
 
             while not stopped.value:
                 now = int(time.time()) + current_drift.value
-
-                queue_size, flushed_segments = buffer.flush_segments(
-                    max_segments=max_flush_segments, now=now
-                )
-                metrics.timing("sentry.spans.buffer.inflight_segments", queue_size)
+                flushed_segments = buffer.flush_segments(max_segments=max_flush_segments, now=now)
 
                 if not flushed_segments:
                     time.sleep(1)
                     continue
 
-                for _, spans_set in flushed_segments.items():
-                    if not spans_set:
+                for _, flushed_segment in flushed_segments.items():
+                    if not flushed_segment.spans:
                         # This is a bug, most likely the input topic is not
                         # partitioned by trace_id so multiple consumers are writing
                         # over each other. The consequence is duplicated segments,
                         # worst-case.
                         metrics.incr("sentry.spans.buffer.empty_segments")
                         continue
 
-                    spans = [span.payload for span in spans_set]
+                    spans = [span.payload for span in flushed_segment.spans]
 
                     kafka_payload = KafkaPayload(
                         None, rapidjson.dumps({"spans": spans}).encode("utf8"), []