Merge pull request #3 from huntflow/time-metrics-collection

Added time metrics collection

Author: WeirdCarrotMonster

tests/conftest.py

@@ -15,7 +15,6 @@ async def create_redis_connection(redis_uri: str):
     async def zadd_single(client: aioredis.Redis, set_name: str, key: str, value: Any):
         await client.zadd(set_name, {key: value})
 
-
 else:
 
     async def create_redis_connection(redis_uri: str):
@@ -118,6 +117,13 @@ async def assert_metric_broken(self, value: int):
 
         assert stored_value == value
 
+    async def assert_metric_time_wait(self, value: int):
+        stored_value = int(
+            await self.queue.client.get(self.queue.metrics_time_wait) or 0
+        )
+
+        assert stored_value == value
+
 
 @pytest.fixture
 def queue_checker(task_queue) -> QueueChecker:

tests/test_task_queue.py

@@ -516,3 +516,14 @@ async def test_task_retry_delay(task_queue, queue_checker, freezer):
     await queue_checker.assert_metric_added(1)
     await queue_checker.assert_metric_taken(2)
     await queue_checker.assert_metric_requeued(1)
+
+
+@pytest.mark.asyncio
+async def test_task_wait_time_metric(task_queue, queue_checker, freezer):
+    await task_queue.add_task({"key": "value"})
+
+    await queue_checker.assert_metric_time_wait(0)
+    freezer.tick(3.0)
+
+    await task_queue.get_task()
+    await queue_checker.assert_metric_time_wait(3000)

yatq/lua/get_template_lua.py

@@ -23,6 +23,14 @@
 	end
 end
 
+local function incr_metric_key_by (key, amount)
+	local incr_success, value = redis.pcall("INCRBY", key, amount)
+
+	if not incr_success then
+		redis.pcall("SET", key, amount)
+	end
+end
+
 
 --[[ Drops invalid task keys from queue
 
@@ -46,7 +54,7 @@
 end
 
 
-local available_tasks = redis.call("ZRANGEBYSCORE", pending_key, 0, time, "LIMIT", 0, 1)
+local available_tasks = redis.call("ZRANGEBYSCORE", pending_key, 0, time, "WITHSCORES", "LIMIT", 0, 1)
 local task_key = available_tasks[1]
 
 if task_key == nil then
@@ -55,6 +63,8 @@
 	})
 end
 
+local task_score = tonumber(available_tasks[2])
+
 redis.call("ZREM", pending_key, task_key)
 local task_id = redis.call("HGET", task_mapping_key, task_key)
 if not task_id then
@@ -106,6 +116,7 @@
 redis.call("PUBLISH", channel, message)
 
 incr_metric_key("$metrics_taken_key")
+incr_metric_key_by("$metrics_time_wait", math.floor((time - task_score) * 1000))
 
 return cjson.encode({
 	success = true,

yatq/queue.py

@@ -82,6 +82,7 @@ def __init__(
         self.metrics_resurrected_key = f"{self._key_prefix}:metrics:resurrected"
         self.metrics_buried_key = f"{self._key_prefix}:metrics:buried"
         self.metrics_broken_key = f"{self._key_prefix}:metrics:broken"
+        self.metrics_time_wait = f"{self._key_prefix}:metrics:time_wait"
         self.environment = {
             "processing_key": self.processing_set_name,
             "pending_key": self.pending_set_name,
@@ -95,6 +96,7 @@ def __init__(
             "metrics_resurrected_key": self.metrics_resurrected_key,
             "metrics_buried_key": self.metrics_buried_key,
             "metrics_broken_key": self.metrics_broken_key,
+            "metrics_time_wait": self.metrics_time_wait,
             "default_timeout": DEFAULT_TIMEOUT,
             "default_task_expiration": DEFAULT_TASK_EXPIRATION,
         }
@@ -227,7 +229,7 @@ async def auto_reschedule_task(
             if task.policy == RetryPolicy.LINEAR:
                 delay = task.delay * task.retry_counter
             else:
-                delay = task.delay ** task.retry_counter
+                delay = task.delay**task.retry_counter
 
         after_time = int(time.time()) + delay
         task.state = TaskState.REQUEUED

yatq/redis_compat.py

@@ -11,7 +11,6 @@ async def eval_sha(
     ):  # pragma: no cover
         return await client.evalsha(digest, 0, *args)
 
-
 else:  # pragma: no cover
     from aioredis.errors import ReplyError