import asyncio
import json
import os
import signal
import sys
import time
from pathlib import Path
import aiohttp
import aioredis
import config
payload = {
"Authorization": f"Bot {config.token}",
"User-Agent": f"DiscordBot (custom, {config.__version__})",
}
async def get_shard_count():
async with aiohttp.ClientSession() as session, session.get(
"https://discordapp.com/api/gateway/bot", headers=payload,
) as req:
response = await req.json()
return response["shards"]
def get_cluster_list(shards):
return [
list(range(0, shards)[i : i + config.shards_per_cluster]) for i in range(0, shards, config.shards_per_cluster)
]
class Instance:
def __init__(self, instance_id, shard_list, shard_count, loop, main, cluster_count):
self.id = instance_id
self.shard_list = shard_list
self.shard_count = shard_count
self.loop = loop
self.main = main
self.cluster_count = cluster_count
self.started_at = None
self.command = (
f"{sys.executable} \"{Path.cwd() / 'main.py'}\" \"{shard_list}\" {shard_count} {self.id} {cluster_count}"
)
self._process = None
self.status = "initialized"
self.started_at = 0.0
self.task = self.loop.create_task(self.start())
self.task.add_done_callback(self.main.dead_process_handler)
@property
def is_active(self):
return self._process is not None and not self._process.returncode
async def read_stream(self, stream):
while self.is_active:
try:
line = await stream.readline()
except (asyncio.LimitOverrunError, ValueError):
continue
if line:
line = line.decode("utf-8")[:-1]
print(f"[Cluster {self.id}] {line}")
else:
break
async def start(self):
if self.is_active:
print(f"[Cluster {self.id}] Already active.")
return
self.started_at = time.time()
self._process = await asyncio.create_subprocess_shell(
self.command,
stdin=asyncio.subprocess.DEVNULL,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
preexec_fn=os.setsid,
limit=1024 * 256,
)
self.status = "running"
self.started_at = time.time()
print(f"[Cluster {self.id}] The cluster is starting.")
await asyncio.wait([self.read_stream(self._process.stdout), self.read_stream(self._process.stderr)])
return self
async def stop(self):
self.status = "stopped"
os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
print(f"[Cluster {self.id}] The cluster is killed.")
await asyncio.sleep(5)
def kill(self):
self.status = "stopped"
os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
async def restart(self):
if self.is_active:
await self.stop()
await self.start()
class Main:
def __init__(self, loop):
self.loop = loop
self.instances = []
self.redis = None
def dead_process_handler(self, result):
instance = result.result()
print(f"[Cluster {instance.id}] The cluster exited with code {instance._process.returncode}.")
if instance._process.returncode == 0 or instance._process.returncode == -15:
print(f"[Cluster {instance.id}] The cluster stopped gracefully.")
else:
print(f"[Cluster {instance.id}] The cluster is restarting.")
instance.loop.create_task(instance.start())
def get_instance(self, iterable, instance_id):
for element in iterable:
if getattr(element, "id") == instance_id:
return element
return None
async def event_handler(self):
self.redis = await aioredis.create_pool("redis://localhost", minsize=1, maxsize=2)
await self.redis.execute_pubsub("SUBSCRIBE", config.ipc_channel)
channel = self.redis.pubsub_channels[bytes(config.ipc_channel, "utf-8")]
while await channel.wait_message():
payload = await channel.get_json(encoding="utf-8")
if payload.get("scope") != "launcher" or not payload.get("action"):
pass
elif payload.get("action") == "restart":
print(f"[Cluster Manager] Received signal to restart cluster {payload.get('id')}.")
self.loop.create_task(self.get_instance(self.instances, payload.get("id")).restart())
elif payload.get("action") == "stop":
print(f"[Cluster Manager] Received signal to stop cluster {payload.get('id')}.")
self.loop.create_task(self.get_instance(self.instances, payload.get("id")).stop())
elif payload.get("action") == "start":
print(f"[Cluster Manager] Received signal to start cluster {payload.get('id')}.")
self.loop.create_task(self.get_instance(self.instances, payload.get("id")).start())
elif payload.get("action") == "statuses" and payload.get("command_id"):
statuses = {}
for instance in self.instances:
statuses[str(instance.id)] = {
"active": instance.is_active,
"status": instance.status,
"started_at": instance.started_at,
}
await self.redis.execute(
"PUBLISH",
config.ipc_channel,
json.dumps({"command_id": payload["command_id"], "output": statuses}),
)
elif payload.get("action") == "roll_restart":
print("[Cluster Manager] Received signal to perform a rolling restart.")
for instance in self.instances:
self.loop.create_task(instance.restart())
await asyncio.sleep(config.shards_per_cluster * 10)
async def close(self):
await self.redis.execute_pubsub("UNSUBSCRIBE", config.ipc_channel)
self.redis.close()
def write_targets(self, clusters):
data = []
for i, shard_list in enumerate(clusters, 1):
if not shard_list:
continue
data.append({"labels": {"cluster": f"{i}"}, "targets": [f"localhost:{6000 + i}"]})
with open("targets.json", "w") as f:
json.dump(data, f, indent=4)
async def launch(self):
self.loop.create_task(self.event_handler())
shard_count = await get_shard_count() + config.additional_shards
clusters = get_cluster_list(shard_count)
if config.testing is False:
self.write_targets(clusters)
print(f"[Cluster Manager] Starting a total of {len(clusters)} clusters.")
for i, shard_list in enumerate(clusters, 1):
if not shard_list:
continue
self.instances.append(
Instance(i, shard_list, shard_count, self.loop, main=self, cluster_count=len(clusters))
)
await asyncio.sleep(config.shards_per_cluster * 10)
loop = asyncio.get_event_loop()
main = Main(loop=loop)
loop.create_task(main.launch())
try:
loop.run_forever()
except KeyboardInterrupt:
def shutdown_handler(_loop, context):
if "exception" not in context or not isinstance(context["exception"], asyncio.CancelledError):
_loop.default_exception_handler(context)
loop.set_exception_handler(shutdown_handler)
for instance in main.instances:
instance.task.remove_done_callback(main.dead_process_handler)
instance.kill()
loop.run_until_complete(main.close())
tasks = asyncio.gather(*asyncio.all_tasks(loop=loop), return_exceptions=True)
tasks.add_done_callback(lambda t: loop.stop())
tasks.cancel()
finally:
loop.run_until_complete(loop.shutdown_asyncgens())
loop.close()
