From 29188e53c16c31521b0d34601cc8f3bc4dcbf369 Mon Sep 17 00:00:00 2001
From: "Nyeki.py" <84998222+Nekidev@users.noreply.github.com>
Date: Tue, 22 Oct 2024 00:42:25 -0300
Subject: [PATCH] perf(bot): :zap: Use one process per shard.

Start one process per shard and initialize a new one if any of them fails.
---
 compose.yml    |  2 +-
 pyproject.toml |  1 +
 run.py         | 52 +++++++++++++++++++++++++++++++++++++++++++++++---
 uv.lock        |  2 ++
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/compose.yml b/compose.yml
index 4199b32..7d7e8f8 100644
--- a/compose.yml
+++ b/compose.yml
@@ -6,7 +6,7 @@ services:
     env_file:
       - .env
     environment:
-      - IS_DEBUG=false
+      IS_DEBUG: "false"
     depends_on:
       - database
       - redis
diff --git a/pyproject.toml b/pyproject.toml
index eb4d8d5..d96af10 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,6 +7,7 @@ requires-python = ">=3.12"
 dependencies = [
     "aerich==0.7.1",
     "aiohttp>=3.10.10",
+    "click>=8.1.7",
     "colorthief>=0.2.1",
     "httpx>=0.27.2",
     "orjson>=3.10.7",
diff --git a/run.py b/run.py
index aaa2c71..79cea41 100644
--- a/run.py
+++ b/run.py
@@ -1,5 +1,51 @@
-from payphone import bot
-from payphone.core import config
+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)
 
 
-bot.bot.run(config.TOKEN)
+if __name__ == "__main__":
+    import multiprocessing as mp
+
+    import time
+
+    import httpx
+
+    from payphone.core import config
+
+    response = httpx.get(
+        "https://discord.com/api/v10/gateway/bot",
+        headers={"Authorization": f"Bot {config.TOKEN}"},
+    )
+
+    info = response.json()
+
+    shard_count = info["shards"]
+
+    ps: list[mp.Process] = [
+        mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
+    ]
+
+    for p in ps:
+        p.start()
+
+        time.sleep(5)
+
+    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()
+
+        time.sleep(5)
diff --git a/uv.lock b/uv.lock
index bb681e7..2c2df91 100644
--- a/uv.lock
+++ b/uv.lock
@@ -371,6 +371,7 @@ source = { virtual = "." }
 dependencies = [
     { name = "aerich" },
     { name = "aiohttp" },
+    { name = "click" },
     { name = "colorthief" },
     { name = "httpx" },
     { name = "orjson" },
@@ -384,6 +385,7 @@ dependencies = [
 requires-dist = [
     { name = "aerich", specifier = "==0.7.1" },
     { name = "aiohttp", specifier = ">=3.10.10" },
+    { name = "click", specifier = ">=8.1.7" },
     { name = "colorthief", specifier = ">=0.2.1" },
     { name = "httpx", specifier = ">=0.27.2" },
     { name = "orjson", specifier = ">=3.10.7" },