diff --git a/run.py b/run.py index 7313d8b..bc98131 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,6 @@ from multiprocessing import managers -from datetime import datetime, timezone +from datetime import datetime, timezone, timedelta def shard( @@ -8,6 +8,7 @@ def shard( shard_count: int, shutdown: managers.ValueProxy[bool], scheduled_shutdown: managers.ValueProxy[bool], + heartbeat: managers.ValueProxy[int] ) -> None: """Starts a new shard. @@ -17,6 +18,7 @@ def shard( shutdown (ValueProxy[bool]): The flag for the main process to shut down all shards. scheduled_shutdown (ValueProxy[bool]): The flag for other shards to know the bot's gonna shut down soon. + heartbeat (ValueProxy[int]): A timestamp updated every second by the shard's process. """ from payphone.bot import bot @@ -26,6 +28,7 @@ def shard( bot.shard_count = shard_count bot.shutdown = shutdown bot.scheduled_shutdown = scheduled_shutdown + bot.heartbeat = heartbeat try: import uvloop # pylint: disable=import-error @@ -106,11 +109,14 @@ if __name__ == "__main__": manager = mp.Manager() shutdown = manager.Value("b", False) scheduled_shutdown = manager.Value("b", False) + heartbeats = [ + manager.Value("i", 0) for i in range(shard_count) + ] print(f"Starting {shard_count} shard{"s" if shard_count != 1 else ""}.") ps: list[mp.Process] = [ - mp.Process(target=shard, args=(i, shard_count, shutdown, scheduled_shutdown)) + mp.Process(target=shard, args=(i, shard_count, shutdown, scheduled_shutdown, heartbeats[i])) for i in range(shard_count) ] flushing_p = mp.Process(target=flush_process) @@ -142,11 +148,22 @@ if __name__ == "__main__": raise ShutdownException() for i in range(shard_count): - if not ps[i].is_alive(): - print(f"Shard {i}'s process exited. Restarting the shard...") + is_heartbeat_updated = heartbeats[i].value > (datetime.now(timezone.utc) - timedelta(seconds=10)).timestamp() + + if not ps[i].is_alive() or not is_heartbeat_updated: + if not ps[i].is_alive(): + print(f"Shard {i}'s process exited. Restarting the shard...") + + else: + print(f"Shard {i} did not update heartbeat. Terminating...") + + ps[i].terminate() + ps[i].join() + + print(f"Shard {i}'s process was terminated. Restarting the shard...") ps[i] = mp.Process( - target=shard, args=(i, shard_count), name=f"payphone-shard-{i}" + target=shard, args=(i, shard_count, shutdown, scheduled_shutdown, heartbeats[i]), name=f"payphone-shard-{i}" ) ps[i].start()