fix: don't KeyError on statistics when begin() ran on another thread (#507)

releasenotes/notes/statistics-thread-local-keys-3f8a1c2d94b7e601.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    Fixed a :exc:`KeyError` on ``statistics["idle_for"]`` when ``begin()``
+    ran on a different thread than the retry loop (for example under
+    Temporal's replay worker): ``next_action()`` now reads ``idle_for`` and
+    ``attempt_number`` with defaults instead of assuming they were already
+    initialized on the current thread. See issue #507.

tenacity/__init__.py

@@ -487,8 +487,13 @@ def next_action(rs: "RetryCallState") -> None:
             sleep = rs.upcoming_sleep
             rs.next_action = RetryAction(sleep)
             rs.idle_for += sleep
-            self.statistics["idle_for"] += sleep
-            self.statistics["attempt_number"] += 1
+            # `begin()` may have run on a different thread than this action
+            # (e.g. Temporal's replay worker re-executing the workflow), in
+            # which case this thread's lazily created statistics dict starts
+            # empty. Read with defaults instead of raising KeyError (#507).
+            stats = self.statistics
+            stats["idle_for"] = stats.get("idle_for", 0) + sleep
+            stats["attempt_number"] = stats.get("attempt_number", 1) + 1
 
         self._add_action_func(next_action)
 

tests/test_tenacity.py

@@ -18,6 +18,7 @@
 import logging
 import pickle
 import re
+import threading
 import time
 import typing
 import unittest
@@ -1620,6 +1621,44 @@ def succeeds_first_try() -> bool:
         succeeds_first_try()
         assert succeeds_first_try.statistics["delay_since_first_attempt"] == 0
 
+    def test_iter_on_fresh_thread_after_begin_elsewhere(self) -> None:
+        """iter()/next_action() must not KeyError when begin() ran elsewhere.
+
+        Reproduces the Temporal scenario from issue #507: begin() runs on one
+        thread, but the retry loop is re-executed on another (replay worker).
+        That thread's lazily created statistics dict starts empty, so
+        `self.statistics["idle_for"] += sleep` used to raise KeyError.
+        """
+        r = tenacity.Retrying(
+            stop=tenacity.stop_after_attempt(5), wait=tenacity.wait_fixed(0)
+        )
+        r.begin()  # simulate begin() having run on a different thread
+
+        observed: dict[str, typing.Any] = {}
+
+        def worker() -> None:
+            retry_state = tenacity.RetryCallState(
+                r, fn=lambda: None, args=(), kwargs={}
+            )
+            retry_state.set_exception((ValueError, ValueError("boom"), None))
+            try:
+                do = r.iter(retry_state=retry_state)
+                observed["do"] = type(do).__name__
+                stats = r.statistics
+                observed["idle_for"] = stats["idle_for"]
+                observed["attempt_number"] = stats["attempt_number"]
+            except Exception as exc:
+                observed["error"] = exc
+
+        thread = threading.Thread(target=worker)
+        thread.start()
+        thread.join()
+
+        assert "error" not in observed, observed["error"]
+        assert observed["do"] == "DoSleep"
+        assert observed["idle_for"] == 0
+        assert observed["attempt_number"] == 2
+
     def test_statistics_visible_through_outer_decorator(self) -> None:
         """Statistics must resolve when @retry is wrapped by another decorator.