Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/sentry/consumers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ def uptime_options() -> list[click.Option]:
options = [
click.Option(
["--mode", "mode"],
type=click.Choice(["serial", "parallel", "batched-parallel"]),
type=click.Choice(["serial", "parallel", "batched-parallel", "thread-queue-parallel"]),
default="serial",
help="The mode to process results in. Parallel uses multithreading.",
),
Expand All @@ -138,7 +138,7 @@ def uptime_options() -> list[click.Option]:
["--max-workers", "max_workers"],
type=int,
default=None,
help="The maximum number of threads to spawn in parallel mode.",
help="The maximum amount of parallelism to use when in a parallel mode.",
),
click.Option(["--processes", "num_processes"], default=1, type=int),
click.Option(["--input-block-size"], type=int, default=None),
Expand Down
345 changes: 345 additions & 0 deletions src/sentry/remote_subscriptions/consumers/queue_consumer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,345 @@
from __future__ import annotations

import logging
import queue
import threading
import time
from collections import defaultdict
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Generic, TypeVar

import sentry_sdk
from arroyo.backends.kafka.consumer import KafkaPayload
from arroyo.processing.strategies import ProcessingStrategy
from arroyo.types import BrokerValue, FilteredPayload, Message, Partition

from sentry.utils import metrics

logger = logging.getLogger(__name__)

T = TypeVar("T")


@dataclass
class WorkItem(Generic[T]):
"""Work item that includes the original message for offset tracking."""

partition: Partition
offset: int
result: T
message: Message[KafkaPayload | FilteredPayload]


class OffsetTracker:
"""
Tracks outstanding offsets and determines which offsets are safe to commit.

- Tracks offsets per partition
- Only commits offsets when all prior offsets are processed
- Thread-safe for concurrent access with per-partition locks
"""

def __init__(self) -> None:
self.all_offsets: dict[Partition, set[int]] = defaultdict(set)
self.outstanding: dict[Partition, set[int]] = defaultdict(set)
self.last_committed: dict[Partition, int] = {}
self.partition_locks: dict[Partition, threading.Lock] = {}

def _get_partition_lock(self, partition: Partition) -> threading.Lock:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

锁获取竞态条件。_get_partition_lock 方法中,多线程同时调用时可能创建多个锁对象。line 49-51 处先检查 lock 是否存在,若不存在则在 line 52 调用 setdefault。两个线程可能同时通过检查,各自创建锁对象,导致线程安全保证失效。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 49

💬 详细说明:

  • 破坏 OffsetTracker 的线程安全性,可能导致偏移量追踪数据损坏,引发消息重复消费或丢失。

📝 问题代码:

        lock = self.partition_locks.get(partition)
        if lock:
            return lock
        return self.partition_locks.setdefault(partition, threading.Lock())

💡 修复建议:

使用单例模式或双重检查锁定确保每个 partition 只创建一个锁对象。推荐使用 threading.Lock() 配合字典原子操作。

✅ 修复示例:

    def _get_partition_lock(self, partition: Partition) -> threading.Lock:
        """Get or create a lock for a partition - thread-safe implementation."""
        lock = self.partition_locks.get(partition)
        if lock is None:
            with threading.Lock():
                lock = self.partition_locks.get(partition)
                if lock is None:
                    lock = threading.Lock()
                    self.partition_locks[partition] = lock
        return lock

🔗 参考链接

"""Get or create a lock for a partition."""
lock = self.partition_locks.get(partition)
if lock:
return lock
return self.partition_locks.setdefault(partition, threading.Lock())

def add_offset(self, partition: Partition, offset: int) -> None:
"""Record that we've started processing an offset."""
with self._get_partition_lock(partition):
self.all_offsets[partition].add(offset)
self.outstanding[partition].add(offset)

def complete_offset(self, partition: Partition, offset: int) -> None:
"""Mark an offset as completed."""
with self._get_partition_lock(partition):
self.outstanding[partition].discard(offset)

def get_committable_offsets(self) -> dict[Partition, int]:
"""
Get the highest offset per partition that can be safely committed.

For each partition, finds the highest contiguous offset that has been processed.
"""
committable = {}
for partition in list(self.all_offsets.keys()):
with self._get_partition_lock(partition):
all_offsets = self.all_offsets[partition]
if not all_offsets:
continue

outstanding = self.outstanding[partition]
last_committed = self.last_committed.get(partition, -1)

min_offset = min(all_offsets)
max_offset = max(all_offsets)

start = max(last_committed + 1, min_offset)

highest_committable = last_committed
for offset in range(start, max_offset + 1):
if offset in all_offsets and offset not in outstanding:
highest_committable = offset
else:
break

if highest_committable > last_committed:
committable[partition] = highest_committable

return committable

def mark_committed(self, partition: Partition, offset: int) -> None:
"""Update the last committed offset for a partition."""
with self._get_partition_lock(partition):
self.last_committed[partition] = offset
# Remove all offsets <= committed offset
self.all_offsets[partition] = {o for o in self.all_offsets[partition] if o > offset}


class OrderedQueueWorker(threading.Thread, Generic[T]):
"""Worker thread that processes items from a queue in order."""

def __init__(
self,
worker_id: int,
work_queue: queue.Queue[WorkItem[T]],
result_processor: Callable[[str, T], None],
identifier: str,
offset_tracker: OffsetTracker,
) -> None:
super().__init__(daemon=True)
self.worker_id = worker_id
self.work_queue = work_queue
self.result_processor = result_processor
self.identifier = identifier
self.offset_tracker = offset_tracker
self.shutdown = False

def run(self) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

worker 线程在队列为空时无限阻塞
line 127-133 的 OrderedQueueWorker.run 方法使用 work_queue.get() 不带 timeout 参数。当 shutdown 标志被设置后,如果队列中仍有未处理的消息,worker 会正确处理;但如果队列为空,worker 将永远阻塞在 get() 调用上,无法响应 shutdown。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 127

💬 详细说明:

  • 在关闭过程中,如果队列恰好为空,worker 线程将永远阻塞,导致 join(timeout=5.0) 超时,线程无法正确终止,可能引发资源泄漏或优雅关闭失败。

📝 问题代码:

    def run(self) -> None:
        """Process items from the queue in order."""
        while not self.shutdown:
            try:
                work_item = self.work_queue.get()

💡 修复建议:

使用带超时的 get() 方法,使 worker 能够定期检查 shutdown 标志。建议修改为:

work_item = self.work_queue.get(timeout=0.5)

这样可以定期轮询 shutdown 状态。

✅ 修复示例:

    def run(self) -> None:
        """Process items from the queue in order."""
        while not self.shutdown:
            try:
                work_item = self.work_queue.get(timeout=0.5)
            except queue.Empty:
                continue
            except queue.ShutDown:
                break

🔗 参考链接

"""Process items from the queue in order."""
while not self.shutdown:
try:
work_item = self.work_queue.get()
except queue.ShutDown:
break

try:
with sentry_sdk.start_transaction(
op="queue_worker.process",
name=f"monitors.{self.identifier}.worker_{self.worker_id}",
):
self.result_processor(self.identifier, work_item.result)

except queue.ShutDown:
break
except Exception:
logger.exception(
"Unexpected error in queue worker", extra={"worker_id": self.worker_id}
)
finally:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

UnboundLocalError 风险。finally 块 line 148 直接访问 work_item,如果 line 127 的 get() 抛出非 ShutDown 异常(如队列已关闭),work_item 变量未定义。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 148

💬 详细说明:

  • 异常处理代码自身抛出异常,掩盖原始错误,增加调试难度。

📝 问题代码:

            finally:
                self.offset_tracker.complete_offset(work_item.partition, work_item.offset)
                metrics.gauge(

💡 修复建议:

在 try 块之前初始化 work_item = None,在 finally 中检查是否为 None 再访问。

✅ 修复示例:

    def run(self) -> None:
        """Process items from the queue in order."""
        work_item = None
        while not self.shutdown:
            try:
                work_item = self.work_queue.get(timeout=0.1)
            except queue.Empty:
                continue
            except queue.ShutDown:
                break
            except Exception:
                logger.exception("Error getting work item")
                continue

            if work_item is None:
                continue

            try:
                with sentry_sdk.start_transaction(
                    op="queue_worker.process",
                    name=f"monitors.{self.identifier}.worker_{self.worker_id}",
                ):
                    self.result_processor(self.identifier, work_item.result)

            except queue.ShutDown:
                break
            except Exception:
                logger.exception(
                    "Unexpected error in queue worker", extra={"worker_id": self.worker_id}
                )
            finally:
                if work_item is not None:
                    self.offset_tracker.complete_offset(work_item.partition, work_item.offset)
                    work_item = None

🔗 参考链接

self.offset_tracker.complete_offset(work_item.partition, work_item.offset)
metrics.gauge(
"remote_subscriptions.queue_worker.queue_depth",
self.work_queue.qsize(),
tags={
"identifier": self.identifier,
},
)


class FixedQueuePool(Generic[T]):
"""
Fixed pool of queues that guarantees order within groups.

Key properties:
- Each group is consistently assigned to the same queue
- Each queue has exactly one worker thread
- Items within a queue are processed in FIFO order
- No dynamic reassignment that could break ordering
- Tracks offset completion for safe commits
"""

def __init__(
self,
result_processor: Callable[[str, T], None],
identifier: str,
num_queues: int = 20,
) -> None:
self.result_processor = result_processor
self.identifier = identifier
self.num_queues = num_queues
self.offset_tracker = OffsetTracker()
self.queues: list[queue.Queue[WorkItem[T]]] = []
self.workers: list[OrderedQueueWorker[T]] = []

for i in range(num_queues):
work_queue: queue.Queue[WorkItem[T]] = queue.Queue()
self.queues.append(work_queue)

worker = OrderedQueueWorker[T](
worker_id=i,
work_queue=work_queue,
result_processor=result_processor,
identifier=identifier,
offset_tracker=self.offset_tracker,
)
worker.start()
self.workers.append(worker)

def get_queue_for_group(self, group_key: str) -> int:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

Python 内置 hash() 函数存在哈希随机化问题
line 202 使用 hash(group_key) % self.num_queues 进行分组键哈希。Python 3.3+ 默认启用哈希随机化(PYTHONHASHSEED),每次进程重启后相同字符串的哈希值不同,导致同一 group_key 在不同进程生命周期中被分配到不同队列,破坏同一组消息的有序性保证。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 198

💬 详细说明:

  • 在消费者重启后,同一分组的消息可能被路由到不同队列,如果旧队列中仍有未处理的消息,会导致同一分组的消息在不同队列中并行处理,破坏有序性保证,可能引发竞态条件或数据不一致。

📝 问题代码:

    def get_queue_for_group(self, group_key: str) -> int:
        """
        Get queue index for a group using consistent hashing.
        """
        return hash(group_key) % self.num_queues

💡 修复建议:

使用稳定的哈希算法替代内置 hash()。建议使用 hashlib.md5 或类似的确定性哈希:

import hashlib

def get_queue_for_group(self, group_key: str) -> int:
    return int(hashlib.md5(group_key.encode()).hexdigest(), 16) % self.num_queues

✅ 修复示例:

    def get_queue_for_group(self, group_key: str) -> int:
        """
        Get queue index for a group using consistent hashing.
        """
        import hashlib
        return int(hashlib.md5(group_key.encode()).hexdigest(), 16) % self.num_queues

🔗 参考链接

"""
Get queue index for a group using consistent hashing.
"""
return hash(group_key) % self.num_queues

def submit(self, group_key: str, work_item: WorkItem[T]) -> None:
"""
Submit a work item to the appropriate queue.
"""
queue_index = self.get_queue_for_group(group_key)
work_queue = self.queues[queue_index]

self.offset_tracker.add_offset(work_item.partition, work_item.offset)
work_queue.put(work_item)

def get_stats(self) -> dict[str, Any]:
"""Get statistics about queue depths."""
queue_depths = [q.qsize() for q in self.queues]
return {
"queue_depths": queue_depths,
"total_items": sum(queue_depths),
}

def wait_until_empty(self, timeout: float = 5.0) -> bool:
"""Wait until all queues are empty. Returns True if successful, False if timeout."""
start_time = time.time()
while time.time() - start_time < timeout:
if self.get_stats()["total_items"] == 0:
return True
time.sleep(0.01)
return False

def shutdown(self) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

shutdown 竞态条件和异常处理不完善。line 231-237 关闭顺序:先设置 worker.shutdown 标志,再关闭队列,最后 join workers。但 worker 可能在队列关闭后才尝试 get(),抛出异常。重复调用 shutdown 会抛出 AttributeError。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 231

💬 详细说明:

  • 关闭流程不稳定,可能抛出未捕获异常,影响优雅关闭。

📝 问题代码:

    def shutdown(self) -> None:
        """Gracefully shutdown all workers."""
        for worker in self.workers:
            worker.shutdown = True

        for q in self.queues:
            try:
                q.shutdown(immediate=False)
            except Exception:
                logger.exception("Error shutting down queue")

        for worker in self.workers:
            worker.join(timeout=5.0)

💡 修复建议:

添加关闭状态标志防止重复关闭,改进关闭顺序:先关闭队列(让 worker 的 get() 抛出 ShutDown),再 join workers。

✅ 修复示例:

    def shutdown(self) -> None:
        """Gracefully shutdown all workers - idempotent."""
        if hasattr(self, '_shutdown_complete') and self._shutdown_complete:
            return

        # First shutdown queues to wake up workers
        for q in self.queues:
            try:
                q.shutdown(immediate=False)
            except Exception:
                logger.exception("Error shutting down queue")

        # Set shutdown flags
        for worker in self.workers:
            worker.shutdown = True

        # Wait for workers to finish
        for worker in self.workers:
            worker.join(timeout=5.0)
            if worker.is_alive():
                logger.warning(f"Worker {worker.worker_id} did not stop gracefully")

        self._shutdown_complete = True

🔗 参考链接

"""Gracefully shutdown all workers."""
for worker in self.workers:
worker.shutdown = True

for q in self.queues:
try:
q.shutdown(immediate=False)
except Exception:
logger.exception("Error shutting down queue")

for worker in self.workers:
worker.join(timeout=5.0)


class SimpleQueueProcessingStrategy(ProcessingStrategy[KafkaPayload], Generic[T]):
"""
Processing strategy that uses a fixed pool of queues.

Guarantees:
- Items for the same group are processed in order
- No item is lost or processed out of order
- Natural backpressure when queues fill up
- Only commits offsets after successful processing
"""

def __init__(
self,
queue_pool: FixedQueuePool[T],
decoder: Callable[[KafkaPayload | FilteredPayload], T | None],
grouping_fn: Callable[[T], str],
commit_function: Callable[[dict[Partition, int]], None],
) -> None:
self.queue_pool = queue_pool
self.decoder = decoder
self.grouping_fn = grouping_fn
self.commit_function = commit_function
self.shutdown_event = threading.Event()

self.commit_thread = threading.Thread(target=self._commit_loop, daemon=True)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

commit_thread 作为 daemon 线程可能在关闭时丢失偏移量提交
line 270-271 的提交线程被设置为 daemon=True。daemon 线程在主线程退出时会立即被终止,不等待其完成。如果在关闭时有待提交的偏移量正在处理中,daemon 线程被强制终止可能导致偏移量提交丢失。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 270

💬 详细说明:

  • 在消费者关闭时,可能丢失已处理但未提交的偏移量,导致消息重复消费;如果业务处理是幂等的,重复消费可能问题不大,但会增加系统负载并可能导致数据不一致。

📝 问题代码:

        self.commit_thread = threading.Thread(target=self._commit_loop, daemon=True)
        self.commit_thread.start()

💡 修复建议:

将 commit_thread 改为非 daemon 线程,并在 close() 方法中确保正确等待其完成(当前代码已调用 join,但 daemon 线程的行为可能导致 join 不可靠)。建议修改为:

self.commit_thread = threading.Thread(target=self._commit_loop, daemon=False)

✅ 修复示例:

        self.commit_thread = threading.Thread(target=self._commit_loop, daemon=False)
        self.commit_thread.start()

🔗 参考链接

self.commit_thread.start()

def _commit_loop(self) -> None:
while not self.shutdown_event.is_set():
try:
self.shutdown_event.wait(1.0)

committable = self.queue_pool.offset_tracker.get_committable_offsets()

if committable:
metrics.incr(
"remote_subscriptions.queue_pool.offsets_committed",
len(committable),
tags={"identifier": self.queue_pool.identifier},
)

self.commit_function(committable)
for partition, offset in committable.items():
self.queue_pool.offset_tracker.mark_committed(partition, offset)
except Exception:
logger.exception("Error in commit loop")

def submit(self, message: Message[KafkaPayload | FilteredPayload]) -> None:
try:
result = self.decoder(message.payload)

assert isinstance(message.value, BrokerValue)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

assert 用于生产环境运行时检查。line 297 使用 assert 检查 message.value 是否为 BrokerValue 类型,但在 Python 优化模式(-O)下 assert 会被移除,导致类型检查失效。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 297

💬 详细说明:

  • 生产环境可能因类型错误导致崩溃,且难以调试。

📝 问题代码:

            assert isinstance(message.value, BrokerValue)
            partition = message.value.partition
            offset = message.value.offset

💡 修复建议:

使用显式的 if 判断和异常抛出替代 assert。

✅ 修复示例:

            if not isinstance(message.value, BrokerValue):
                raise TypeError(f"Expected BrokerValue, got {type(message.value).__name__}")
            partition = message.value.partition
            offset = message.value.offset

🔗 参考链接

partition = message.value.partition
offset = message.value.offset

if result is None:
self.queue_pool.offset_tracker.add_offset(partition, offset)
self.queue_pool.offset_tracker.complete_offset(partition, offset)
return

group_key = self.grouping_fn(result)

work_item = WorkItem(
partition=partition,
offset=offset,
result=result,
message=message,
)

self.queue_pool.submit(group_key, work_item)

except Exception:
logger.exception("Error submitting message to queue")
if isinstance(message.value, BrokerValue):
self.queue_pool.offset_tracker.add_offset(
message.value.partition, message.value.offset
)
self.queue_pool.offset_tracker.complete_offset(
message.value.partition, message.value.offset
)

def poll(self) -> None:
stats = self.queue_pool.get_stats()
metrics.gauge(
"remote_subscriptions.queue_pool.total_queued",
stats["total_items"],
tags={"identifier": self.queue_pool.identifier},
)

def close(self) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

close() 方法中 commit_thread.join 失败未处理
line 335-338 的 close 方法调用 commit_thread.join(timeout=5.0),如果超时(线程未在 5 秒内完成),join 返回但不会抛出异常,代码继续执行 queue_pool.shutdown()。此时提交线程可能仍在运行,访问正在被关闭的资源,引发竞态条件。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 335

💬 详细说明:

  • 提交线程可能在 queue_pool 关闭后仍在尝试访问 offset_tracker,导致异常或数据不一致;未完成的偏移量提交可能丢失。

📝 问题代码:

    def close(self) -> None:
        self.shutdown_event.set()
        self.commit_thread.join(timeout=5.0)
        self.queue_pool.shutdown()

💡 修复建议:

检查 join 是否成功,并处理超时情况:

def close(self) -> None:
    self.shutdown_event.set()
    self.commit_thread.join(timeout=5.0)
    if self.commit_thread.is_alive():
        logger.warning("Commit thread did not stop in time, proceeding with shutdown")
    self.queue_pool.shutdown()

✅ 修复示例:

    def close(self) -> None:
        self.shutdown_event.set()
        self.commit_thread.join(timeout=5.0)
        if self.commit_thread.is_alive():
            logger.warning("Commit thread did not stop in time, proceeding with shutdown")
        self.queue_pool.shutdown()

🔗 参考链接

self.shutdown_event.set()
self.commit_thread.join(timeout=5.0)
self.queue_pool.shutdown()

def terminate(self) -> None:
self.shutdown_event.set()
self.queue_pool.shutdown()

def join(self, timeout: float | None = None) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 AI 代码审查发现问题

📋 问题概述

join 方法未使用 timeout 参数
line 344-345 的 join 方法接收 timeout 参数但直接调用 self.close(),未将 timeout 传递给 close 或内部调用。这违反了 ProcessingStrategy 接口契约,调用者期望能够控制等待时间,但实际被忽略。

📍 问题详情

🟡 问题 1 | 严重程度: MEDIUM | 行号: 344

💬 详细说明:

  • 调用者无法通过 timeout 参数控制关闭等待时间,可能导致调用方期望的超时行为与实际不符,影响优雅关闭流程。

📝 问题代码:

    def join(self, timeout: float | None = None) -> None:
        self.close()

💡 修复建议:

在 join 方法中传递或使用 timeout 参数:

def join(self, timeout: float | None = None) -> None:
    # close 已经包含 join,但我们需要尊重 timeout 参数
    self.shutdown_event.set()
    self.commit_thread.join(timeout=timeout)
    self.queue_pool.wait_until_empty(timeout=timeout if timeout else 5.0)

✅ 修复示例:

    def join(self, timeout: float | None = None) -> None:
        self.shutdown_event.set()
        self.commit_thread.join(timeout=timeout)
        self.queue_pool.wait_until_empty(timeout=timeout if timeout else 5.0)

🔗 参考链接

self.close()
Loading