feat(bot): ✨ Terminate all processes after p.shutdown.
This commit is contained in:
@@ -1,12 +1,15 @@
|
||||
from multiprocessing import managers
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
def shard(shard_id: int, shard_count: int) -> None:
|
||||
def shard(shard_id: int, shard_count: int, shutdown: managers.ValueProxy[bool]) -> None:
|
||||
"""Starts a new shard.
|
||||
|
||||
Args:
|
||||
shard_id (int): The shard's ID.
|
||||
shard_count (int): The total amount of shards.
|
||||
shutdown (ValueProxy[bool]): The flag for the main process to shut down all shards.
|
||||
"""
|
||||
|
||||
from payphone.bot import bot
|
||||
@@ -14,6 +17,7 @@ def shard(shard_id: int, shard_count: int) -> None:
|
||||
|
||||
bot.shard_ids = [shard_id]
|
||||
bot.shard_count = shard_count
|
||||
bot.shutdown = shutdown
|
||||
|
||||
try:
|
||||
import uvloop # pylint: disable=import-error
|
||||
@@ -46,6 +50,10 @@ def flush_process():
|
||||
asyncio.run(loop())
|
||||
|
||||
|
||||
class ShutdownException(Exception):
|
||||
"""Raised when a shard has triggered the shutdown of all shards."""
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import multiprocessing as mp
|
||||
|
||||
@@ -87,10 +95,14 @@ if __name__ == "__main__":
|
||||
|
||||
shard_count = info["shards"]
|
||||
|
||||
manager = mp.Manager()
|
||||
shutdown = manager.Value("b", False)
|
||||
|
||||
print(f"Starting {shard_count} shard{"s" if shard_count != 1 else ""}.")
|
||||
|
||||
ps: list[mp.Process] = [
|
||||
mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
|
||||
mp.Process(target=shard, args=(i, shard_count, shutdown))
|
||||
for i in range(shard_count)
|
||||
]
|
||||
flushing_p = mp.Process(target=flush_process)
|
||||
|
||||
@@ -101,12 +113,12 @@ if __name__ == "__main__":
|
||||
f"Shard {i}'s process started."
|
||||
+ (
|
||||
" Waiting 5 seconds start initialize next shard..."
|
||||
if len(ps) - 1 < i
|
||||
if i != len(ps) - 1
|
||||
else " All shards were started."
|
||||
)
|
||||
)
|
||||
|
||||
if i + 1 != len(ps):
|
||||
if i != len(ps) - 1:
|
||||
time.sleep(5)
|
||||
|
||||
print("Starting messages flushing process...")
|
||||
@@ -117,11 +129,16 @@ if __name__ == "__main__":
|
||||
|
||||
try:
|
||||
while True:
|
||||
if shutdown.value:
|
||||
raise ShutdownException()
|
||||
|
||||
for i in range(shard_count):
|
||||
if not ps[i].is_alive():
|
||||
print(f"Shard {i}'s process exited. Restarting the shard...")
|
||||
|
||||
ps[i] = mp.Process(target=shard, args=(i, shard_count))
|
||||
ps[i] = mp.Process(
|
||||
target=shard, args=(i, shard_count), name=f"payphone-shard-{i}"
|
||||
)
|
||||
ps[i].start()
|
||||
|
||||
print(
|
||||
@@ -132,7 +149,7 @@ if __name__ == "__main__":
|
||||
if not flushing_p.is_alive():
|
||||
print("Messages flushing process was stopped. Restarting it...")
|
||||
|
||||
flushing_p = mp.Process(target=flush_process)
|
||||
flushing_p = mp.Process(target=flush_process, name="payphone-flushing")
|
||||
flushing_p.start()
|
||||
|
||||
print("Messages flushing process started.")
|
||||
@@ -140,7 +157,13 @@ if __name__ == "__main__":
|
||||
time.sleep(5)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"Stopping {shard_count} shards...")
|
||||
print("Ctrl-C was pressed.")
|
||||
|
||||
except ShutdownException:
|
||||
print("A shard shut down the bot.")
|
||||
|
||||
finally:
|
||||
print(f"Terminating {shard_count} shards...")
|
||||
|
||||
for i, p in enumerate(ps):
|
||||
p.terminate()
|
||||
|
||||
Reference in New Issue
Block a user