Merge pull request #209 from DABND19/fix/uvicorn-service

fix: Uvicorn graceful shutdown.
aiokitchen · May 14, 2024 · 3d38ade · 3d38ade
2 parents 6d2f617 + a431c33
commit 3d38ade
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/aiomisc/service/uvicorn.py b/aiomisc/service/uvicorn.py
@@ -24,6 +24,7 @@
 class UvicornService(Service, abc.ABC):
     __async_required__: Tuple[str, ...] = (
         "start",
+        "stop",
         "create_application",
     )
 
@@ -114,8 +115,11 @@ async def start(self) -> Any:
         )
         if not self.sock:
             self.sock = config.bind_socket()
-        server = Server(config)
-
-        self.start_event.set()
+        self.server = Server(config)
+        self.serve_task = asyncio.create_task(
+            self.server.serve(sockets=[self.sock])
+        )
 
-        await server.serve(sockets=[self.sock])
+    async def stop(self, exception: Optional[Exception] = None) -> None:
+        self.server.should_exit = True
+        await self.serve_task