1 files changed, 133 insertions, 0 deletions

diff --git a/proxy/queue.py b/proxy/queue.py
new file mode 100644
index 0000000..cfa0f3c
--- /dev/null
+++ b/proxy/queue.py
@@ -0,0 +1,133 @@
+import asyncio
+import json
+import logging
+import os
+import traceback
+from typing import Type
+
+from ttun_server.redis import RedisConnectionPool
+from ttun_server.types import Message
+
+logger = logging.getLogger(__name__)
+
+
+class BaseProxyQueue:
+    def __init__(self, identifier: str):
+        self.identifier = identifier
+
+    @classmethod
+    async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue':
+        raise NotImplementedError(f'Please implement create_for_identifier')
+
+    @classmethod
+    async def get_for_identifier(cls, identifier: str) -> Type['self']:
+        assert await cls.has_connection(identifier)
+        return cls(identifier)
+
+    @classmethod
+    async def has_connection(cls, identifier) -> bool:
+        raise NotImplementedError(f'Please implement has_connection')
+
+    async def enqueue(self, message: Message):
+        raise NotImplementedError(f'Please implement send_request')
+
+    async def dequeue(self) -> Message:
+        raise NotImplementedError(f'Please implement handle_requests')
+
+    async def delete(self):
+        raise NotImplementedError(f'Please implement delete')
+
+
+class MemoryProxyQueue(BaseProxyQueue):
+    connections: dict[str, asyncio.Queue] = {}
+
+    @classmethod
+    async def has_connection(cls, identifier) -> bool:
+        return identifier in cls.connections
+
+    @classmethod
+    async def create_for_identifier(cls, identifier: str) -> 'MemoryProxyQueue':
+        instance = cls(identifier)
+        cls.connections[identifier] = asyncio.Queue()
+
+        return instance
+
+    async def enqueue(self, message: Message):
+        return await self.__class__.connections[self.identifier].put(message)
+
+    async def dequeue(self) -> Message:
+        return await self.__class__.connections[self.identifier].get()
+
+    async def delete(self):
+        del self.__class__.connections[self.identifier]
+
+
+class RedisProxyQueue(BaseProxyQueue):
+    def __init__(self, identifier):
+        super().__init__(identifier)
+
+        self.pubsub = RedisConnectionPool()\
+            .get_connection()\
+            .pubsub()
+
+        self.subscription_queue = asyncio.Queue()
+
+    @classmethod
+    async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue':
+        instance = cls(identifier)
+
+        await instance.pubsub.subscribe(f'request_{identifier}')
+        return instance
+
+    @classmethod
+    async def get_for_identifier(cls, identifier: str) -> 'RedisProxyQueue':
+        instance: 'RedisProxyQueue' = await super().get_for_identifier(identifier)
+
+        await instance.pubsub.subscribe(f'response_{identifier}')
+
+        return instance
+
+    @classmethod
+    async def has_connection(cls, identifier) -> bool:
+        logger.debug(await RedisConnectionPool.get_connection().pubsub_channels())
+        return f'request_{identifier}' in {
+            channel.decode()
+            for channel
+            in await RedisConnectionPool \
+                .get_connection() \
+                .pubsub_channels()
+        }
+
+    async def wait_for_message(self):
+        async for message in self.pubsub.listen():
+            match message['type']:
+                case 'subscribe':
+                    continue
+                case _:
+                    return message['data']
+
+    async def enqueue(self, message: Message):
+        await RedisConnectionPool \
+            .get_connection() \
+            .publish(f'request_{self.identifier}', json.dumps(message))
+
+    async def dequeue(self) -> Message:
+        message = await self.wait_for_message()
+        return json.loads(message)
+
+    async def delete(self):
+        await self.pubsub.unsubscribe(f'request_{self.identifier}')
+
+        await RedisConnectionPool.get_connection()\
+            .srem('connections', self.identifier)
+
+
+class ProxyQueueMeta(type):
+    def __new__(cls, name, superclasses, attributes):
+        return RedisProxyQueue \
+            if 'REDIS_URL' in os.environ \
+            else MemoryProxyQueue
+
+
+class ProxyQueue(BaseProxyQueue, metaclass=ProxyQueueMeta):
+    pass