feat(bot): Update stuff to migrate to k8s.
This commit is contained in:
Nyeki.py
+211
@@ -0,0 +1,211 @@
|
||||
from multiprocessing import managers
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from payphone.core import logging
|
||||
|
||||
|
||||
log = logging.logger
|
||||
|
||||
|
||||
def shard(
|
||||
shard_id: int,
|
||||
shard_count: int,
|
||||
shutdown: managers.ValueProxy[bool],
|
||||
scheduled_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.
|
||||
scheduled_shutdown (ValueProxy[bool]): The flag for other shards to know the bot's gonna
|
||||
shut down soon.
|
||||
"""
|
||||
|
||||
from payphone.db import setup
|
||||
from payphone.bot import bot
|
||||
from payphone.core import config, logging
|
||||
|
||||
bot.shard_ids = [shard_id]
|
||||
bot.shard_count = shard_count
|
||||
bot.shutdown = shutdown
|
||||
bot.scheduled_shutdown = scheduled_shutdown
|
||||
|
||||
logging.setup(f"Shard {shard_id}")
|
||||
|
||||
log.info("Running shard %s with default loop.", shard_id)
|
||||
|
||||
bot.loop.set_debug(True)
|
||||
|
||||
async def initialize():
|
||||
await setup.init()
|
||||
await bot.redis.ping()
|
||||
|
||||
bot.loop.run_until_complete(initialize())
|
||||
|
||||
bot.run(config.TOKEN)
|
||||
|
||||
|
||||
def flush_process():
|
||||
import asyncio
|
||||
|
||||
from procs.flush_messages import loop
|
||||
|
||||
try:
|
||||
import uvloop # pylint: disable=import-error
|
||||
|
||||
uvloop.install()
|
||||
log.info("Flushing messages with uvloop.")
|
||||
|
||||
except ImportError:
|
||||
log.info("Flushing messages with default loop.")
|
||||
|
||||
finally:
|
||||
asyncio.run(loop())
|
||||
|
||||
|
||||
class ShutdownException(Exception):
|
||||
"""Raised when a shard has triggered the shutdown of all shards."""
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import multiprocessing as mp
|
||||
|
||||
import time
|
||||
|
||||
from tortoise import run_async
|
||||
|
||||
import httpx
|
||||
|
||||
import redis.asyncio as aredis
|
||||
|
||||
from payphone.db import setup, models
|
||||
from payphone.core import config
|
||||
|
||||
logging.setup("Payphone")
|
||||
|
||||
redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")
|
||||
|
||||
async def init():
|
||||
await setup.init()
|
||||
|
||||
await models.Call.filter(ended_at=None).update(
|
||||
ended_at=datetime.now(timezone.utc)
|
||||
)
|
||||
await models.Channel.exclude(call_id=None).update(call_id=None)
|
||||
|
||||
await redis.flushall()
|
||||
|
||||
log.info("Initializing DB and running Redis cleanup...")
|
||||
|
||||
run_async(init())
|
||||
|
||||
log.info("Cleanup done! Loading recommended shard amount...")
|
||||
|
||||
response = httpx.get(
|
||||
"https://discord.com/api/v10/gateway/bot",
|
||||
headers={"Authorization": f"Bot {config.TOKEN}"},
|
||||
)
|
||||
|
||||
info = response.json()
|
||||
|
||||
shard_count = info["shards"]
|
||||
|
||||
manager = mp.Manager()
|
||||
shutdown = manager.Value("b", False)
|
||||
scheduled_shutdown = manager.Value("b", False)
|
||||
|
||||
log.info("Starting %s shards.", shard_count)
|
||||
|
||||
ps: list[mp.Process] = [
|
||||
mp.Process(
|
||||
target=shard,
|
||||
args=(i, shard_count, shutdown, scheduled_shutdown),
|
||||
)
|
||||
for i in range(shard_count)
|
||||
]
|
||||
flushing_p = mp.Process(target=flush_process)
|
||||
|
||||
for i, p in enumerate(ps):
|
||||
p.start()
|
||||
|
||||
log.info("Shard %s's process started.", i)
|
||||
log.info(
|
||||
"Waiting 5 seconds start initialize next shard..."
|
||||
if i != len(ps) - 1
|
||||
else "All shards were started."
|
||||
)
|
||||
|
||||
if i != len(ps) - 1:
|
||||
time.sleep(5)
|
||||
|
||||
log.info("Starting messages flushing process...")
|
||||
|
||||
flushing_p.start()
|
||||
|
||||
log.info("Messages flushing process started.")
|
||||
|
||||
try:
|
||||
while True:
|
||||
if shutdown.value:
|
||||
raise ShutdownException()
|
||||
|
||||
for i in range(shard_count):
|
||||
if not ps[i].is_alive():
|
||||
log.warning("Shard %s's process exited. Restarting the shard...", i)
|
||||
|
||||
ps[i] = mp.Process(
|
||||
target=shard,
|
||||
args=(
|
||||
i,
|
||||
shard_count,
|
||||
shutdown,
|
||||
scheduled_shutdown,
|
||||
),
|
||||
name=f"payphone-shard-{i}",
|
||||
)
|
||||
ps[i].start()
|
||||
|
||||
log.info(
|
||||
"Shard %s was restarted. Waiting 5 seconds to check next shard...",
|
||||
i,
|
||||
)
|
||||
time.sleep(5)
|
||||
|
||||
if not flushing_p.is_alive():
|
||||
log.info("Messages flushing process was stopped. Restarting it...")
|
||||
|
||||
flushing_p = mp.Process(target=flush_process, name="payphone-flushing")
|
||||
flushing_p.start()
|
||||
|
||||
log.info("Messages flushing process started.")
|
||||
|
||||
time.sleep(5)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
log.info("Ctrl-C was pressed.")
|
||||
|
||||
except ShutdownException:
|
||||
log.info("A shard shut down the bot.")
|
||||
|
||||
finally:
|
||||
log.info("Terminating %s shards...", shard_count)
|
||||
|
||||
for i, p in enumerate(ps):
|
||||
p.terminate()
|
||||
log.info("Terminated shard %s.", i)
|
||||
|
||||
log.info("Stopping messages flushing process...")
|
||||
|
||||
flushing_p.terminate()
|
||||
|
||||
log.info("Checking that all processes have finished...")
|
||||
|
||||
for p in ps:
|
||||
p.join()
|
||||
|
||||
flushing_p.join()
|
||||
|
||||
log.info("All processes have been terminated.")
|
||||
Reference in New Issue
Block a user