diff --git a/.sampo/changesets/stalwart-bard-goulven.md b/.sampo/changesets/stalwart-bard-goulven.md new file mode 100644 index 000000000..b51a6e3a1 --- /dev/null +++ b/.sampo/changesets/stalwart-bard-goulven.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Reset PostHog context after fork. Forked children no longer retain the parent process's active lexical context; they start without inherited context and can establish a new child-local context. diff --git a/posthog/contexts.py b/posthog/contexts.py index a5ccae376..f897d1d96 100644 --- a/posthog/contexts.py +++ b/posthog/contexts.py @@ -1,4 +1,5 @@ import contextvars +import os from contextlib import contextmanager from typing import Optional, Any, Callable, Dict, TypeVar, cast, TYPE_CHECKING @@ -7,6 +8,9 @@ from posthog.client import Client +_context_generation = 0 + + class ContextScope: def __init__( self, @@ -16,6 +20,7 @@ def __init__( client: Optional["Client"] = None, ): self.client: Optional[Client] = client + self._generation = _context_generation self.parent = parent self.fresh = fresh self.capture_exceptions = capture_exceptions @@ -130,8 +135,25 @@ def get_code_variables_detect_secrets(self) -> Optional[bool]: ) +def _reset_context_after_fork() -> None: + global _context_generation + + _context_generation += 1 + _context_stack.set(None) + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_context_after_fork) + + def _get_current_context() -> Optional[ContextScope]: - return _context_stack.get() + current_context = _context_stack.get() + if ( + current_context is not None + and current_context._generation != _context_generation + ): + return None + return current_context def _default_capture_exceptions(client: Optional["Client"] = None) -> bool: @@ -199,13 +221,17 @@ def new_context( from . import capture_exception current_context = _get_current_context() + context_generation = _context_generation resolved_capture_exceptions = ( capture_exceptions if capture_exceptions is not None else _default_capture_exceptions(client) ) new_context = ContextScope( - current_context, fresh, resolved_capture_exceptions, client + current_context, + fresh, + resolved_capture_exceptions, + client, ) _context_stack.set(new_context) @@ -219,7 +245,8 @@ def new_context( capture_exception(e) raise finally: - _context_stack.set(new_context.get_parent()) + if context_generation == _context_generation: + _context_stack.set(new_context.get_parent()) def tag(key: str, value: Any) -> None: diff --git a/posthog/test/test_contexts.py b/posthog/test/test_contexts.py index 4c7aa5577..e75a6610d 100644 --- a/posthog/test/test_contexts.py +++ b/posthog/test/test_contexts.py @@ -1,4 +1,6 @@ import asyncio +import contextvars +import os import unittest from unittest.mock import MagicMock, patch @@ -11,8 +13,10 @@ tag, identify_context, set_context_session, + set_context_device_id, get_context_session_id, get_context_distinct_id, + get_context_device_id, ) @@ -275,6 +279,111 @@ def test_context_inheritance_non_fresh_context(self): assert get_context_distinct_id() == "user123" assert get_context_session_id() == "session456" + @unittest.skipUnless( + hasattr(os, "fork") and hasattr(os, "register_at_fork"), + "requires os.fork and os.register_at_fork", + ) + def test_fork_clears_context_in_child_and_preserves_parent(self): + def context_state(): + return ( + get_context_distinct_id(), + get_context_session_id(), + get_context_device_id(), + get_tags(), + ) + + read_fd, write_fd = os.pipe() + pid = -1 + child_result = b"" + child_exit_code = 1 + try: + with new_context(fresh=True): + identify_context("parent-user") + set_context_session("parent-session") + set_context_device_id("parent-device") + tag("parent-tag", "parent-value") + + with new_context(): + copied_context = contextvars.copy_context() + pid = os.fork() + if pid == 0: + os.close(read_fd) + with new_context(): + identify_context("child-user") + set_context_session("child-session") + set_context_device_id("child-device") + tag("child-tag", "child-value") + child_local_state = context_state() + child_state_after_local_scope = context_state() + else: + os.close(write_fd) + + if pid == 0: + child_state_after_inner_scope = context_state() + else: + parent_state = context_state() + + if pid == 0: + child_states = ( + child_local_state, + child_state_after_local_scope, + child_state_after_inner_scope, + context_state(), + copied_context.run(context_state), + ) + child_result = repr(child_states).encode() + child_exit_code = 0 + except BaseException as error: + if pid != 0: + raise + child_result = f"{type(error).__name__}: {error}".encode() + finally: + if pid == 0: + try: + os.write(write_fd, child_result) + finally: + try: + os.close(write_fd) + finally: + os._exit(child_exit_code) + + child_states = os.read(read_fd, 4096) + os.close(read_fd) + _, status = os.waitpid(pid, 0) + + empty_state = (None, None, None, {}) + self.assertTrue( + os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, + child_states.decode(errors="replace"), + ) + child_local_state = ( + "child-user", + "child-session", + "child-device", + {"child-tag": "child-value"}, + ) + self.assertEqual( + child_states, + repr( + ( + child_local_state, + empty_state, + empty_state, + empty_state, + empty_state, + ) + ).encode(), + ) + self.assertEqual( + parent_state, + ( + "parent-user", + "parent-session", + "parent-device", + {"parent-tag": "parent-value"}, + ), + ) + def test_child_tags_override_parent_tags_in_non_fresh_context(self): with new_context(fresh=True): tag("shared_key", "parent_value")