From 00abb50c3e137c64811bb5eb581d276a133845e8 Mon Sep 17 00:00:00 2001
From: "Nyeki.py" <84998222+Nekidev@users.noreply.github.com>
Date: Fri, 25 Oct 2024 20:29:29 -0300
Subject: [PATCH] feat(bot): :sparkles: Add a heartbeat to every shard process.

---
 payphone/bot.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/payphone/bot.py b/payphone/bot.py
index af9658b..0715fba 100644
--- a/payphone/bot.py
+++ b/payphone/bot.py
@@ -8,7 +8,7 @@ import sys
 
 import discord
 from discord import abc
-from discord.ext import commands
+from discord.ext import commands, tasks
 from discord.ext.commands import errors
 
 from tortoise import expressions, transactions, exceptions as tortoise_exc
@@ -42,6 +42,14 @@ class Payphone(commands.AutoShardedBot):
     `shutdown`.
     """
 
+    heartbeat: managers.ValueProxy[int]
+    """A value proxy that connects with the main process.
+
+    It's a timestamp, updated every second and checked by the main process every 5 seconds. If the
+    timestamp hasn't been updated for the past 10 seconds, it means that the process was stuck, so
+    the shard's process is restarted.
+    """
+
     redis: aredis.Redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")
 
     def __init__(self) -> None:
@@ -72,6 +80,8 @@ class Payphone(commands.AutoShardedBot):
 
         await self.redis.ping()
 
+        self.update_heartbeat.start()
+
         print(f"Ready! Logged in as {self.user.display_name}.")
 
     async def on_shard_ready(self, shard_id: int) -> None:
@@ -441,6 +451,12 @@ class Payphone(commands.AutoShardedBot):
         return discord.EmbedAuthor(
             name=self.user.display_name, icon_url=self.user.display_avatar.url
         )
+    
+    @tasks.loop(seconds=1)
+    async def update_heartbeat(self) -> None:
+        """Updates the heartbeat timestamp."""
+
+        self.heartbeat.value = int(datetime.now(timezone.utc).timestamp())
 
 
 bot = Payphone()