Merge pull request #2 from MRoci/master

Add CustomState instead of celery.event.State & refac task collection logic

Author: SharpEdgeMarshall

celery_exporter/monitor.py

@@ -9,6 +9,7 @@
 import celery.events
 
 from .metrics import TASKS, TASKS_RUNTIME, LATENCY, WORKERS
+from .state import CeleryState
 
 class TaskThread(threading.Thread):
     """
@@ -20,7 +21,7 @@ def __init__(self, app, namespace, max_tasks_in_memory, *args, **kwargs):
         self._app = app
         self._namespace = namespace
         self.log = logging.getLogger('task-thread')
-        self._state = self._app.events.State(max_tasks_in_memory=max_tasks_in_memory)
+        self._state = CeleryState(max_tasks_in_memory=max_tasks_in_memory)
         self._known_states = set()
         self._known_states_names = set()
         self._tasks_started = dict()
@@ -32,13 +33,12 @@ def run(self):  # pragma: no cover
     def _process_event(self, evt):
         # Events might come in in parallel. Celery already has a lock
         # that deals with this exact situation so we'll use that for now.
-        with self._state._mutex:
-            if celery.events.group_from(evt['type']) == 'task':
-                evt_state = evt['type'][5:]
-                state = celery.events.state.TASK_EVENT_TO_STATE[evt_state]
-                if state == celery.states.STARTED:
-                    self._observe_latency(evt)
-                self._collect_tasks(evt, state)
+        if celery.events.group_from(evt['type']) == 'task':
+            evt_state = evt['type'][5:]
+            state = celery.events.state.TASK_EVENT_TO_STATE[evt_state]
+            if state == celery.states.STARTED:
+                self._observe_latency(evt)
+            self._collect_tasks(evt, state)
 
     def _observe_latency(self, evt):
         try:
@@ -54,31 +54,12 @@ def _observe_latency(self, evt):
                     evt['local_received'] - prev_evt.local_received)
 
     def _collect_tasks(self, evt, state):
-        if state in celery.states.READY_STATES:
-            self._incr_ready_task(evt, state)
-        else:
-            self._incr_unready_task(evt, state)
+        (name, state, runtime) = self._state.collect(evt, state)
+        if runtime is not None:
+            TASKS_RUNTIME.labels(namespace=self._namespace, name=name).observe(runtime)
+
+        TASKS.labels(namespace=self._namespace, name=name, state=state).inc()
 
-    def _incr_ready_task(self, evt, state):
-        try:
-            # remove event from list of in-progress tasks
-            name = self._state.tasks.pop(evt['uuid']).name or ''
-        except (KeyError, AttributeError):  # pragma: no cover
-            name = ''
-        finally:
-            TASKS.labels(namespace=self._namespace, name=name, state=state).inc()
-            if 'runtime' in evt:
-                TASKS_RUNTIME.labels(namespace=self._namespace, name=name) \
-                             .observe(evt['runtime'])
-
-    def _incr_unready_task(self, evt, state):
-        self._state._event(evt)
-        try:
-            name = self._state.tasks[evt['uuid']].name or ''
-        except (KeyError, AttributeError):  # pragma: no cover
-            name = ''
-        finally:
-            TASKS.labels(namespace=self._namespace, name=name, state=state).inc()
 
     def _monitor(self):  # pragma: no cover
         while True:

celery_exporter/state.py

@@ -0,0 +1,63 @@
+import threading
+from operator import itemgetter
+
+from celery.states import READY_STATES
+from celery.utils.functional import LRUCache
+from celery.events.state import Task
+
+
+class CeleryState:
+    event_count = 0
+    task_count = 0
+
+    def __init__(self, max_tasks_in_memory=10000):
+        self.tasks = LRUCache(max_tasks_in_memory)
+        self._mutex = threading.Lock()
+
+    def _event(self, evt):
+        tfields = itemgetter("uuid", "hostname", "timestamp", "local_received", "clock")
+        get_task = self.tasks.__getitem__
+        self.event_count += 1
+        group, _, subject = evt["type"].partition("-")
+        if group == "task":
+            (uuid, hostname, timestamp, local_received, clock) = tfields(evt)
+            is_client_event = subject == "sent"
+            try:
+                task, task_created = get_task(uuid), False
+            except KeyError:
+                task = self.tasks[uuid] = Task(uuid, cluster_state=None)
+                task_created = True
+            if is_client_event:
+                task.client = hostname
+
+            if subject == "received":
+                self.task_count += 1
+
+            task.event(subject, timestamp, local_received, evt)
+
+            return (task, task_created), subject
+
+    def event(self, evt):
+        with self._mutex:
+            return self._event(evt)
+
+    def collect(self, evt, state):
+        runtime = None
+        if state in READY_STATES:
+            try:
+                with self._mutex:
+                    name = self.tasks.pop(evt["uuid"]).name or ""
+            except (KeyError, AttributeError): # pragma: no cover
+                name = ""
+            finally:
+                if "runtime" in evt:
+                    runtime = evt["runtime"]
+                return (name, state, runtime)
+        else:
+            self.event(evt)
+            try:
+                name = self.tasks[evt["uuid"]].name or ""
+            except (KeyError, AttributeError): # pragma: no cover
+                name = ""
+            finally:
+                return (name, state, runtime)

test/test_unit.py

@@ -92,6 +92,13 @@ def test_tasks_events(self):
             runtime=runtime, hostname=hostname, clock=2,
             local_received=local_received + latency_before_started + runtime))
         self._assert_all_states({celery.states.RECEIVED, celery.states.STARTED, celery.states.SUCCESS})
+        
+        m._process_event(Event(
+            'task-sent', uuid=task_uuid,  name=self.task,
+            args='()', kwargs='{}', retries=0, eta=None, hostname=hostname,
+            clock=0,
+            local_received=local_received))
+        self._assert_task_states({celery.states.PENDING}, 1)
 
         assert REGISTRY.get_sample_value('celery_tasks_latency_seconds_count',
             labels=dict(namespace=self.namespace)) == 1