diff --git a/.sampo/changesets/resolute-bard-goulven.md b/.sampo/changesets/resolute-bard-goulven.md new file mode 100644 index 000000000..9edc9c3ba --- /dev/null +++ b/.sampo/changesets/resolute-bard-goulven.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Reset the client registry lock after fork diff --git a/posthog/client.py b/posthog/client.py index 9469407bd..1f81eecc1 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -433,6 +433,7 @@ class Client(object): log = logging.getLogger("posthog") _client_registry_lock = threading.Lock() + _client_registry_pid = os.getpid() _client_registry: dict[tuple[str, str], weakref.WeakSet] = {} _duplicate_client_warnings: set[tuple[str, str]] = set() @@ -1826,12 +1827,26 @@ def capture_exception( self.log.exception(f"Failed to capture exception: {e}") return None + @classmethod + def _reinit_client_registry_after_fork(cls): + """Replace the inherited registry lock once in each forked child.""" + child_pid = os.getpid() + if cls._client_registry_pid == child_pid: + return + + # The lock may have been held by a parent thread at fork time. Replace it + # without acquiring it, while preserving inherited active-client records. + cls._client_registry_lock = threading.Lock() + cls._client_registry_pid = child_pid + @staticmethod def _reinit_after_fork_weak(weak_self): """ Reinitialize the client after a fork. Garbage collected if the client is deleted. """ + Client._reinit_client_registry_after_fork() + self = weak_self() if self is None: return diff --git a/posthog/test/test_client_fork.py b/posthog/test/test_client_fork.py index 294dcffb0..bd35cc487 100644 --- a/posthog/test/test_client_fork.py +++ b/posthog/test/test_client_fork.py @@ -238,6 +238,59 @@ def child_probe(): ) self.assertEqual(result, "ok") + def test_register_at_fork_replaces_duplicate_registry_lock_in_child_process(self): + Client._client_registry.clear() + Client._duplicate_client_warnings.clear() + host = "https://fork-registry.example.com" + registry_key = (FAKE_TEST_API_KEY, host) + + with mock.patch("posthog.client.Consumer.start"): + inherited_client = Client(FAKE_TEST_API_KEY, host=host) + inherited_lock = Client._client_registry_lock + + def child_probe(): + # A deadlocked child gets killed by SIGALRM instead of hanging the + # test when constructing or unregistering a client. + signal.alarm(5) + try: + if Client._client_registry_lock is inherited_lock: + return "registry lock was not replaced" + + child_client = Client(FAKE_TEST_API_KEY, host=host) + clients = Client._client_registry.get(registry_key) + if clients is None or set(clients) != { + inherited_client, + child_client, + }: + return "child client was not registered" + + child_client.shutdown() + clients = Client._client_registry.get(registry_key) + if clients is None or set(clients) != {inherited_client}: + return "child client was not unregistered" + + inherited_client.shutdown() + if registry_key in Client._client_registry: + return "inherited client was not unregistered" + finally: + signal.alarm(0) + + return "ok" + + inherited_lock.acquire() + try: + status, result = self._run_fork_probe(child_probe) + finally: + inherited_lock.release() + inherited_client.shutdown() + Client._client_registry.clear() + Client._duplicate_client_warnings.clear() + + self.assertTrue( + os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, msg=result + ) + self.assertEqual(result, "ok") + def test_register_at_fork_replaces_metrics_locks_in_child_process(self): # Locks held at fork time are inherited locked, and their holders don't # exist in the child — the metrics path must not deadlock on them.