diff --git a/coco/core.py b/coco/core.py index 43b8347a..c1deda5f 100644 --- a/coco/core.py +++ b/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, diff --git a/old_tests/test_queue.py b/old_tests/test_queue.py deleted file mode 100644 index 5da05ec2..00000000 --- a/old_tests/test_queue.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Test the limited length queue.""" - -import asyncio -import time - -import pytest -import requests -from aiohttp import request -from coco.test import coco_runner, endpoint_farm -from prometheus_client.parser import text_string_to_metric_families - -PORT = 12055 -METRIC_PORT = 12056 -T_WAIT = 2 -QUEUE_LEN = 3 -CONFIG = { - "log_level": "DEBUG", - "queue_length": QUEUE_LEN, - "port": PORT, - "metrics_port": METRIC_PORT, -} -ENDPOINTS = { - "do_wait": { - "group": "test", - "call": {"coco": {"name": "wait", "request": {"duration": T_WAIT}}}, - }, - "test": {"group": "test"}, -} - - -def callback(data): - """Reply with the incoming json request.""" - return data - - -N_HOSTS = 2 -CALLBACKS = dict.fromkeys(ENDPOINTS, callback) - - -@pytest.fixture -def farm(): - """Create an endpoint test farm.""" - return endpoint_farm.Farm(N_HOSTS, CALLBACKS) - - -@pytest.fixture -def runner(farm): - """Create a coco runner.""" - CONFIG["groups"] = {"test": farm.hosts} - with coco_runner.Runner(CONFIG, ENDPOINTS, reset_on_shutdown=False) as runner: - yield runner - - -@pytest.fixture -def loop(): - """Use a separate event loop.""" - orig_loop = asyncio.get_event_loop() - try: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - yield loop - loop.close() - finally: - asyncio.set_event_loop(orig_loop) - - -async def _client(config, endpoint, sleep=None): - if sleep: - await asyncio.sleep(sleep) - async with request( - "get", f"http://localhost:{PORT}/{endpoint}", json={"coco_report_type": "FULL"} - ) as r: - return await r.json() - - -def test_queue(farm, runner, loop): - """Test queue limit.""" - # Wait for coco to start up - time.sleep(3) - - # Set up client tasks - wait = _client(runner.configfile.name, "do_wait") - clients = [] - for i in range(QUEUE_LEN + 1): - clients.append(_client(runner.configfile.name, "test", sleep=0.1)) - - # Send requests - replies = loop.run_until_complete(asyncio.gather(wait, *clients)) - - # Check responses - failed = 0 - for r in replies[1:]: - if "status" in r: - assert r["status"] == 503 - failed += 1 - else: - for h in farm.hosts: - assert h in r["test"] - assert r["test"][h]["status"] == 200 - # Not certain they came in order, but only one should have been dropped - assert failed == 1 - - # Check metrics record dropped requests - metrics = requests.get(f"http://localhost:{METRIC_PORT}/metrics") - assert metrics.status_code == 200 - metrics = text_string_to_metric_families(metrics.text) - - # 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 - - runner.stop_coco() diff --git a/tests/coco_runner.py b/tests/coco_runner.py index 09ddc847..eac2ef9c 100644 --- a/tests/coco_runner.py +++ b/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, diff --git a/tests/test_queue.py b/tests/test_queue.py new file mode 100644 index 00000000..83d7cf6a --- /dev/null +++ b/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