use worker_init/worker_process_init to start event processor threads in celery (#444)

Celery uses a prefork model for its worker processes. The app is loaded
in a "worker", which then forks one or more "worker processes". Similiarly
to the uwsgi prefork issues we had to work around, this fork happens
_after_ we created the event processor thread, so the worker processes
don't have an event processor thread of their own, and just fill up the
event queue without it ever being processed.

To work around this, we use worker_init and worker_process_init signals
to start the event processor thread _after_ the fork.

This only became an issue with version 4.2.0 of the agent. In earlier
versions of the agent, threads were started relatively late, after
celery already forked the worker processes.

closes #444

Author: beniwohli

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+[Check the diff](https://github.com/elastic/apm-agent-python/compare/v4.2.1...master)
+ * fixed an issue with Celery and the prefork worker pool (#444)
+
 ## v4.2.1 
 [Check the diff](https://github.com/elastic/apm-agent-python/compare/v4.2.0...v4.2.1)
  * fixed an issue with the certificate pinning feature introduced in 4.2.0 (#433, #434)

elasticapm/contrib/celery/__init__.py

@@ -52,6 +52,7 @@ def process_failure_signal(sender, task_id, exception, args, kwargs, traceback,
 
     signals.task_failure.disconnect(process_failure_signal, dispatch_uid=dispatch_uid)
     signals.task_failure.connect(process_failure_signal, weak=False, dispatch_uid=dispatch_uid)
+    _register_worker_signals(client)
 
 
 def register_instrumentation(client):
@@ -71,3 +72,20 @@ def end_transaction(task_id, task, *args, **kwargs):
     # register for this client
     signals.task_prerun.connect(begin_transaction, dispatch_uid=dispatch_uid % "prerun", weak=False)
     signals.task_postrun.connect(end_transaction, weak=False, dispatch_uid=dispatch_uid % "postrun")
+    _register_worker_signals(client)
+
+
+def _register_worker_signals(client):
+    def worker_startup(*args, **kwargs):
+        client._transport._start_event_processor()
+
+    def worker_shutdown(*args, **kwargs):
+        client.close()
+
+    def connect_worker_process_init(*args, **kwargs):
+        signals.worker_process_init.connect(worker_startup, dispatch_uid="elasticapm-start-worker", weak=False)
+        signals.worker_process_shutdown.connect(worker_shutdown, dispatch_uid="elasticapm-shutdown-worker", weak=False)
+
+    signals.worker_init.connect(
+        connect_worker_process_init, dispatch_uid="elasticapm-connect-start-threads", weak=False
+    )

elasticapm/transport/base.py

@@ -90,8 +90,7 @@ def __init__(
         self._queued_data = None
         self._event_queue = self._init_event_queue(chill_until=queue_chill_count, max_chill_time=queue_chill_time)
         self._is_chilled_queue = isinstance(self._event_queue, ChilledQueue)
-        self._event_process_thread = threading.Thread(target=self._process_queue, name="eapm event processor thread")
-        self._event_process_thread.daemon = True
+        self._event_process_thread = None
         self._last_flush = timeit.default_timer()
         self._counts = defaultdict(int)
         self._flushed = threading.Event()
@@ -212,8 +211,12 @@ def _flush(self, buffer):
                 self.handle_transport_fail(e)
 
     def _start_event_processor(self):
-        if not self._event_process_thread.is_alive() and not self._closed:
+        if (not self._event_process_thread or not self._event_process_thread.is_alive()) and not self._closed:
             try:
+                self._event_process_thread = threading.Thread(
+                    target=self._process_queue, name="eapm event processor thread"
+                )
+                self._event_process_thread.daemon = True
                 self._event_process_thread.start()
             except RuntimeError:
                 pass

tests/contrib/celery/django_tests.py

@@ -35,8 +35,6 @@
 django = pytest.importorskip("django")  # isort:skip
 celery = pytest.importorskip("celery")  # isort:skip
 
-import mock
-
 from elasticapm.conf.constants import ERROR, TRANSACTION
 from elasticapm.contrib.celery import register_exception_tracking, register_instrumentation
 from tests.contrib.django.testapp.tasks import failing_task, successful_task