Coverage for src/lexigram/web/server/shutdown.py: 33%

1"""Graceful shutdown management for the web server. 

2 

3Provides GracefulShutdownManager for connection draining. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9from contextlib import asynccontextmanager 

10from typing import TYPE_CHECKING, Any 

11 

12if TYPE_CHECKING: 

13 from starlette.types import ASGIApp, Receive, Scope, Send 

14 

15 

16class GracefulShutdownManager: 

17 """Manages graceful shutdown with connection draining.""" 

18 

19 def __init__(self, timeout: float = 30.0): 

20 self._timeout = timeout 

21 self._active_connections: int = 0 

22 self._shutting_down: bool = False 

23 self._drain_event: asyncio.Event | None = None 

24 

25 @property 

26 def is_shutting_down(self) -> bool: 

27 """Check if shutdown is in progress.""" 

28 return self._shutting_down 

29 

30 @property 

31 def active_connections(self) -> int: 

32 """Get the number of active connections.""" 

33 return self._active_connections 

34 

35 async def wait_for_drain(self) -> None: 

36 """Wait for active connections to complete or timeout.""" 

37 if self._drain_event is None: 

38 self._drain_event = asyncio.Event() 

39 

40 try: 

41 await asyncio.wait_for( 

42 self._drain_event.wait(), 

43 timeout=self._timeout, 

44 ) 

45 except TimeoutError: 

46 # Timeout reached, proceed with shutdown anyway 

47 pass 

48 

49 def begin_shutdown(self) -> None: 

50 """Begin the shutdown process.""" 

51 self._shutting_down = True 

52 if self._drain_event: 

53 # If no active connections, signal immediately 

54 if self._active_connections == 0: 

55 self._drain_event.set() 

56 

57 def complete_shutdown(self) -> None: 

58 """Complete the shutdown process.""" 

59 if self._drain_event: 

60 self._drain_event.set() 

61 

62 @asynccontextmanager 

63 async def track_connection(self) -> Any: 

64 """Context manager to track active connections.""" 

65 self._active_connections += 1 

66 try: 

67 yield 

68 finally: 

69 self._active_connections -= 1 

70 # If shutting down and no more connections, signal 

71 if self._shutting_down and self._active_connections == 0: 

72 if self._drain_event: 

73 self._drain_event.set() 

74 

75 async def serve_503_during_drain(self, request: Any) -> Any: 

76 """Return a 503 response during drain period.""" 

77 from starlette.responses import JSONResponse 

78 

79 return JSONResponse( 

80 content={ 

81 "error": "service_unavailable", 

82 "message": "Server is shutting down", 

83 }, 

84 status_code=503, 

85 headers={"Retry-After": str(int(self._timeout))}, 

86 ) 

87 

88 

89class ShutdownMiddleware: 

90 """Middleware that tracks connections for graceful shutdown.""" 

91 

92 def __init__(self, app: ASGIApp, shutdown_manager: GracefulShutdownManager) -> None: 

93 self.app = app 

94 self.manager = shutdown_manager 

95 

96 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

97 if scope["type"] != "http": 

98 await self.app(scope, receive, send) 

99 return 

100 

101 # If shutting down, return 503 

102 if self.manager.is_shutting_down: 

103 response = await self.manager.serve_503_during_drain(None) 

104 await response(scope, receive, send) 

105 return 

106 

107 # Track the connection 

108 async with self.manager.track_connection(): 

109 await self.app(scope, receive, send)