strict order sharded job queue
from asyncio import get_event_loop, sleep
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from functools import cache
from importlib import import_module
from signal import SIGTERM
from typing import (Any, AsyncGenerator, Generic, NamedTuple, Optional, Self,
                    TypeVar, get_type_hints)

from sharded_queue.drivers import JsonTupleSerializer
from sharded_queue.protocols import Lock, Serializer, Storage
from sharded_queue.settings import WorkerSettings

T = TypeVar('T')


class Route(NamedTuple):
    thread: int = 0
    priority: int = 0


class Handler(Generic[T]):
    priorities: Optional[list[int]] = None

    @classmethod
    async def create(cls) -> Self:
        return cls()

    @classmethod
    def request_cls(cls) -> type[T]:
        request_cls = list(get_type_hints(cls.handle).values())[0]
        if isinstance(request_cls, TypeVar):
            raise NotImplementedError(cls)
        return request_cls

    @classmethod
    async def route(cls, *requests: T) -> list[Route]:
        return [Route() for _ in requests]

    async def start(self) -> None:
        pass

    async def handle(self, *requests: T) -> None:
        raise NotImplementedError(self.__class__)

    async def stop(self) -> None:
        pass


class Tube(NamedTuple):
    handler: type[Handler]
    route: Route

    @classmethod
    @cache
    def parse_pipe(cls, tube: str) -> Self:
        [*module, name, thread, priority] = tube.split('/')
        handler = getattr(import_module("_".join(module)), name)
        return cls(handler, Route(int(thread), int(priority)))

    @property
    def pipe(self) -> str:
        return '/'.join([
            self.handler.__module__,
            self.handler.__name__,
            str(self.route.thread),
            str(self.route.priority),
        ])

    @asynccontextmanager
    async def context(self) -> AsyncGenerator:
        instance = await self.handler.create()
        await instance.start()
        try:
            yield instance
        finally:
            await instance.stop()


@dataclass
class Queue(Generic[T]):
    def __init__(
        self,
        storage: Storage,
        serializer: Optional[Serializer] = None,
    ):
        self.storage = storage
        self.serializer = serializer or JsonTupleSerializer()

    async def register(
        self, handler: type[Handler], *requests: T,
        defer: Optional[float | int | timedelta] = None,
        if_not_exists: bool = False,
        recurrent: Optional[float | int | timedelta] = None,
    ) -> None:
        routes = await handler.route(*requests)
        pipe_messages: list[tuple[str, Any]] = [
            (Tube(handler, routes[n]).pipe, requests[n])
            for n in range(len(routes))
        ]

        if recurrent:
            if_not_exists = True
            pipe_messages = RecurrentHandler.transform(
                pipe_messages, recurrent, self.serializer
            )

        if defer:
            if_not_exists = True
            pipe_messages = DeferredHandler.transform(
                pipe_messages, defer, self.serializer
            )

        for pipe in set([pipe for (pipe, _) in pipe_messages]):
            msgs = [
                msg for msg in
                [
                    self.serializer.serialize(request)
                    for (pipe_candidate, request) in pipe_messages
                    if pipe_candidate == pipe
                ]
                if not if_not_exists
                or not await self.storage.contains(pipe, msg)
            ]
            if len(msgs):
                await self.storage.append(pipe, *msgs)


@dataclass
class Worker:
    lock: Lock
    queue: Queue
    pipe: Optional[str] = None
    settings: WorkerSettings = field(default_factory=WorkerSettings)

    async def acquire_tube(
        self, handler: Optional[type[Handler]] = None
    ) -> Optional[Tube]:
        all_pipes = False
        while get_event_loop().is_running():
            for pipe in await self.queue.storage.pipes():
                if not await self.queue.storage.length(pipe):
                    continue
                tube: Tube = Tube.parse_pipe(pipe)
                if handler and tube.handler is not handler:
                    continue
                if tube.handler.priorities:
                    if tube.route.priority != tube.handler.priorities[0]:
                        if not all_pipes:
                            continue
                        tube = Tube(
                            handler=tube.handler,
                            route=Route(
                                thread=tube.route.thread,
                                priority=tube.handler.priorities[0]
                            )
                        )
                if not await self.lock.acquire(tube.pipe):
                    continue
                return tube

            if all_pipes:
                await sleep(self.settings.acquire_delay)
            else:
                all_pipes = True

        return None

    def page_size(self, limit: Optional[int] = None) -> int:
        if limit is None:
            return self.settings.batch_size

        return min(limit, self.settings.batch_size)

    async def loop(
        self,
        limit: Optional[int] = None,
        handler: Optional[type[Handler]] = None,
    ) -> None:
        get_event_loop().add_signal_handler(SIGTERM, self.housekeep)
        processed = 0
        while get_event_loop().is_running() and (
            limit is None or limit > processed
        ):
            tube = await self.acquire_tube(handler)
            if not tube:
                break
            self.pipe = tube.pipe
            processed = processed + await self.process(tube, limit)
            self.pipe = None

        get_event_loop().remove_signal_handler(SIGTERM)

    def housekeep(self) -> None:
        if self.pipe:
            get_event_loop().create_task(self.lock.release(self.pipe))

    async def process(self, tube: Tube, limit: Optional[int] = None) -> int:
        deserialize = self.queue.serializer.deserialize
        storage = self.queue.storage

        cls = tube.handler.request_cls()
        pipes = [tube.pipe]
        empty_counter = 0
        processed_counter = 0

        if tube.handler.priorities:
            pipes = [
                Tube(tube.handler, Route(tube.route.thread, priority)).pipe
                for priority in tube.handler.priorities
            ]

        async with tube.context() as instance:
            if isinstance(instance, DeferredHandler | RecurrentHandler):
                instance.worker = self

            while get_event_loop().is_running() and (
                limit is None or limit > processed_counter
            ):
                if tube.handler is RecurrentHandler:
                    page_size = self.settings.recurrent_tasks_limit
                else:
                    page_size = self.page_size(limit)
                processed = False
                for pipe in pipes:
                    msgs = await storage.range(pipe, page_size)
                    if not len(msgs):
                        continue

                    await instance.handle(*[
                        deserialize(cls, msg) for msg in msgs
                    ])

                    if tube.handler is RecurrentHandler:
                        await self.prolongate_lock(
                            self.settings.recurrent_check_interval
                        )
                        return len(msgs)

                    await storage.pop(pipe, len(msgs))

                    processed = True
                    processed_counter = processed_counter + len(msgs)
                    empty_counter = 0
                    break

                if not processed:
                    empty_counter = empty_counter + 1
                    if empty_counter >= self.settings.empty_limit:
                        break
                    await sleep(self.settings.empty_pause)

                if tube.handler is DeferredHandler:
                    await self.prolongate_lock(
                        self.settings.deferred_retry_delay
                    )
                    return processed_counter
                else:
                    await self.prolongate_lock()

        if self.pipe:
            await self.lock.release(self.pipe)

        return processed_counter

    async def prolongate_lock(self, ttl: Optional[int] = None):
        if not self.pipe:
            raise RuntimeError('No active pipe')
        if ttl is None:
            ttl = self.lock.settings.timeout
        await self.lock.ttl(self.pipe, ttl)


class DeferredRequest(NamedTuple):
    timestamp: float
    pipe: str
    msg: list

    @classmethod
    def calculate_timestamp(cls, delta: float | int | timedelta) -> float:
        now: datetime = datetime.now()
        if isinstance(delta, timedelta):
            now = now + delta

        timestamp: float = now.timestamp()
        if not isinstance(delta, timedelta):
            timestamp = delta = delta

        return timestamp


class DeferredHandler(Handler):
    worker: Worker

    async def handle(self, *requests: DeferredRequest) -> None:
        now: float = datetime.now().timestamp()
        pending: list[DeferredRequest] = [
            request for request in requests
            if request.timestamp > now
        ]

        if len(pending):
            await self.worker.queue.register(DeferredHandler, *pending)

        todo: list[tuple[str, list]] = [
            (request.pipe, request.msg)
            for request in requests
            if request.timestamp <= now
        ]

        for pipe in set([pipe for (pipe, _) in todo]):
            ready = [
                msg for msg in
                [
                    self.worker.queue.serializer.serialize(request)
                    for (candidate_pipe, request) in todo
                    if candidate_pipe == pipe
                ]
                if not await self.worker.queue.storage.contains(pipe, msg)
            ]

            if len(ready):
                await self.worker.queue.storage.append(pipe, *ready)

        if len(pending) and not len(todo):
            await sleep(self.worker.settings.deferred_retry_delay)

    @classmethod
    def transform(
        cls,
        pipe_messages: list[tuple[str, T]],
        defer: float | int | timedelta,
        serializer: Serializer,
    ) -> list[tuple[str, DeferredRequest]]:
        timestamp = DeferredRequest.calculate_timestamp(defer)
        return [
            (
                Tube(DeferredHandler, Route()).pipe,
                DeferredRequest(timestamp, pipe, values),
            )
            for (pipe, values)
            in [
                (pipe, serializer.get_values(request))
                for (pipe, request) in pipe_messages
            ]
        ]


class RecurrentRequest(NamedTuple):
    interval: float
    pipe: str
    msg: list

    @classmethod
    def get_interval(cls, interval: int | float | timedelta) -> float:
        if isinstance(interval, timedelta):
            return float(interval.seconds)

        return float(interval)


class RecurrentHandler(Handler):
    worker: Worker

    async def handle(self, *requests: RecurrentRequest) -> None:
        deferred_pipe: str = Tube(DeferredHandler, Route()).pipe
        deferred_requests: list[tuple[str, str]] = [
            (request.pipe, request.msg)
            for request in [
                self.worker.queue.serializer.deserialize(DeferredRequest, msg)
                for msg in await self.worker.queue.storage.range(
                    deferred_pipe, self.worker.settings.recurrent_tasks_limit
                )
            ]
        ]

        todo: list[DeferredRequest] = [
            DeferredRequest(
                DeferredRequest.calculate_timestamp(request.interval),
                request.pipe,
                request.msg,
            )
            for request in requests
            if (request.pipe, request.msg) not in deferred_requests
        ]

        if len(todo):
            await self.worker.queue.register(
                DeferredHandler, *todo, if_not_exists=True
            )

    @classmethod
    def transform(
        cls,
        pipe_messages: list[tuple[str, T]],
        recurrent: float | int | timedelta,
        serializer: Serializer,
    ) -> list[tuple[str, RecurrentRequest]]:
        interval: float = RecurrentRequest.get_interval(recurrent)
        return [
            (
                Tube(RecurrentHandler, Route()).pipe,
                RecurrentRequest(interval, pipe, values)
            )
            for (pipe, values)
            in [
                (pipe, serializer.get_values(request))
                for (pipe, request) in pipe_messages
            ]
        ]

sharded_queue/__init__.py
basis-company/sharded-queue.py
from json import dumps, loads
from typing import Any, List, Optional, Sequence

from redis.asyncio import Redis

from sharded_queue.protocols import Lock, Serializer, Storage
from sharded_queue.settings import LockSettings, StorageSettings


class JsonTupleSerializer(Serializer):
    def get_values(self, request) -> list[Any]:
        if isinstance(request, Sequence):
            return [k for k in request]
        return list(request.__dict__.values())

    def serialize(self, request: Any) -> str:
        return dumps(self.get_values(request))

    def deserialize(self, cls: type[Any], source: str) -> Any:
        values = loads(source)
        if hasattr(cls, 'model_fields'):
            return cls(**dict(zip(cls.model_fields, values)))

        return cls(*values)


class RuntimeLock(Lock):
    def __init__(self) -> None:
        self.storage: dict[str, bool] = {}

    async def acquire(self, key: str) -> bool:
        if key in self.storage:
            return False
        self.storage[key] = True
        return True

    async def exists(self, key: str) -> bool:
        return key in self.storage

    async def release(self, key: str) -> None:
        del self.storage[key]

    async def ttl(self, key: str, ttl: int) -> bool:
        if ttl == 0:
            await self.release(key)
            return True
        return await self.exists(key)


class RuntimeStorage(Storage):
    data: dict[str, List[str]]

    def __init__(self) -> None:
        self.data = {}

    async def append(self, tube: str, *msgs: str) -> int:
        if tube not in self.data:
            self.data[tube] = list(msgs)
        else:
            self.data[tube].extend(list(msgs))

        return len(self.data[tube])

    async def contains(self, tube: str, msg: str) -> bool:
        return tube in self.data and msg in self.data[tube]

    async def length(self, tube: str) -> int:
        return len(self.data[tube]) if tube in self.data else 0

    async def pop(self, tube: str, max: int) -> list[str]:
        res = await self.range(tube, max)
        if len(res):
            self.data[tube] = self.data[tube][len(res):]
        return res

    async def pipes(self) -> list[str]:
        return list(self.data.keys())

    async def range(self, tube: str, max: int) -> list[str]:
        return self.data[tube][0:max] if tube in self.data else []


class RedisLock(Lock):
    def __init__(
        self, redis: Redis, settings: Optional[LockSettings] = None
    ) -> None:
        self.redis = redis
        self.settings = settings or LockSettings()

    async def acquire(self, key: str) -> bool:
        return None is not await self.redis.set(
            name=self.settings.prefix + key,
            ex=self.settings.timeout,
            nx=True,
            value=1,
        )

    async def exists(self, key: str) -> bool:
        checker = await self.redis.exists(
            self.settings.prefix + key
        )
        return bool(checker)

    async def release(self, key: str) -> None:
        await self.redis.delete(self.settings.prefix + key)

    async def ttl(self, key: str, ttl: int) -> bool:
        setter = await self.redis.set(
            self.settings.prefix + key,
            value=key,
            ex=ttl,
            xx=True
        )
        return bool(setter)


class RedisStorage(Storage):
    def __init__(
        self, redis: Redis, settings: Optional[StorageSettings] = None
    ) -> None:
        self.redis = redis
        self.settings = settings or StorageSettings()

    async def append(self, tube: str, *msgs: str) -> int:
        return await self.redis.rpush(self.key(tube), *msgs)

    async def contains(self, tube: str, msg: str) -> bool:
        return await self.redis.lpos(self.key(tube), msg) is not None

    def key(self, tube):
        return self.settings.prefix + tube

    async def length(self, tube: str) -> int:
        return await self.redis.llen(self.key(tube))

    async def pipes(self) -> list[str]:
        return [
            key[len(self.settings.prefix):]
            for key in await self.redis.keys(self.key('*'))
        ]

    async def pop(self, tube: str, max: int) -> list[str]:
        return await self.redis.lpop(self.key(tube), max) or []

    async def range(self, tube: str, max: int) -> list[str]:
        return await self.redis.lrange(self.key(tube), 0, max-1) or []

sharded_queue/drivers.py
basis-company/sharded-queue.py
from typing import Any, Protocol

from sharded_queue.settings import LockSettings, StorageSettings


class Lock(Protocol):
    settings: LockSettings = LockSettings()

    async def acquire(self, key: str) -> bool: ...
    async def exists(self, key: str) -> bool: ...
    async def release(self, key: str) -> None: ...
    async def ttl(self, key: str, ttl: int) -> bool: ...


class Serializer(Protocol):
    def get_values(self, request) -> list[Any]: ...
    def serialize(self, request: Any) -> str: ...
    def deserialize(self, cls: type[Any], source: str) -> Any: ...


class Storage(Protocol):
    settings: StorageSettings = StorageSettings()

    async def append(self, tube: str, *msgs: str) -> int: ...
    async def contains(self, tube: str, msg: str) -> bool: ...
    async def length(self, tube: str) -> int: ...
    async def pipes(self) -> list[str]: ...
    async def pop(self, tube: str, max: int) -> list[str]: ...
    async def range(self, tube: str, max: int) -> list[str]: ...

sharded_queue/protocols.py
basis-company/sharded-queue.py
from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class LockSettings(BaseSettings):
    model_config = SettingsConfigDict(env_prefix='queue_lock_')

    prefix: str = Field(
        default="lock_",
        title="Lock key prefix"
    )

    timeout: int = Field(
        default=24*60*60,
        title="Lock key ttl"
    )


class StorageSettings(BaseSettings):
    model_config = SettingsConfigDict(env_prefix='queue_storage_')

    prefix: str = Field(
        default="tube_",
        title="Tube name prefix"
    )


class WorkerSettings(BaseSettings):

    acquire_delay: float = Field(
        default=1,
        title="Acquire delay in seconds on empty queues"
    )

    batch_size: int = Field(
        default=128,
        title='Batch processing size'
    )

    deferred_retry_delay: int = Field(
        default=1,
        title='Defereed tasks retry delay'
    )

    empty_limit: int = Field(
        default=16,
        title="Empty queue attempt limit berfore queue rebind",
    )

    empty_pause: float = Field(
        default=0.1,
        title="Pause in seconds on empty queue",
    )

    model_config = SettingsConfigDict(env_prefix='queue_worker_')

    recurrent_check_interval: int = Field(
        default=30,
        title='Recurrent interval check in seconds'
    )

    recurrent_tasks_limit: int = Field(
        default=1024,
        title='Recurrent tasks limit count'
    )

sharded_queue/settings.py
basis-company/sharded-queue.py
from asyncio import sleep
from datetime import timedelta
from typing import NamedTuple

from pytest import mark

from sharded_queue import DeferredHandler, Handler, Queue, Route, Tube, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage
from sharded_queue.protocols import Storage


class BucketRequest(NamedTuple):
    bucket: int


class DropBucket(Handler):
    async def handle(self, *requests: BucketRequest) -> None:
        pass


@mark.asyncio
async def test_deferred() -> None:
    storage: Storage = RuntimeStorage()
    queue: Queue = Queue(storage)
    await queue.register(
        DropBucket,
        BucketRequest(1),
        defer=timedelta(milliseconds=10),
    )

    deferred_pipe: str = Tube(DeferredHandler, Route()).pipe
    drop_pipe: str = Tube(DropBucket, Route()).pipe
    worker = Worker(RuntimeLock(), queue)
    worker.settings.deferred_retry_delay = 0

    assert await queue.storage.length(drop_pipe) == 0
    assert await queue.storage.length(deferred_pipe) == 1

    await worker.loop(1)
    assert await queue.storage.length(drop_pipe) == 0
    assert await queue.storage.length(deferred_pipe) == 1

    await sleep(0.01)

    await worker.loop(1)
    assert await queue.storage.length(drop_pipe) == 1
    assert await queue.storage.length(deferred_pipe) == 0

    await worker.loop(1)
    assert await queue.storage.length(drop_pipe) == 0
    assert await queue.storage.length(deferred_pipe) == 0

tests/test_deferred.py
basis-company/sharded-queue.py
from dataclasses import dataclass

from pytest import mark

from sharded_queue import Handler, Queue, Route, Tube
from sharded_queue.drivers import RuntimeStorage


@dataclass
class CompanyRequest:
    company: int


class ProlongateContracts(Handler):
    async def handle(self, *requests: CompanyRequest) -> None:
        pass


@mark.asyncio
async def test_duplicate() -> None:
    storage = RuntimeStorage()
    queue: Queue = Queue(storage)

    await queue.register(
        ProlongateContracts,
        CompanyRequest(1),
        if_not_exists=True,
    )

    assert 1 == await storage.length(Tube(ProlongateContracts, Route()).pipe)

    await queue.register(
        ProlongateContracts,
        CompanyRequest(1),
        if_not_exists=True,
    )

    assert 1 == await storage.length(Tube(ProlongateContracts, Route()).pipe)

    await queue.register(
        ProlongateContracts,
        CompanyRequest(1),
        if_not_exists=False,
    )

    assert 2 == await storage.length(Tube(ProlongateContracts, Route()).pipe)

tests/test_duplicate.py
basis-company/sharded-queue.py
from pytest import mark, raises

from sharded_queue import Handler, Queue, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage
from sharded_queue.protocols import Storage


class AbstractRequest:
    ...


class AbstractHandler(Handler):
    ...


@mark.asyncio
async def test_handler() -> None:
    storage: Storage = RuntimeStorage()
    queue: Queue = Queue(storage)
    await queue.register(AbstractHandler, AbstractRequest())

    with raises(NotImplementedError):
        await Worker(RuntimeLock(), queue).loop(1)

tests/test_handler.py
basis-company/sharded-queue.py
from asyncio import gather, sleep

from pytest import mark
from redis.asyncio import Redis

from sharded_queue import Handler, Queue, Route, Tube, Worker
from sharded_queue.drivers import RedisLock, RuntimeLock, RuntimeStorage
from sharded_queue.protocols import Lock


class Request:
    ...


class TestHandler(Handler):
    async def handle(self, *requests: Request) -> None:
        ...


@mark.asyncio
async def test_lock_prolongate() -> None:
    queue: Queue = Queue(RuntimeStorage())
    redis: Redis = Redis(decode_responses=True)
    await redis.flushall()
    worker: Worker = Worker(RedisLock(redis), queue)
    pipe: str = worker.lock.settings.prefix + Tube(TestHandler, Route()).pipe

    async def register_task():
        ttl1 = await redis.ttl(pipe)
        await sleep(1)
        ttl2 = await redis.ttl(pipe)
        assert ttl1 == ttl2, 'ttl is updated'
        await queue.register(TestHandler, Request())

    await queue.register(TestHandler, Request())
    await gather(worker.loop(2), register_task())
    await redis.close()


@mark.asyncio
async def test_runtime_storage() -> None:
    await runner(RuntimeLock())


@mark.asyncio
async def test_redis_storage() -> None:
    redis: Redis = Redis(decode_responses=True)
    await redis.flushall()
    await runner(RedisLock(redis))
    await redis.close()


async def runner(lock: Lock):
    assert await lock.acquire('tester')
    assert await lock.exists('tester')
    assert not await lock.acquire('tester')
    assert await lock.exists('tester')
    assert await lock.ttl('tester', 1)

    await lock.release('tester')
    assert not await lock.exists('tester')
    assert not await lock.ttl('tester', 1)

    assert await lock.acquire('tester')
    await lock.release('tester')

tests/test_lock.py
basis-company/sharded-queue.py
from asyncio import gather
from typing import NamedTuple

from pytest import mark

from sharded_queue import Handler, Queue, Route, Tube, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage


class ActionMessage(NamedTuple):
    id: int


class Action(NamedTuple):
    id: int


class BaseAction(Handler):
    async def handle(self, *requests: ActionMessage) -> None:
        await gather(*[
            self.perform(Action(request.id))
            for request in requests
        ])

    async def perform(self, action: Action) -> None:
        ...


class CreateUserAction(BaseAction):
    async def perform(self, action: Action) -> None:
        ...


class UpdateUserAction(BaseAction):
    async def perform(self, action: Action) -> None:
        ...


@mark.asyncio
async def test_nested() -> None:
    queue: Queue = Queue(RuntimeStorage())

    await queue.register(CreateUserAction, Action(1))
    await queue.register(UpdateUserAction, Action(2))
    create_user_pipe = Tube(CreateUserAction, Route(0, 0)).pipe
    update_user_pipe = Tube(UpdateUserAction, Route(0, 0)).pipe

    assert await queue.storage.length(create_user_pipe) == 1
    assert await queue.storage.length(update_user_pipe) == 1

    worker = Worker(RuntimeLock(), queue)
    worker.settings.empty_pause = 0
    await worker.loop(2)

    assert await queue.storage.length(create_user_pipe) == 0
    assert await queue.storage.length(update_user_pipe) == 0

tests/test_nested.py
basis-company/sharded-queue.py
from asyncio import gather
from contextlib import asynccontextmanager
from datetime import datetime
from typing import AsyncGenerator, NamedTuple

from pydantic_settings import BaseSettings
from pytest import mark

from sharded_queue import Handler, Queue, Route, Tube, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage


class BenchmarkSettings(BaseSettings):
    threads: int = 1
    requests: int = 100_000


class PerformaceMatrix:
    requests: list[int] = [100, 10_000]
    threads: list[int] = [1, 2, 5, 10, 20, 50]


benchmark = BenchmarkSettings()


class DummyRequest(NamedTuple):
    n: int


class DummyHandler(Handler):
    @classmethod
    async def route(cls, *requests: DummyRequest) -> list[Route]:
        return [Route(r.n % benchmark.threads) for r in requests]

    async def handle(self, *requests: DummyRequest) -> None:
        pass


@mark.asyncio
async def test_performance() -> None:
    print('start benchmark')
    queue: Queue = Queue(RuntimeStorage())
    worker: Worker = Worker(RuntimeLock(), queue)
    worker.settings.batch_size = 512
    worker.settings.empty_limit = 0
    worker.settings.empty_pause = 0
    for requests in PerformaceMatrix.requests:
        benchmark.requests = requests
        for threads in PerformaceMatrix.threads:
            benchmark.threads = threads
            print(f'test {requests} requests using {threads} thread(s)')
            async with measure('registration'):
                await queue.register(
                    DummyHandler,
                    *[DummyRequest(n) for n in range(1, requests)]
                )
            async with measure('one worker'):
                await worker.loop(requests-1)

            for thread in range(1, threads):
                pipe = Tube(DummyHandler, Route(thread)).pipe
                assert await queue.storage.length(pipe) == 0

            await queue.register(
                DummyHandler,
                *[DummyRequest(n) for n in range(1, requests)]
            )

            async with measure('worker threads'):
                await gather(*[
                    worker.loop(
                        int(requests-1/threads) - 1
                    )
                ])


@asynccontextmanager
async def measure(label: str) -> AsyncGenerator:
    start = datetime.now().timestamp()
    yield
    time = datetime.now().timestamp() - start
    rps = round(benchmark.requests / time)
    time = round(time, 3)

    print(f'{label} [rps: {rps}]')

tests/test_performance.py
basis-company/sharded-queue.py
from typing import NamedTuple

from pytest import mark

from sharded_queue import Handler, Queue, Route, Tube, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage


class TimesheetSign(NamedTuple):
    timesheet: int


processed: list[TimesheetSign] = []


class TimesheetSignHandler(Handler):
    priorities: list[int] = [3, 2, 1]

    @classmethod
    async def route(cls, *requests: TimesheetSign) -> list[Route]:
        return [
            Route(priority=r.timesheet) for r in requests
        ]

    async def handle(self, *requests: TimesheetSign) -> None:
        processed.extend(requests)


def get_pipe(priority: int) -> str:
    return Tube(TimesheetSignHandler, Route(0, priority)).pipe


@mark.asyncio
async def test_piorities() -> None:
    storage = RuntimeStorage()
    queue: Queue = Queue(storage)
    await queue.register(TimesheetSignHandler, *[
        TimesheetSign(n) for n in range(1, 4)
    ])

    worker = Worker(RuntimeLock(), queue)
    await worker.loop(1)
    assert len(processed) == 1

    assert await storage.length(get_pipe(1)) == 1
    assert await storage.length(get_pipe(2)) == 1
    assert await storage.length(get_pipe(3)) == 0

    await worker.loop(1)
    assert len(processed) == 2

    assert await storage.length(get_pipe(1)) == 1
    assert await storage.length(get_pipe(2)) == 0
    assert await storage.length(get_pipe(3)) == 0

    await worker.loop(1)
    assert len(processed) == 3

    assert await storage.length(get_pipe(1)) == 0
    assert await storage.length(get_pipe(2)) == 0
    assert await storage.length(get_pipe(3)) == 0

tests/test_priority.py
basis-company/sharded-queue.py
from asyncio import sleep
from datetime import timedelta
from typing import NamedTuple

from pytest import mark

from sharded_queue import (DeferredHandler, Handler, Queue, RecurrentHandler,
                           Route, Tube, Worker)
from sharded_queue.drivers import RuntimeLock, RuntimeStorage
from sharded_queue.protocols import Lock, Storage


class CompanyRequest(NamedTuple):
    company: int


class ValidateAccess(Handler):
    async def handle(self, *requests: CompanyRequest) -> None:
        pass


@mark.asyncio
async def test_recurrent() -> None:
    storage: Storage = RuntimeStorage()
    lock: Lock = RuntimeLock()
    queue: Queue = Queue(storage)

    worker = Worker(lock, queue)
    worker.settings.deferred_retry_delay = 0
    worker.settings.recurrent_check_interval = 1

    deferred_pipe: str = Tube(DeferredHandler, Route()).pipe
    recurrent_pipe: str = Tube(RecurrentHandler, Route()).pipe
    validatation_pipe: str = Tube(ValidateAccess, Route()).pipe

    async def stats() -> tuple[int, int, int]:
        return (
            await queue.storage.length(deferred_pipe),
            await queue.storage.length(recurrent_pipe),
            await queue.storage.length(validatation_pipe),
        )

    await queue.register(
        ValidateAccess, CompanyRequest(1), recurrent=timedelta(milliseconds=10)
    )
    assert await stats() == (0, 1, 0), 'recurrent pipe contains request'

    await Worker(lock, queue).loop(1)
    assert await stats() == (1, 1, 0), 'added defered request'
    assert await lock.exists(recurrent_pipe)

    await lock.release(recurrent_pipe)
    await Worker(lock, queue).loop(1, handler=RecurrentHandler)
    assert await stats() == (1, 1, 0), 'no deffered duplicates'

    await sleep(0.01)

    await lock.release(recurrent_pipe)
    await Worker(lock, queue).loop(1, handler=RecurrentHandler)
    assert await stats() == (1, 1, 0), 'no deffered duplicates'

    await worker.loop(1, handler=DeferredHandler)
    assert await stats() == (0, 1, 1), 'deferred processed, validation added'

    await lock.release(recurrent_pipe)
    await worker.loop(1, handler=RecurrentHandler)
    assert await stats() == (1, 1, 1), 'deferred added'

    await sleep(0.01)

    await worker.loop(1, handler=DeferredHandler)
    assert await stats() == (0, 1, 1), 'no validation duplicate'

tests/test_recurrent.py
basis-company/sharded-queue.py
from dataclasses import dataclass
from typing import Any, NamedTuple

from pydantic import BaseModel

from sharded_queue.drivers import JsonTupleSerializer
from sharded_queue.protocols import Serializer


class NamedTupleRequest(NamedTuple):
    name: str


@dataclass
class DataclassRequest:
    name: str


class PydanticRequest(BaseModel):
    name: str


def test_dataclass() -> None:
    runner(JsonTupleSerializer(), DataclassRequest('nekufa'))


def test_named_tuple() -> None:
    runner(JsonTupleSerializer(), NamedTupleRequest('nekufa'))


def test_pydantic() -> None:
    runner(JsonTupleSerializer(), PydanticRequest(name='nekufa'))


def runner(serializer: Serializer, request: Any):
    serialized = serializer.serialize(request)
    parsed = serializer.deserialize(type(request), serialized)
    assert isinstance(parsed, type(request))
    assert parsed.name == request.name

tests/test_serializer.py
basis-company/sharded-queue.py
from datetime import datetime
from typing import NamedTuple

from pytest import mark

from sharded_queue import Handler, Queue, Route, Tube, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage


class Request(NamedTuple):
    bucket: int


context: dict[str, float] = {
    'started': 0,
    'stopped': 0
}

synced: list[Request] = []


class SyncHandler(Handler):
    @classmethod
    async def route(cls, *requests: Request) -> list[Route]:
        return [
            Route(thread=request.bucket % 2) for request in requests
        ]

    async def start(self) -> None:
        context['started'] = datetime.now().timestamp()

    async def stop(self) -> None:
        context['stopped'] = datetime.now().timestamp()

    async def handle(self, *requests: Request) -> None:
        synced.extend(requests)


@mark.asyncio
async def test_queue() -> None:
    storage = RuntimeStorage()
    queue: Queue = Queue(storage)

    def get_pipe(thread: int) -> str:
        return Tube(SyncHandler, Route(thread=thread)).pipe

    await queue.register(
        SyncHandler,
        Request(1),
        Request(2),
        Request(3),
        Request(4),
        Request(5),
    )

    assert get_pipe(0) in storage.data
    assert get_pipe(1) in storage.data
    assert get_pipe(2) not in storage.data

    assert await storage.length(get_pipe(0)) == 2
    assert await storage.length(get_pipe(1)) == 3
    assert await storage.length(get_pipe(2)) == 0

    assert context['started'] == 0
    assert context['stopped'] == 0

    worker = Worker(RuntimeLock(), queue)
    await worker.loop(2)

    assert context['started'] != 0
    assert context['stopped'] != 0
    assert context['started'] != context['stopped']

    assert len(synced) == 2
    thread = str(synced[0].bucket % 2)

    if thread == '0':
        assert await storage.length(get_pipe(0)) == 0
        assert await storage.length(get_pipe(1)) == 3
    else:
        assert await storage.length(get_pipe(0)) == 2
        assert await storage.length(get_pipe(1)) == 1

    # test worker switch
    worker.settings.empty_pause = 0
    (last_started, last_stopped) = tuple(context.values())
    await worker.loop(3)
    assert len(synced) == 5
    assert await storage.length(get_pipe(0)) == 0
    assert await storage.length(get_pipe(1)) == 0
    assert last_started != context['started']
    assert last_stopped != context['stopped']

tests/test_sharding.py
basis-company/sharded-queue.py
from pytest import mark
from redis.asyncio import Redis

from sharded_queue.drivers import RedisStorage, RuntimeStorage


@mark.asyncio
async def test_runtime_storage() -> None:
    await runner(RuntimeStorage())


@mark.asyncio
async def test_redis_storage() -> None:
    redis: Redis = Redis(decode_responses=True)
    await redis.flushall()
    await runner(RedisStorage(redis))
    await redis.close()


async def runner(storage) -> None:
    await storage.append('tester', 'q')
    assert not await storage.contains('tester', 'w')
    await storage.append('tester', 'w')
    assert await storage.contains('tester', 'w')
    await storage.append('tester', 'e')
    assert await storage.length('tester') == 3
    assert await storage.pipes() == ['tester']
    await storage.append('tester', 'r', 't', 'y')
    assert await storage.length('tester') == 6
    assert await storage.length('tester2') == 0
    assert await storage.range('tester', 1) == ['q']
    assert await storage.range('tester', 2) == ['q', 'w']
    assert await storage.range('tester', 3) == ['q', 'w', 'e']
    assert await storage.pop('tester', 1) == ['q']
    assert await storage.range('tester', 1) == ['w']
    assert await storage.range('tester', 2) == ['w', 'e']
    assert await storage.pop('tester', 2) == ['w', 'e']
    assert await storage.pop('tester', 10) == ['r', 't', 'y']
    assert await storage.pop('tester', 1) == []

tests/test_storage.py
basis-company/sharded-queue.py
from asyncio import get_event_loop, sleep
from typing import NamedTuple

from pytest import mark, raises

from sharded_queue import Handler, Queue, Worker
from sharded_queue.drivers import RuntimeLock, RuntimeStorage


class SignContractRequest(NamedTuple):
    document: int


signed: list[SignContractRequest] = []


class SignContract(Handler):
    async def handle(self, *requests: SignContractRequest) -> None:
        signed.extend(requests)


@mark.asyncio
async def test_worker_pause() -> None:
    queue: Queue = Queue(RuntimeStorage())
    worker = Worker(RuntimeLock(), queue)
    worker.settings.empty_pause = 0.1
    working_loop = worker.loop(2)
    worker_task = get_event_loop().create_task(working_loop)
    get_event_loop().create_task(working_loop)
    await queue.register(SignContract, SignContractRequest(1))
    await sleep(0.2)
    assert len(signed) == 1
    await queue.register(SignContract, SignContractRequest(2))
    await worker_task
    assert len(signed) == 2


@mark.asyncio
async def test_worker_prolongate_invalid_pipe() -> None:
    queue: Queue = Queue(RuntimeStorage())
    worker = Worker(RuntimeLock(), queue)
    with raises(RuntimeError):
        await worker.prolongate_lock()

tests/test_worker.py
basis-company/sharded-queue.py