Address code review comments

Author: ericapisani

sentry_sdk/_batcher.py

@@ -51,8 +51,11 @@ def _reset_in_child() -> None:
             os.register_at_fork(after_in_child=_reset_in_child)
 
     def _reset_thread_state(self) -> None:
-        self._flusher = None
+        self._buffer = []
         self._lock = threading.Lock()
+        self._active = threading.local()
+        self._flush_event = threading.Event()
+        self._flusher = None
         self._flusher_pid = None
 
     def _ensure_thread(self) -> bool:

sentry_sdk/_span_batcher.py

@@ -63,6 +63,18 @@ def _reset_in_child() -> None:
 
             os.register_at_fork(after_in_child=_reset_in_child)
 
+    def _reset_thread_state(self) -> None:
+        self._span_buffer = defaultdict(list)
+        self._running_size = defaultdict(lambda: 0)
+
+        self._lock = threading.Lock()
+        self._active = threading.local()
+
+        self._flush_event = threading.Event()
+
+        self._flusher = None
+        self._flusher_pid = None
+
     def add(self, span: "StreamedSpan") -> None:
         # Bail out if the current thread is already executing batcher code.
         # This prevents deadlocks when code running inside the batcher (e.g.

tests/test_logs.py

@@ -835,14 +835,18 @@ def test_log_batcher_lock_reset_in_child_after_fork(sentry_init):
     child inherits the lock locked. The holding thread does not exist in
     the child, so the lock can never be released and _ensure_thread
     deadlocks forever. The after-fork hook must replace the lock with a
-    fresh one in the child and reset _flusher / _flusher_pid.
+    fresh one in the child and reset
+    _flusher / _flusher_pid / _buffer / _active / _flush_event.
     """
     sentry_init(enable_logs=True)
     batcher = sentry_sdk.get_client().log_batcher
     assert batcher is not None
 
     original_lock = batcher._lock
     original_lock.acquire()
+    batcher._buffer.append(object())
+    batcher._active.flag = True
+    batcher._flush_event.set()
     pid = os.fork()
     if pid == 0:
         # Child: was the lock object replaced and is the new one not
@@ -852,7 +856,19 @@ def test_log_batcher_lock_reset_in_child_after_fork(sentry_init):
         replaced = batcher._lock is not original_lock
         unheld = batcher._lock.acquire(blocking=False)
         flusher_reset = batcher._flusher is None and batcher._flusher_pid is None
-        os._exit(0 if replaced and unheld and flusher_reset else 1)
+        buffer_reset = len(batcher._buffer) == 0
+        active_reset = not getattr(batcher._active, "flag", False)
+        event_reset = not batcher._flush_event.is_set()
+        os._exit(
+            0
+            if replaced
+            and unheld
+            and flusher_reset
+            and buffer_reset
+            and active_reset
+            and event_reset
+            else 1
+        )
 
     original_lock.release()
     _, status = os.waitpid(pid, 0)

tests/test_metrics.py

@@ -531,24 +531,39 @@ def test_metrics_batcher_lock_reset_in_child_after_fork(sentry_init):
     exist in the child, so the lock can never be released and
     _ensure_thread deadlocks forever. The after-fork hook must replace
     the lock with a fresh one in the child and reset
-    _flusher / _flusher_pid.
+    _flusher / _flusher_pid / _buffer / _active / _flush_event.
     """
     sentry_init()
     batcher = sentry_sdk.get_client().metrics_batcher
     assert batcher is not None
 
     original_lock = batcher._lock
     original_lock.acquire()
+
+    batcher._buffer.append(object())
+    batcher._active.flag = True
+    batcher._flush_event.set()
+
     pid = os.fork()
     if pid == 0:
-        # Child: was the lock object replaced and is the new one not
-        # held? Without the fix, _lock is `original_lock` inherited
-        # locked, so `replaced` is False. blocking=False guarantees the
-        # child can't hang on a regression.
         replaced = batcher._lock is not original_lock
         unheld = batcher._lock.acquire(blocking=False)
+
         flusher_reset = batcher._flusher is None and batcher._flusher_pid is None
-        os._exit(0 if replaced and unheld and flusher_reset else 1)
+        buffer_reset = len(batcher._buffer) == 0
+        active_reset = not getattr(batcher._active, "flag", False)
+        event_reset = not batcher._flush_event.is_set()
+
+        os._exit(
+            0
+            if replaced
+            and unheld
+            and flusher_reset
+            and buffer_reset
+            and active_reset
+            and event_reset
+            else 1
+        )
 
     original_lock.release()
     _, status = os.waitpid(pid, 0)

tests/tracing/test_span_streaming.py

@@ -1570,7 +1570,9 @@ def test_span_batcher_lock_reset_in_child_after_fork(sentry_init):
     child inherits the lock locked. The holding thread does not exist in
     the child, so the lock can never be released and _ensure_thread
     deadlocks forever. The after-fork hook must replace the lock with a
-    fresh one in the child and reset _flusher / _flusher_pid.
+    fresh one in the child and reset
+    _flusher / _flusher_pid / _span_buffer / _running_size / _active /
+    _flush_event.
     """
     sentry_init(
         traces_sample_rate=1.0,
@@ -1581,16 +1583,35 @@ def test_span_batcher_lock_reset_in_child_after_fork(sentry_init):
 
     original_lock = batcher._lock
     original_lock.acquire()
+
+    batcher._span_buffer["test-trace-id"].append(object())
+    batcher._running_size["test-trace-id"] = 42
+    batcher._active.flag = True
+    batcher._flush_event.set()
+
     pid = os.fork()
     if pid == 0:
-        # Child: was the lock object replaced and is the new one not
-        # held? Without the fix, _lock is `original_lock` inherited
-        # locked, so `replaced` is False. blocking=False guarantees the
-        # child can't hang on a regression.
         replaced = batcher._lock is not original_lock
         unheld = batcher._lock.acquire(blocking=False)
+
         flusher_reset = batcher._flusher is None and batcher._flusher_pid is None
-        os._exit(0 if replaced and unheld and flusher_reset else 1)
+        span_buffer_reset = len(batcher._span_buffer) == 0
+        running_size_reset = len(batcher._running_size) == 0
+
+        active_reset = not getattr(batcher._active, "flag", False)
+        event_reset = not batcher._flush_event.is_set()
+
+        os._exit(
+            0
+            if replaced
+            and unheld
+            and flusher_reset
+            and span_buffer_reset
+            and running_size_reset
+            and active_reset
+            and event_reset
+            else 1
+        )
 
     original_lock.release()
     _, status = os.waitpid(pid, 0)