138 lines
3.3 KiB
Python
138 lines
3.3 KiB
Python
from datetime import datetime, timezone
|
|
|
|
|
|
def shard(shard_id: int, shard_count: int) -> None:
|
|
"""Starts a new shard.
|
|
|
|
Args:
|
|
shard_id (int): The shard's ID.
|
|
shard_count (int): The total amount of shards.
|
|
"""
|
|
|
|
from payphone.bot import bot
|
|
from payphone.core import config
|
|
|
|
bot.shard_ids = [shard_id]
|
|
bot.shard_count = shard_count
|
|
|
|
bot.run(config.TOKEN)
|
|
|
|
|
|
def flush_process():
|
|
import asyncio
|
|
|
|
from procs.flush_messages import loop
|
|
|
|
asyncio.run(loop())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import multiprocessing as mp
|
|
|
|
import time
|
|
|
|
from tortoise import run_async
|
|
|
|
import httpx
|
|
|
|
from payphone.db import setup, models
|
|
from payphone.core import config
|
|
from payphone.core.redis import 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()
|
|
|
|
print("Running DB and Redis cleanup...")
|
|
|
|
run_async(init())
|
|
|
|
print("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"]
|
|
|
|
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)
|
|
]
|
|
flushing_p = mp.Process(target=flush_process)
|
|
|
|
for i, p in enumerate(ps):
|
|
p.start()
|
|
|
|
print(
|
|
f"Shard {i}'s process started."
|
|
+ (
|
|
" Waiting 5 seconds start initialize next shard..."
|
|
if len(ps) - 1 < i
|
|
else " All shards were started."
|
|
)
|
|
)
|
|
|
|
if i + 1 != len(ps):
|
|
time.sleep(5)
|
|
|
|
print("Starting messages flushing process...")
|
|
|
|
flushing_p.start()
|
|
|
|
print("Messages flushing process started.")
|
|
|
|
try:
|
|
while True:
|
|
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].start()
|
|
|
|
print(
|
|
f"Shard {i} was restarted. Waiting 5 seconds to check next shard..."
|
|
)
|
|
time.sleep(5)
|
|
|
|
if not flushing_p.is_alive():
|
|
print("Messages flushing process was stopped. Restarting it...")
|
|
|
|
flushing_p = mp.Process(target=flush_process)
|
|
flushing_p.start()
|
|
|
|
print("Messages flushing process started.")
|
|
|
|
time.sleep(5)
|
|
|
|
except KeyboardInterrupt:
|
|
print(f"Stopping {shard_count} shards...")
|
|
|
|
for i, p in enumerate(ps):
|
|
p.terminate()
|
|
print(f"Terminated shard {i}.")
|
|
|
|
print("Stopping messages flushing process...")
|
|
|
|
flushing_p.terminate()
|
|
|
|
print("Checking that all processes have finished...")
|
|
|
|
for p in ps:
|
|
p.join()
|
|
|
|
flushing_p.join()
|
|
|
|
print("All processes have been terminated.")
|