nyeki/payphone
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(bot): ✨ Terminate all processes after p.shutdown.

Browse Source
This commit is contained in:
Nyeki.py
2024-10-24 00:22:07 -03:00
parent b6ac3d70c4
commit b5872df034
4 changed files with 57 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
payphone/bot.py
+15 -1
View File
@@ -2,6 +2,8 @@ from typing import Literal
from datetime import datetime, timezone, timedelta
from multiprocessing import managers
import sys
import discord
@@ -20,6 +22,13 @@ from payphone.core import constants, exceptions, formatting, help_command, confi
class Payphone(commands.AutoShardedBot):
shutdown: managers.ValueProxy[bool]
"""A value proxy that connects with all other shard processes.
This value is shared between shard processes, and if any of them sets this proxy's value to
`True`, then all shard processes will be terminated by the main process (`run.py`).
"""
redis: aredis.Redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")
def __init__(self) -> None:
@@ -398,7 +407,12 @@ class Payphone(commands.AutoShardedBot):
) -> None:
_, exception, _ = sys.exc_info()
if isinstance(exception, asyncpg.exceptions.InterfaceError):
if isinstance(
exception,
(asyncpg.exceptions.InterfaceError, asyncpg.exceptions.PostgresError),
):
# IMPORTANT! These are just spam at the moment, but may cause trouble in the future as
# the bot grows.
return
return await super().on_error(event_method, *args, **kwargs)
payphone/cogs/calls/cog.py
+5 -4
View File
@@ -9,8 +9,6 @@ from discord.utils import find
from tortoise import transactions, exceptions
import asyncpg.exceptions
import aiohttp
import orjson
@@ -19,7 +17,7 @@ from payphone import bot
from payphone.db import utils as dbutils, models
from payphone.cogs import base
from payphone.cogs.calls import embeds
from payphone.core import constants, checks, censorship, formatting, utils
from payphone.core import constants, checks, censorship, formatting
class Calls(base.Cog):
@@ -877,7 +875,10 @@ class Calls(base.Cog):
orjson.dumps(message_dict),
)
other_channel = await models.Channel.get(id=message.propagation_channel_id)
other_channel = await models.Channel.get_or_none(id=message.propagation_channel_id)
if other_channel is None:
return
if other_channel.call_id != message.call_id:
return
payphone/cogs/mods/cog.py
+7 -3
View File
@@ -243,7 +243,9 @@ class Mods(base.Cog):
- `[<reason>]`: The reason of the shutdown.
"""
oncall_channel_ids = list(await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET))
oncall_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET)
)
searching_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
)
@@ -270,7 +272,9 @@ class Mods(base.Cog):
await asyncio.sleep(seconds)
oncall_channel_ids = list(await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET))
oncall_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET)
)
searching_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
)
@@ -310,7 +314,7 @@ class Mods(base.Cog):
await ctx.send("⚡ Messages were flushed, payphone will shut down.")
sys.exit()
self.bot.shutdown.value = True
@commands.command()
@commands.is_owner()
run.py
+30 -7
View File
@@ -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()