use worker_init/worker_process_init to start event processor threads in celery

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