address review comments -

- make fork safety complete in the client
- add shutdown mechanism to the integration
- better test coverage
- better docs on usage

Author: parinporecha

examples/celery_integration.py

@@ -2,38 +2,35 @@
 Celery integration example for the PostHog Python SDK.
 
 Demonstrates how to use ``PosthogCeleryIntegration`` with:
-- producer-side instrumentation (publishing events and context propagation)
-- worker-side instrumentation via ``worker_process_init`` (prefork-safe)
+- producer-side and worker-side instrumentation (publishing events and context propagation)
 - context propagation (distinct ID, session ID, tags) from producer to worker
 - task lifecycle events (published, started, success, failure, retry)
 - exception capture from failed tasks
 - ``task_filter`` customization hook
 
 Setup:
-    1. Update POSTHOG_PROJECT_API_KEY and POSTHOG_HOST here with your credentials
-       (environment variables won't work as it's better if Celery forks worker into
-        separate process for the example to prove context propagation)
+    1. Set ``POSTHOG_PROJECT_API_KEY`` and ``POSTHOG_HOST`` in your environment
     2. Install dependencies: pip install posthog celery redis
     3. Start Redis: redis-server
     4. Start the worker: celery -A examples.celery_integration worker --loglevel=info
     5. Run the producer: python -m examples.celery_integration
 """
 
+import os
 import time
 from typing import Any, Optional
 
 from celery import Celery
 from celery.signals import worker_process_init, worker_process_shutdown
 
 import posthog
-from posthog.client import Client
 from posthog.integrations.celery import PosthogCeleryIntegration
 
 
 # --- Configuration ---
 
-POSTHOG_PROJECT_API_KEY = "phc_..."
-POSTHOG_HOST = "http://localhost:8000"
+POSTHOG_PROJECT_API_KEY = os.getenv("POSTHOG_PROJECT_API_KEY", "phc_...")
+POSTHOG_HOST = os.getenv("POSTHOG_HOST", "http://localhost:8000")
 
 app = Celery(
     "examples.celery_integration",
@@ -43,11 +40,11 @@
 
 # --- Integration wiring ---
 
-def create_client() -> Client:
-    return Client(
-        project_api_key=POSTHOG_PROJECT_API_KEY,
-        host=POSTHOG_HOST
-    )
+def configure_posthog() -> None:
+    posthog.api_key = POSTHOG_PROJECT_API_KEY
+    posthog.host = POSTHOG_HOST
+    posthog.enable_local_evaluation = False     # to not require personal_api_key for this example
+    posthog.setup()
 
 
 def task_filter(task_name: Optional[str], task_properties: dict[str, Any]) -> bool:
@@ -56,40 +53,42 @@ def task_filter(task_name: Optional[str], task_properties: dict[str, Any]) -> bo
     return True
 
 
-def create_integration(client: Client) -> PosthogCeleryIntegration:
+def create_integration() -> PosthogCeleryIntegration:
     return PosthogCeleryIntegration(
-        client=client,
         capture_exceptions=True,
         capture_task_lifecycle_events=True,
         propagate_context=True,
         task_filter=task_filter,
     )
 
-
-# Worker process setup.
-# Celery's default prefork pool runs tasks in child processes, so initialize
-# PostHog per child using worker_process_init.
+configure_posthog()
+integration = create_integration()
+integration.instrument()
 
 
+# --- Worker process setup ---
+# Celery's default prefork pool runs tasks in child processes. This example
+# runs on a single host, so the inherited PostHog client and Celery
+# integration are fork-safe and do not need to be recreated in each child.
+# If workers run across multiple hosts, configure PostHog and instrument a
+# worker-local integration in worker_process_init.
 @worker_process_init.connect
 def on_worker_process_init(**kwargs) -> None:
-    worker_posthog_client = create_client()
-    worker_integration = create_integration(worker_posthog_client)
-    worker_integration.instrument()
-    
-    app._posthog_client = worker_posthog_client
-    app._posthog_integration = worker_integration
+    # global integration
+
+    # configure_posthog()
+    # integration = create_integration()
+    # integration.instrument()
+    return
 
 
+# Use this signal to shutdown the integration and PostHog client
+# Calling shutdown() is important to flush any pending events
 @worker_process_shutdown.connect
 def on_worker_process_shutdown(**kwargs) -> None:
-    worker_integration = getattr(app, "_posthog_integration", None)
-    if worker_integration:
-        worker_integration.uninstrument()
+    integration.shutdown()
+    posthog.shutdown()
 
-    worker_posthog_client = getattr(app, "_posthog_client", None)
-    if worker_posthog_client:
-        worker_posthog_client.shutdown()
 
 # --- Example tasks ---
 
@@ -98,8 +97,8 @@ def health_check() -> dict[str, str]:
     return {"status": "ok"}
 
 
-@app.task(bind=True, max_retries=3)
-def process_order(self, order_id: str) -> dict:
+@app.task(max_retries=3)
+def process_order(order_id: str) -> dict:
     """A task that processes an order successfully."""
 
     # simulate work
@@ -108,7 +107,7 @@ def process_order(self, order_id: str) -> dict:
     # Custom event inside the task - context tags propagated from the
     # producer (e.g. "source", "release") should appear on this event
     # and this should be attributed to the correct distinct ID and session.
-    app._posthog_client.capture(
+    posthog.capture(
         "celery example order processed",
         properties={"order_id": order_id, "amount": 99.99},
     )
@@ -136,17 +135,13 @@ def failing_task() -> None:
 # --- Producer code ---
 
 if __name__ == "__main__":
-    posthog_client = create_client()
-    integration = create_integration(posthog_client)
-    integration.instrument()
-
     print("PostHog Celery Integration Example")
     print("=" * 40)
     print()
 
     # Set up PostHog context before dispatching tasks.
     # The integration propagates this context to workers via task headers.
-    with posthog.new_context(fresh=True, client=posthog_client):
+    with posthog.new_context(fresh=True):
         posthog.identify_context("user-123")
         posthog.set_context_session("session-user-123-abc")
         posthog.tag("source", "celery_integration_example_script")
@@ -186,6 +181,5 @@ def failing_task() -> None:
     print("Tasks dispatched. Check your Celery worker logs and PostHog for events.")
     print()
 
-    posthog_client.flush()
-    integration.uninstrument()
-    posthog_client.shutdown()
+    integration.shutdown()
+    posthog.shutdown()

posthog/client.py

@@ -57,6 +57,7 @@
     flags,
     get,
     remote_config,
+    reset_sessions,
 )
 from posthog.types import (
     FeatureFlag,
@@ -245,6 +246,7 @@ def __init__(
         )
         self.poller = None
         self.distinct_ids_feature_flags_reported = SizeLimitedDict(MAX_DICT_SIZE, set)
+        self.flag_fallback_cache_url = flag_fallback_cache_url
         self.flag_cache = self._initialize_flag_cache(flag_fallback_cache_url)
         self.flag_definition_version = 0
         self._flags_etag: Optional[str] = None
@@ -1098,42 +1100,56 @@ def _reinit_after_fork_weak(weak_self):
         self._reinit_after_fork()
 
     def _reinit_after_fork(self):
-        """Reinitialize queue and consumer threads in a forked child process.
+        """Reinitialize fork-unsafe client state in a forked child process.
 
         Registered via os.register_at_fork(after_in_child=...) so it runs
         exactly once in each child, before any user code, covering all code
         paths (capture, flush, join, etc.).
 
         Python threads do not survive fork() and queue.Queue internal locks
-        may be in an inconsistent state, so both are replaced.
-        Inherited queue items are intentionally discarded as they'll be
-        handled by the parent process's consumers.
+        may be in an inconsistent state, so the event queue, consumer threads
+        and other state are replaced. Inherited queue items are not retained
+        as they'll be handled by the parent process's consumers.
         """
-        if self.consumers is None:
-            return
+        if self.consumers:
+            self.queue = queue.Queue(self._max_queue_size)
+
+            new_consumers = []
+            for old in self.consumers:
+                consumer = Consumer(
+                    self.queue,
+                    old.api_key,
+                    flush_at=old.flush_at,
+                    host=old.host,
+                    on_error=old.on_error,
+                    flush_interval=old.flush_interval,
+                    gzip=old.gzip,
+                    retries=old.retries,
+                    timeout=old.timeout,
+                    historical_migration=old.historical_migration,
+                )
+                new_consumers.append(consumer)
+
+                if self.send:
+                    consumer.start()
+
+            self.consumers = new_consumers
 
-        self.queue = queue.Queue(self._max_queue_size)
-
-        new_consumers = []
-        for old in self.consumers:
-            consumer = Consumer(
-                self.queue,
-                old.api_key,
-                flush_at=old.flush_at,
-                host=old.host,
-                on_error=old.on_error,
-                flush_interval=old.flush_interval,
-                gzip=old.gzip,
-                retries=old.retries,
-                timeout=old.timeout,
-                historical_migration=old.historical_migration,
+        if self.enable_local_evaluation:
+            self.poller = Poller(
+                interval=timedelta(seconds=self.poll_interval),
+                execute=self._load_feature_flags,
             )
-            new_consumers.append(consumer)
+            self.poller.start()
+        else:
+            self.poller = None
 
-            if self.send:
-                consumer.start()
+        # If using Redis cache, we must reinitialize to get a fresh connection (fork-safe).
+        # If using Memory cache, we keep it as-is to benefit from the inherited warm cache.
+        if isinstance(self.flag_cache, RedisFlagCache):
+            self.flag_cache = self._initialize_flag_cache(self.flag_fallback_cache_url)
 
-        self.consumers = new_consumers
+        reset_sessions()
 
     def _enqueue(self, msg, disable_geoip):
         # type: (...) -> Optional[str]

posthog/integrations/celery.py

@@ -28,11 +28,10 @@
     integration = PosthogCeleryIntegration(client=posthog)
     integration.instrument()
 
-Both the producer process and each worker process must initialize the
-PostHog client and instrument the integration because the worker needs
-to bind to Celery signals, and the PostHog client may use background threads
-to send captured events (depending on ``sync_mode``). Celery provides a signal
-called ``worker_process_init`` that can be used to accomplish this.
+    # ... publish tasks or run workers ...
+
+    integration.shutdown()
+    posthog.shutdown()
 
 See ``examples/celery_integration.py`` for a complete working example.
 
@@ -64,6 +63,7 @@
     - **retry**: ``celery_reason``
 """
 
+import atexit
 import json
 import logging
 import time
@@ -113,16 +113,28 @@ def __init__(
         self.task_filter = task_filter
 
         self._instrumented = False
+        self._shut_down = False
         self._signals: Optional[Any] = None
         self._celery_version: Optional[str] = None
 
     def instrument(self) -> None:
+        """Connect Celery signal handlers to capture task events and exceptions.
+        Call this after initializing the PostHog client and this integration.
+
+        If Celery runs on a single host, reinstrumenting in worker children is
+        not strictly necessary because the PostHog client and this integration
+        are fork-safe. If Celery workers run across multiple hosts, each worker
+        process must initialize PostHog, this integration, and call
+        ``instrument()``. Celery provides ``worker_process_init`` signal to help
+        with this.
+        """
         if self._instrumented:
             return
 
         from celery import signals
         from celery import __version__ as celery_version
 
+        self._shut_down = False
         self._signals = signals
         self._celery_version = celery_version
 
@@ -133,9 +145,12 @@ def instrument(self) -> None:
         signals.before_task_publish.connect(self._on_before_task_publish, weak=False)
         signals.after_task_publish.connect(self._on_after_task_publish, weak=False)
 
+        signals.worker_process_shutdown.connect(self._on_worker_process_shutdown, weak=False)
+        atexit.register(self.shutdown)
+
         self._instrumented = True
 
-    def uninstrument(self) -> None:
+    def _disconnect_signals(self) -> None:
         if not self._instrumented or not self._signals:
             return
 
@@ -146,9 +161,48 @@ def uninstrument(self) -> None:
         self._signals.before_task_publish.disconnect(self._on_before_task_publish)
         self._signals.after_task_publish.disconnect(self._on_after_task_publish)
 
+        self._signals.worker_process_shutdown.disconnect(self._on_worker_process_shutdown)
+
         self._signals = None
         self._instrumented = False
 
+    def uninstrument(self) -> None:
+        """Disconnect Celery signal handlers and unregister exit cleanup.
+
+        Do not use directly, call `shutdown()` instead.
+        """
+        self._disconnect_signals()
+        atexit.unregister(self.shutdown)
+
+    def shutdown(self) -> None:
+        """Disconnect all signal handlers registered by ``instrument()``, flush all pending events
+        and cleanly shutdown the integration.
+
+        ``shutdown()`` is also registered on ``worker_process_shutdown`` and ``atexit``, but
+        there is no guarantee those will always be called, so we strongly recommend calling
+        it manually when the integration is no longer needed to avoid data loss.
+        """
+        if self._shut_down:
+            return
+
+        try:
+            self._disconnect_signals()
+
+            if self.client:
+                self.client.flush()
+            else:
+                import posthog
+
+                posthog.flush()
+
+            self.uninstrument()
+            self._shut_down = True
+        except Exception:
+            logger.exception("Failed to shut down PostHog Celery integration")
+
+    def _on_worker_process_shutdown(self, *args, **kwargs) -> None:
+        self.shutdown()
+
     def _on_before_task_publish(self, *args, **kwargs):
         try:
             if not self.propagate_context:
@@ -298,7 +352,7 @@ def _handle_task_end(
             if self.capture_task_lifecycle_events and self._should_track(task_name, task_properties):
                 self._capture_event(f"celery task {state}", properties=task_properties)
         except Exception:
-            logger.exception("Failed to process Celery %s", state)
+            logger.exception("Failed to process Celery %s state", state)
         finally:
             ctx = getattr(request, "_posthog_ctx", None)
             if ctx is not None: