test: recover test_queue

coco/core.py

@@ -94,6 +94,7 @@ def __init__(
 
         self.state = None
         self.redis_sync = None
+        self.queue_sha = None
 
         # Set to True when a coco_shutdown is happening
         self.coco_shutdown = False
@@ -180,18 +181,6 @@ def __init__(
             "queue_wait_time",
         )
 
-        # Load queue update script into redis cache
-        self.queue_sha = self.redis_sync.script_load(
-            """ if redis.call('llen', KEYS[1]) >= tonumber(ARGV[1]) then
-                        return true
-                    else
-                        redis.call('hset', KEYS[2], ARGV[2], ARGV[3], ARGV[4], ARGV[5], ARGV[6], ARGV[7], ARGV[8], ARGV[9], ARGV[10], ARGV[11])
-                        redis.call('rpush', KEYS[1], KEYS[2])
-                        return false
-                    end
-            """  # noqa: E501
-        )
-
         if testing:
             # Create TCP/IP socket
             sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -612,8 +601,26 @@ async def external_endpoint(self, request, endpoint):
         name = f"{os.getpid()}-{now}"
 
         async with self.redis_async.client() as ra_cli:
-            # Check if queue is full. If not, add this task.
             if self.config["queue_length"] > 0:
+                # If we're using a limited queue, check the queue length before
+                # adding the task.  This needs to be done atomically, so we
+                # use a redis script in this case.
+
+                # Load the script into the redis cache if not already done
+                if not self.queue_sha:
+                    self.queue_sha = self.redis_sync.script_load(
+                        """ if redis.call('llen', KEYS[1]) >= tonumber(ARGV[1]) then
+                                    return true
+                                else
+                                    redis.call('hset', KEYS[2], ARGV[2], ARGV[3],
+                                        ARGV[4], ARGV[5], ARGV[6], ARGV[7], ARGV[8],
+                                        ARGV[9], ARGV[10], ARGV[11])
+                                    redis.call('rpush', KEYS[1], KEYS[2])
+                                    return false
+                                end
+                        """
+                    )
+
                 full = await ra_cli.evalsha(
                     self.queue_sha,
                     2,

tests/coco_runner.py

@@ -30,7 +30,7 @@ def testN(coco_runner):
 import multiprocessing
 import os
 import socket
-from time import sleep
+import time
 
 import aiohttp
 import pytest
@@ -362,7 +362,7 @@ def start_daemon(self, *args):
             finally:
                 # This is in a finally block because we want to wait a bit
                 # even after it successfully fetches the port.
-                sleep(0.2)
+                time.sleep(0.2)
 
     def redis_conn(self):
         """Return a connection to the redis server.
@@ -380,6 +380,35 @@ def redis_conn(self):
 
         return redis.Redis(port=redis_port)
 
+    def enqueue(self, endpoint, method="GET", request="", params=""):
+        """Add an endpoint to the redis queue directly.
+
+        Parameters
+        ----------
+        endpoint:
+            The endpoint to enqueue
+        method:
+            The endpoint request method
+        request:
+            The endpoint request body
+        params:
+            The endpoint request query
+        """
+        redis_conn = self.redis_conn()
+        now = time.perf_counter()
+        tag = f"0-{now}"
+        redis_conn.hset(
+            tag,
+            mapping={
+                "method": method,
+                "endpoint": endpoint,
+                "request": request,
+                "params": params,
+                "received": now,
+            },
+        )
+        redis_conn.rpush("queue", tag)
+
     def call_endpoint(
         self,
         endpoint: str,

tests/test_queue.py

@@ -0,0 +1,50 @@
+"""Redis queue tests."""
+
+from prometheus_client.parser import text_string_to_metric_families
+
+
+def test_limited_queue(coco_runner):
+    """Test limited queue length."""
+    QUEUE_LEN = 3
+
+    coco_runner.add_config(queue_length=QUEUE_LEN)
+    coco_runner.add_targets("test", 1)
+    coco_runner.add_endpoint(
+        "do_wait",
+        {
+            "group": "test",
+            "call": {"coco": {"name": "wait", "request": {"duration": "3s"}}},
+        },
+    )
+    coco_runner.add_endpoint("test", {"group": "test"})
+
+    # Block the qworker
+    coco_runner.enqueue("do_wait")
+
+    # Fill up the queue (the qworker has already popped the do_wait)
+    for _ in range(QUEUE_LEN):
+        coco_runner.enqueue("test")
+
+    # Queue now full
+    result, _ = coco_runner.call_endpoint("test")
+    assert result.status == 503
+
+    # Check metrics record dropped requests
+    result, metrics = coco_runner.call_endpoint("metrics")
+    assert result.status == 200
+
+    metrics = text_string_to_metric_families(metrics)
+    # parse metrics
+    count_coco = []
+    for metric in metrics:
+        for sample in metric.samples:
+            if sample.name == "coco_dropped_request_total":
+                count_coco.append(sample)
+
+    # Find test endpoint metric
+    missing = True
+    for sample in count_coco:
+        if sample.labels["endpoint"] == "test":
+            assert sample.value == 1.0
+            missing = False
+    assert not missing