feat(bot): ✨ Improve the startup script.

This commit is contained in:
Nyeki.py
2024-10-22 18:43:54 -03:00
parent edae71827b
commit c13b351649
3 changed files with 60 additions and 18 deletions
+58 -8
View File
@@ -1,3 +1,6 @@
from datetime import datetime, timezone
def shard(shard_id: int, shard_count: int) -> None:
"""Starts a new shard.
@@ -20,9 +23,29 @@ if __name__ == "__main__":
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",
@@ -39,15 +62,42 @@ if __name__ == "__main__":
mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
]
for p in ps:
for i, p in enumerate(ps):
p.start()
time.sleep(5)
print(
f"Shard {i}'s process started." + (" Waiting 5 seconds start initialize next shard..." if len(ps) - 1 < i else " All shards were started.")
)
while True:
for i in range(shard_count):
if not ps[i].is_alive():
ps[i] = mp.Process(target=shard, args=(i, shard_count))
ps[i].start()
if i + 1 != len(ps):
time.sleep(5)
time.sleep(5)
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)
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("Checking that all processes have finished...")
for p in ps:
p.join()
print("All processes have been terminated.")