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
+2 -2
View File
@@ -64,7 +64,7 @@ class Payphone(commands.AutoShardedBot):
shard_id=shard_id,
)
print(f"Shard is ready! Shard ID: {shard_id}")
print(f"Shard {shard_id} is ready!")
async def on_shard_resumed(self, shard_id: int) -> None:
"""Sets the activity for the shard.
@@ -81,7 +81,7 @@ class Payphone(commands.AutoShardedBot):
shard_id=shard_id,
)
print(f"Shard was resumed! Shard ID: {shard_id}")
print(f"Shard {shard_id} was resumed!")
async def get_or_fetch_channel(
self, channel_id: int
-8
View File
@@ -1,8 +1,5 @@
from datetime import datetime, timezone
import tortoise
from payphone.db import models
from payphone.core import config
@@ -10,8 +7,3 @@ async def init() -> None:
"""Initialize the database and it's connections."""
await tortoise.Tortoise.init(config.TORTOISE_ORM, use_tz=True)
await tortoise.Tortoise.generate_schemas(safe=True)
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)
+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.")