fix(bot): 🐛 Main process crashing on shard restart.

This commit is contained in:
Nyeki.py
2024-10-25 20:29:57 -03:00
parent 00abb50c3e
commit 1c77532a7f
+20 -3
View File
@@ -1,6 +1,6 @@
from multiprocessing import managers from multiprocessing import managers
from datetime import datetime, timezone from datetime import datetime, timezone, timedelta
def shard( def shard(
@@ -8,6 +8,7 @@ def shard(
shard_count: int, shard_count: int,
shutdown: managers.ValueProxy[bool], shutdown: managers.ValueProxy[bool],
scheduled_shutdown: managers.ValueProxy[bool], scheduled_shutdown: managers.ValueProxy[bool],
heartbeat: managers.ValueProxy[int]
) -> None: ) -> None:
"""Starts a new shard. """Starts a new shard.
@@ -17,6 +18,7 @@ def shard(
shutdown (ValueProxy[bool]): The flag for the main process to shut down all shards. 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 scheduled_shutdown (ValueProxy[bool]): The flag for other shards to know the bot's gonna
shut down soon. shut down soon.
heartbeat (ValueProxy[int]): A timestamp updated every second by the shard's process.
""" """
from payphone.bot import bot from payphone.bot import bot
@@ -26,6 +28,7 @@ def shard(
bot.shard_count = shard_count bot.shard_count = shard_count
bot.shutdown = shutdown bot.shutdown = shutdown
bot.scheduled_shutdown = scheduled_shutdown bot.scheduled_shutdown = scheduled_shutdown
bot.heartbeat = heartbeat
try: try:
import uvloop # pylint: disable=import-error import uvloop # pylint: disable=import-error
@@ -106,11 +109,14 @@ if __name__ == "__main__":
manager = mp.Manager() manager = mp.Manager()
shutdown = manager.Value("b", False) shutdown = manager.Value("b", False)
scheduled_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 ""}.") print(f"Starting {shard_count} shard{"s" if shard_count != 1 else ""}.")
ps: list[mp.Process] = [ 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) for i in range(shard_count)
] ]
flushing_p = mp.Process(target=flush_process) flushing_p = mp.Process(target=flush_process)
@@ -142,11 +148,22 @@ if __name__ == "__main__":
raise ShutdownException() raise ShutdownException()
for i in range(shard_count): for i in range(shard_count):
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(): if not ps[i].is_alive():
print(f"Shard {i}'s process exited. Restarting the shard...") 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( 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() ps[i].start()