wehub-resource-sync
663 lines
22 KiB
Python
663 lines
22 KiB
Python
import asyncio
|
|
import sys
|
|
from typing import List
|
|
|
|
import pytest
|
|
|
|
from ray.serve._private.common import DeploymentID
|
|
from ray.serve.experimental.capacity_queue import (
|
|
CapacityQueue,
|
|
CapacityQueueStats,
|
|
ReplicaCapacityInfo,
|
|
)
|
|
|
|
|
|
def _get_raw_class():
|
|
"""Get the underlying class without the @ray.remote decorator."""
|
|
return CapacityQueue.__ray_actor_class__
|
|
|
|
|
|
def _create_queue(acquire_timeout_s=30.0, token_ttl_s=None):
|
|
"""Create a CapacityQueue instance for local testing (no Ray)."""
|
|
raw_class = _get_raw_class()
|
|
return raw_class(
|
|
acquire_timeout_s=acquire_timeout_s,
|
|
token_ttl_s=token_ttl_s,
|
|
_enable_long_poll=False,
|
|
)
|
|
|
|
|
|
class TestCapacityQueueStats:
|
|
"""Tests for CapacityQueueStats dataclass."""
|
|
|
|
def test_stats_fields(self):
|
|
stats = CapacityQueueStats(
|
|
queue_size=10,
|
|
num_waiters=2,
|
|
num_replicas=3,
|
|
total_capacity=15,
|
|
total_in_flight=5,
|
|
total_acquires=100,
|
|
total_releases=95,
|
|
total_timeouts=1,
|
|
max_waiters_seen=5,
|
|
)
|
|
|
|
assert stats.queue_size == 10
|
|
assert stats.num_waiters == 2
|
|
assert stats.num_replicas == 3
|
|
assert stats.total_capacity == 15
|
|
assert stats.total_in_flight == 5
|
|
assert stats.total_acquires == 100
|
|
assert stats.total_releases == 95
|
|
assert stats.total_timeouts == 1
|
|
assert stats.max_waiters_seen == 5
|
|
|
|
|
|
class TestReplicaCapacityInfo:
|
|
"""Tests for ReplicaCapacityInfo dataclass."""
|
|
|
|
def test_default_values(self):
|
|
info = ReplicaCapacityInfo(replica_id="replica1", max_capacity=5)
|
|
assert info.replica_id == "replica1"
|
|
assert info.max_capacity == 5
|
|
assert info.in_flight == 0
|
|
|
|
|
|
class TestCapacityQueueLocal:
|
|
"""Tests for CapacityQueue without Ray (local logic only)."""
|
|
|
|
def test_initialization(self):
|
|
queue = _create_queue(acquire_timeout_s=10.0)
|
|
assert queue._acquire_timeout_s == 10.0
|
|
assert queue.get_queue_length() == 0
|
|
assert len(queue._waiters) == 0
|
|
assert len(queue._replicas) == 0
|
|
|
|
def test_register_replica(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
assert queue.get_queue_length() == 5
|
|
assert "replica1" in queue._replicas
|
|
assert queue._replicas["replica1"].max_capacity == 5
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
|
|
def test_register_replica_duplicate(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
queue.register_replica("replica1", 10) # Should be ignored
|
|
assert queue.get_queue_length() == 5
|
|
assert queue._replicas["replica1"].max_capacity == 5
|
|
|
|
def test_register_multiple_replicas(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 3)
|
|
queue.register_replica("replica2", 5)
|
|
queue.register_replica("replica3", 2)
|
|
assert queue.get_queue_length() == 10
|
|
assert len(queue._replicas) == 3
|
|
|
|
def test_unregister_replica(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
queue.register_replica("replica2", 3)
|
|
queue.unregister_replica("replica1")
|
|
assert queue.get_queue_length() == 3
|
|
assert "replica1" not in queue._replicas
|
|
assert "replica2" in queue._replicas
|
|
|
|
def test_unregister_nonexistent_replica(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
queue.unregister_replica("nonexistent") # Should not raise
|
|
assert queue.get_queue_length() == 5
|
|
|
|
def test_release_to_unregistered_replica(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
initial_tokens = queue.get_queue_length()
|
|
queue.release("nonexistent")
|
|
assert queue.get_queue_length() == initial_tokens
|
|
|
|
def test_get_stats(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
queue.register_replica("replica2", 3)
|
|
stats = queue.get_stats()
|
|
assert stats.queue_size == 8
|
|
assert stats.num_replicas == 2
|
|
assert stats.total_capacity == 8
|
|
assert stats.total_in_flight == 0
|
|
|
|
def test_get_registered_replicas(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
queue.register_replica("replica2", 3)
|
|
replicas = queue.get_registered_replicas()
|
|
assert set(replicas) == {"replica1", "replica2"}
|
|
|
|
|
|
class TestCapacityQueueAsync:
|
|
"""Async tests for CapacityQueue."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_immediate(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 5)
|
|
replica_id = await queue.acquire()
|
|
assert replica_id == "replica1"
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
assert queue.get_queue_length() == 4
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_and_release(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 2)
|
|
|
|
r1 = await queue.acquire()
|
|
r2 = await queue.acquire()
|
|
assert queue.get_queue_length() == 0
|
|
assert queue._replicas["replica1"].in_flight == 2
|
|
|
|
queue.release(r1)
|
|
assert queue.get_queue_length() == 1
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
|
|
queue.release(r2)
|
|
assert queue.get_queue_length() == 2
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_timeout(self):
|
|
queue = _create_queue(acquire_timeout_s=0.1)
|
|
# No replicas registered
|
|
replica_id = await queue.acquire()
|
|
assert replica_id is None
|
|
assert queue._total_timeouts == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_timeout_custom(self):
|
|
queue = _create_queue(acquire_timeout_s=10.0)
|
|
replica_id = await queue.acquire(timeout_s=0.1)
|
|
assert replica_id is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_waits_for_release(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 1)
|
|
|
|
await queue.acquire()
|
|
assert queue.get_queue_length() == 0
|
|
|
|
acquire_task = asyncio.create_task(queue.acquire(timeout_s=5.0))
|
|
await asyncio.sleep(0.01)
|
|
assert queue.get_num_waiters() == 1
|
|
|
|
queue.release("replica1")
|
|
replica_id = await acquire_task
|
|
assert replica_id == "replica1"
|
|
assert queue.get_num_waiters() == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_waits_for_registration(self):
|
|
queue = _create_queue()
|
|
|
|
acquire_task = asyncio.create_task(queue.acquire(timeout_s=5.0))
|
|
await asyncio.sleep(0.01)
|
|
assert queue.get_num_waiters() == 1
|
|
|
|
queue.register_replica("replica1", 3)
|
|
replica_id = await acquire_task
|
|
assert replica_id == "replica1"
|
|
assert queue.get_num_waiters() == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_cancellation(self):
|
|
queue = _create_queue()
|
|
|
|
acquire_task = asyncio.create_task(queue.acquire(timeout_s=5.0))
|
|
await asyncio.sleep(0.01)
|
|
assert queue.get_num_waiters() == 1
|
|
|
|
acquire_task.cancel()
|
|
try:
|
|
await acquire_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
await asyncio.sleep(0.01)
|
|
assert queue.get_num_waiters() == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_least_loaded_distribution(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 3)
|
|
queue.register_replica("replica2", 3)
|
|
|
|
# Acquire 2 tokens - should distribute evenly
|
|
await queue.acquire()
|
|
await queue.acquire()
|
|
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
assert queue._replicas["replica2"].in_flight == 1
|
|
|
|
# Acquire 2 more - again evenly
|
|
await queue.acquire()
|
|
await queue.acquire()
|
|
|
|
assert queue._replicas["replica1"].in_flight == 2
|
|
assert queue._replicas["replica2"].in_flight == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_max_waiters_tracking(self):
|
|
queue = _create_queue()
|
|
|
|
tasks = [asyncio.create_task(queue.acquire(timeout_s=0.2)) for _ in range(5)]
|
|
await asyncio.sleep(0.05)
|
|
await asyncio.gather(*tasks)
|
|
|
|
stats = queue.get_stats()
|
|
assert stats.max_waiters_seen >= 5
|
|
|
|
|
|
class TestCapacityQueueStatsTracking:
|
|
"""Tests for statistics tracking during operations."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_total_acquires_increments(self):
|
|
queue = _create_queue(acquire_timeout_s=0.1)
|
|
queue.register_replica("replica1", 2)
|
|
|
|
assert queue._total_acquires == 0
|
|
|
|
await queue.acquire()
|
|
assert queue._total_acquires == 1
|
|
|
|
await queue.acquire()
|
|
assert queue._total_acquires == 2
|
|
|
|
# Timeout also counts
|
|
await queue.acquire()
|
|
assert queue._total_acquires == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_total_releases_increments(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 3)
|
|
|
|
assert queue._total_releases == 0
|
|
|
|
r1 = await queue.acquire()
|
|
r2 = await queue.acquire()
|
|
|
|
queue.release(r1)
|
|
assert queue._total_releases == 1
|
|
|
|
queue.release(r2)
|
|
assert queue._total_releases == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_in_flight_tracking(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 3)
|
|
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
|
|
r1 = await queue.acquire()
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
|
|
r2 = await queue.acquire()
|
|
assert queue._replicas["replica1"].in_flight == 2
|
|
|
|
queue.release(r1)
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
|
|
queue.release(r2)
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stats_reflect_operations(self):
|
|
queue = _create_queue(acquire_timeout_s=0.1)
|
|
queue.register_replica("replica1", 2)
|
|
queue.register_replica("replica2", 1)
|
|
|
|
r1 = await queue.acquire()
|
|
await queue.acquire()
|
|
|
|
stats = queue.get_stats()
|
|
assert stats.total_capacity == 3
|
|
assert stats.total_in_flight == 2
|
|
assert stats.queue_size == 1
|
|
assert stats.total_acquires == 2
|
|
|
|
queue.release(r1)
|
|
|
|
stats = queue.get_stats()
|
|
assert stats.total_in_flight == 1
|
|
assert stats.queue_size == 2
|
|
assert stats.total_releases == 1
|
|
|
|
|
|
class TestCapacityQueueEdgeCases:
|
|
"""Tests for edge cases and error handling."""
|
|
|
|
def test_register_zero_capacity(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 0)
|
|
assert "replica1" in queue._replicas
|
|
assert queue._replicas["replica1"].max_capacity == 0
|
|
assert queue.get_queue_length() == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_while_in_flight(self):
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 3)
|
|
|
|
r1 = await queue.acquire()
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
|
|
queue.unregister_replica("replica1")
|
|
assert "replica1" not in queue._replicas
|
|
assert queue.get_queue_length() == 0
|
|
|
|
# Release should be handled gracefully (discarded)
|
|
queue.release(r1)
|
|
assert queue.get_queue_length() == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_same_token_twice(self):
|
|
"""Releasing the same token twice doesn't inflate capacity."""
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 1)
|
|
|
|
r1 = await queue.acquire()
|
|
assert queue.get_queue_length() == 0
|
|
|
|
queue.release(r1)
|
|
assert queue.get_queue_length() == 1
|
|
|
|
# Release again - in_flight clamped to 0
|
|
queue.release(r1)
|
|
assert queue.get_queue_length() == 1 # Correctly doesn't inflate
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fifo_waiter_ordering(self):
|
|
"""Waiters are fulfilled in FIFO order."""
|
|
queue = _create_queue()
|
|
|
|
completion_order: List[int] = []
|
|
|
|
async def waiter(waiter_id: int):
|
|
await queue.acquire(timeout_s=5.0)
|
|
completion_order.append(waiter_id)
|
|
|
|
tasks = [asyncio.create_task(waiter(i)) for i in range(3)]
|
|
await asyncio.sleep(0.05)
|
|
assert queue.get_num_waiters() == 3
|
|
|
|
queue.register_replica("replica1", 3)
|
|
await asyncio.gather(*tasks)
|
|
|
|
assert completion_order == [0, 1, 2]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_replicas_least_loaded(self):
|
|
"""Least-loaded replica is selected."""
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 2)
|
|
queue.register_replica("replica2", 2)
|
|
|
|
r1 = await queue.acquire()
|
|
r2 = await queue.acquire()
|
|
|
|
# Should be different replicas (load balanced)
|
|
assert r1 != r2
|
|
|
|
# Both should have 1 in-flight
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
assert queue._replicas["replica2"].in_flight == 1
|
|
|
|
# Acquire two more
|
|
await queue.acquire()
|
|
await queue.acquire()
|
|
|
|
assert queue._replicas["replica1"].in_flight == 2
|
|
assert queue._replicas["replica2"].in_flight == 2
|
|
|
|
# Clean up
|
|
queue.release(r1)
|
|
queue.release(r2)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_waiter_timeout_cleanup(self):
|
|
"""Timed-out waiters are properly cleaned up."""
|
|
queue = _create_queue(acquire_timeout_s=0.1)
|
|
|
|
tasks = [asyncio.create_task(queue.acquire()) for _ in range(3)]
|
|
results = await asyncio.gather(*tasks)
|
|
|
|
assert all(r is None for r in results)
|
|
assert queue.get_num_waiters() == 0
|
|
assert queue._total_timeouts == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_before_token_acquired(self):
|
|
"""Cancelling a request while waiting for a token (tokens exhausted)
|
|
should clean up the waiter without affecting capacity."""
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 1)
|
|
|
|
# Exhaust the single token.
|
|
r1 = await queue.acquire()
|
|
assert queue.get_queue_length() == 0
|
|
|
|
# Second acquire blocks — no capacity available.
|
|
waiting_task = asyncio.create_task(queue.acquire(timeout_s=5.0))
|
|
await asyncio.sleep(0.01)
|
|
assert queue.get_num_waiters() == 1
|
|
|
|
# Cancel before a token is granted.
|
|
waiting_task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiting_task
|
|
|
|
await asyncio.sleep(0.01)
|
|
assert queue.get_num_waiters() == 0
|
|
# Capacity unchanged — still 0 available (1 in-flight).
|
|
assert queue.get_queue_length() == 0
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
|
|
# Release the held token — full capacity restored.
|
|
queue.release(r1)
|
|
assert queue.get_queue_length() == 1
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_after_token_acquired(self):
|
|
"""Cancelling a request after a token was acquired should still allow
|
|
the caller to release the token, preserving capacity."""
|
|
queue = _create_queue()
|
|
queue.register_replica("replica1", 2)
|
|
|
|
# Acquire a token normally.
|
|
r1 = await queue.acquire()
|
|
assert r1 == "replica1"
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
assert queue.get_queue_length() == 1
|
|
|
|
# Simulate the caller being cancelled after acquiring — the token is
|
|
# held. The caller (router) is responsible for releasing it.
|
|
# Verify capacity is reduced while the token is held.
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
|
|
# Release the token (as the router's on_request_cancelled would do).
|
|
queue.release(r1)
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
assert queue.get_queue_length() == 2
|
|
|
|
|
|
class TestReplicaLifecycleOnDeploymentTargetUpdate:
|
|
"""Tests that the queue reflects replica changes from the controller.
|
|
|
|
The controller sends deployment target updates (via long poll) when replicas
|
|
are added or removed. The queue must unregister removed replicas and
|
|
register new ones so that tokens are only issued for live replicas.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _make_target_info(replica_specs):
|
|
"""Build a DeploymentTargetInfo from a list of (unique_id, capacity) tuples."""
|
|
from ray.serve._private.common import (
|
|
DeploymentTargetInfo,
|
|
ReplicaID,
|
|
RunningReplicaInfo,
|
|
)
|
|
|
|
dep_id = DeploymentID(name="test_deployment", app_name="test_app")
|
|
replicas = []
|
|
for uid, cap in replica_specs:
|
|
replicas.append(
|
|
RunningReplicaInfo(
|
|
replica_id=ReplicaID(unique_id=uid, deployment_id=dep_id),
|
|
node_id="node1",
|
|
node_ip="127.0.0.1",
|
|
availability_zone=None,
|
|
actor_name=f"actor_{uid}",
|
|
max_ongoing_requests=cap,
|
|
is_cross_language=False,
|
|
)
|
|
)
|
|
return DeploymentTargetInfo(is_available=True, running_replicas=replicas)
|
|
|
|
def test_update_adds_new_replicas(self):
|
|
"""New replicas in the update are registered in the queue."""
|
|
queue = _create_queue()
|
|
target = self._make_target_info([("r1", 5)])
|
|
|
|
queue._update_deployment_targets(target)
|
|
|
|
assert set(queue.get_registered_replicas()) == {"r1"}
|
|
assert queue.get_queue_length() == 5
|
|
|
|
def test_update_removes_old_replicas(self):
|
|
"""Replicas absent from the update are unregistered."""
|
|
queue = _create_queue()
|
|
queue.register_replica("r1", 3)
|
|
queue.register_replica("r2", 2)
|
|
assert len(queue.get_registered_replicas()) == 2
|
|
|
|
# Update with only r1 — r2 should be removed
|
|
target = self._make_target_info([("r1", 3)])
|
|
queue._update_deployment_targets(target)
|
|
|
|
assert set(queue.get_registered_replicas()) == {"r1"}
|
|
assert queue.get_queue_length() == 3
|
|
|
|
def test_update_empty_removes_all(self):
|
|
"""An update with no replicas clears the queue."""
|
|
queue = _create_queue()
|
|
queue.register_replica("r1", 3)
|
|
|
|
target = self._make_target_info([])
|
|
queue._update_deployment_targets(target)
|
|
|
|
assert queue.get_registered_replicas() == []
|
|
assert queue.get_queue_length() == 0
|
|
|
|
def test_update_is_idempotent(self):
|
|
"""Calling update twice with the same replicas is a no-op."""
|
|
queue = _create_queue()
|
|
target = self._make_target_info([("r1", 5)])
|
|
|
|
queue._update_deployment_targets(target)
|
|
queue._update_deployment_targets(target)
|
|
|
|
assert set(queue.get_registered_replicas()) == {"r1"}
|
|
assert queue.get_queue_length() == 5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scale_down_frees_dead_capacity(self):
|
|
"""After scale-down, tokens for removed replicas are no longer issued."""
|
|
queue = _create_queue()
|
|
queue.register_replica("r1", 1)
|
|
queue.register_replica("r2", 1)
|
|
|
|
# Saturate both replicas
|
|
t1 = await queue.acquire()
|
|
t2 = await queue.acquire()
|
|
assert queue.get_queue_length() == 0
|
|
|
|
# Scale down: remove r2
|
|
target = self._make_target_info([("r1", 1)])
|
|
queue._update_deployment_targets(target)
|
|
|
|
# Release both tokens — only r1's should restore capacity
|
|
queue.release(t1)
|
|
queue.release(t2) # r2 is gone, release is discarded
|
|
|
|
assert queue.get_queue_length() == 1
|
|
assert set(queue.get_registered_replicas()) == {"r1"}
|
|
|
|
# Next acquire must go to r1 (no phantom r2)
|
|
token = await queue.acquire(timeout_s=1.0)
|
|
assert token == "r1"
|
|
|
|
|
|
class TestTokenTTL:
|
|
"""Tests for the token TTL auto-reclaim feature."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expired_token_reclaimed(self):
|
|
"""Tokens exceeding the TTL are automatically reclaimed."""
|
|
queue = _create_queue(token_ttl_s=0.2)
|
|
queue.register_replica("replica1", 2)
|
|
|
|
# Acquire a token (never release it).
|
|
r = await queue.acquire()
|
|
assert r == "replica1"
|
|
assert queue._replicas["replica1"].in_flight == 1
|
|
assert queue.get_queue_length() == 1
|
|
|
|
# Wait for the TTL reaper to reclaim it.
|
|
await asyncio.sleep(0.5)
|
|
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
assert queue.get_queue_length() == 2
|
|
assert queue._total_ttl_reclaims == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ttl_does_not_reclaim_released_tokens(self):
|
|
"""Properly released tokens are not double-reclaimed by the reaper."""
|
|
queue = _create_queue(token_ttl_s=0.2)
|
|
queue.register_replica("replica1", 2)
|
|
|
|
r = await queue.acquire()
|
|
queue.release(r)
|
|
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
|
|
# Wait past TTL — reaper should find nothing to reclaim.
|
|
await asyncio.sleep(0.5)
|
|
|
|
assert queue._replicas["replica1"].in_flight == 0
|
|
assert queue.get_queue_length() == 2
|
|
assert queue._total_ttl_reclaims == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ttl_reclaim_wakes_waiters(self):
|
|
"""Reclaimed capacity from TTL should wake blocked waiters."""
|
|
queue = _create_queue(token_ttl_s=0.2)
|
|
queue.register_replica("replica1", 1)
|
|
|
|
# Exhaust capacity.
|
|
await queue.acquire()
|
|
assert queue.get_queue_length() == 0
|
|
|
|
# This waiter should be woken once the TTL reaper reclaims.
|
|
result = await asyncio.wait_for(queue.acquire(), timeout=2.0)
|
|
assert result == "replica1"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__, "-v"]))
|