distributed.camp
A Python port of the Invisible Internet Project (I2P)
1"""GracefulShutdown — orderly router shutdown with tunnel drain.
2
3Ported from net.i2p.router.Router (shutdown logic) and
4net.i2p.router.tasks.GracefulShutdown.
5
6Graceful shutdown polls participating tunnel count at intervals.
7When count reaches 0 (all tunnels expired), the final shutdown proceeds.
8Hard shutdown waits a short delay then forces immediate stop.
9"""
10
11from __future__ import annotations
12
13import asyncio
14import logging
15import time
16
17logger = logging.getLogger(__name__)
18
19
20class GracefulShutdown:
21 """Manages orderly router shutdown with tunnel drain."""
22
23 POLL_INTERVAL_SECONDS = 10
24 HARD_SHUTDOWN_DELAY_SECONDS = 2
25
26 def __init__(self, router) -> None:
27 self._router = router
28 self._shutting_down = False
29 self._shutdown_initiated_at: float | None = None
30
31 @property
32 def is_shutting_down(self) -> bool:
33 return self._shutting_down
34
35 def get_participating_count(self) -> int:
36 """Get current participating tunnel count from router context."""
37 ctx = getattr(self._router, '_context', None)
38 if ctx is None:
39 return 0
40 tunnel_mgr = getattr(ctx, 'tunnel_manager', None)
41 if tunnel_mgr is None:
42 return 0
43 return getattr(tunnel_mgr, 'participating_count', 0)
44
45 async def start_graceful(self) -> None:
46 """Begin graceful shutdown: drain tunnels then stop.
47
48 Polls participating count every POLL_INTERVAL_SECONDS.
49 Once all tunnels have expired, calls router.shutdown().
50 """
51 if self._shutting_down:
52 return
53
54 self._shutting_down = True
55 self._shutdown_initiated_at = time.monotonic()
56 logger.info("Graceful shutdown initiated, draining tunnels...")
57
58 while self.get_participating_count() > 0:
59 count = self.get_participating_count()
60 elapsed = time.monotonic() - self._shutdown_initiated_at
61 logger.info(
62 "Graceful shutdown: %d tunnels remaining (%.0fs elapsed)",
63 count, elapsed,
64 )
65 await asyncio.sleep(self.POLL_INTERVAL_SECONDS)
66
67 logger.info("All tunnels drained, completing shutdown")
68 self._router.shutdown()
69
70 async def start_hard(self) -> None:
71 """Immediate shutdown after a short delay."""
72 if self._shutting_down:
73 return
74
75 self._shutting_down = True
76 logger.info(
77 "Hard shutdown in %ds...", self.HARD_SHUTDOWN_DELAY_SECONDS,
78 )
79 await asyncio.sleep(self.HARD_SHUTDOWN_DELAY_SECONDS)
80 self._router.shutdown()