beam: circuit_breaker HALF_OPEN single-probe gating

wizard-work/beam/circuit_breaker.py

@@ -47,7 +47,7 @@ class CircuitBreaker:
         self._consecutive_failures = 0
         self._consecutive_successes = 0
         self._opened_at: Optional[float] = None
-        self._half_open_probe: bool = False
+        self._probe_in_flight = False
         self._lock = threading.Lock()
 
     @property
@@ -61,7 +61,6 @@ class CircuitBreaker:
             if self._clock() - self._opened_at >= self._recovery_timeout:
                 self._state = CircuitState.HALF_OPEN
                 self._consecutive_successes = 0
-                self._half_open_probe = False
 
     def _on_success(self) -> None:
         with self._lock:
@@ -73,7 +72,6 @@ class CircuitBreaker:
                     self._consecutive_failures = 0
                     self._consecutive_successes = 0
                     self._opened_at = None
-                self._half_open_probe = False
             else:
                 self._consecutive_failures = 0
 
@@ -85,7 +83,6 @@ class CircuitBreaker:
                 self._opened_at = self._clock()
                 self._consecutive_successes = 0
                 self._consecutive_failures = 0
-                self._half_open_probe = False
                 return
             self._consecutive_failures += 1
             if self._consecutive_failures >= self._failure_threshold:
@@ -99,16 +96,20 @@ class CircuitBreaker:
             if self._state is CircuitState.OPEN:
                 raise CircuitOpenError("circuit is open")
             if self._state is CircuitState.HALF_OPEN:
-                if self._half_open_probe:
-                    raise CircuitOpenError("circuit is half-open (probe in flight)")
-                self._half_open_probe = True
+                if self._probe_in_flight:
+                    raise CircuitOpenError("circuit is open (probe in flight)")
+                self._probe_in_flight = True
         try:
             result = fn(*args, **kwargs)
         except BaseException:
             self._on_failure()
             raise
-        self._on_success()
-        return result
+        else:
+            self._on_success()
+            return result
+        finally:
+            with self._lock:
+                self._probe_in_flight = False
 
     def reset(self) -> None:
         with self._lock:
@@ -116,4 +117,4 @@ class CircuitBreaker:
             self._consecutive_failures = 0
             self._consecutive_successes = 0
             self._opened_at = None
-            self._half_open_probe = False
+            self._probe_in_flight = False