app/routers/take.py

@@ -1,5 +1,3 @@
-import asyncio
-import collections
 import fastapi
 import pydantic
 import typing
@@ -7,7 +5,7 @@ import typing
 from app.storage.mongo import tasks
 
 
-locks = collections.defaultdict(asyncio.Lock)
+DEFAULT_RETRY_AFTER = 0.2
 router = fastapi.APIRouter()
 
 
@@ -21,10 +19,13 @@ class Response(pydantic.BaseModel):
     task: Task|None
 
 
-@router.get('/api/v1/take', responses={404: {'description': 'Not found'}})
+@router.get('/api/v1/take')
 async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
-    async with locks[queue]:
+    try:
         task = await tasks.take_task(queue)
-    if not task:
+        if not task:
+            return Response(task=None)
+        return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))
+    except Exception as e:
+        print("GOT ERROR", e)
         return Response(task=None)
-    return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))