beam: circuit breaker half-open single-probe fix

wizard-work/beam/circuit_breaker.py

@@ -47,6 +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._lock = threading.Lock()
 
     @property
@@ -60,6 +61,7 @@ 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:
@@ -71,6 +73,7 @@ class CircuitBreaker:
                     self._consecutive_failures = 0
                     self._consecutive_successes = 0
                     self._opened_at = None
+                self._half_open_probe = False
             else:
                 self._consecutive_failures = 0
 
@@ -82,6 +85,7 @@ 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:
@@ -94,6 +98,10 @@ class CircuitBreaker:
             self._maybe_transition_to_half_open()
             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
         try:
             result = fn(*args, **kwargs)
         except BaseException:
@@ -107,4 +115,5 @@ class CircuitBreaker:
             self._state = CircuitState.CLOSED
             self._consecutive_failures = 0
             self._consecutive_successes = 0
-            self._opened_at = None
+            self._opened_at = None
+            self._half_open_probe = False