payphone/run.py

from multiprocessing import managers

from datetime import datetime, timezone


def shard(
    shard_id: int,
    shard_count: int,
    shutdown: managers.ValueProxy[bool],
    scheduled_shutdown: managers.ValueProxy[bool],
) -> None:
    """Starts a new shard.

    Args:
        shard_id (int): The shard's ID.
        shard_count (int): The total amount of shards.
        shutdown (ValueProxy[bool]): The flag for the main process to shut down all shards.
        scheduled_shutdown (ValueProxy[bool]): The flag for other shards to know the bot's gonna
          shut down soon.
    """

    from payphone.bot import bot
    from payphone.core import config

    bot.shard_ids = [shard_id]
    bot.shard_count = shard_count
    bot.shutdown = shutdown
    bot.scheduled_shutdown = scheduled_shutdown

    print(f"Running shard {shard_id} with default loop.")
    bot.run(config.TOKEN)


def flush_process():
    import asyncio

    from procs.flush_messages import loop

    try:
        import uvloop  # pylint: disable=import-error

        uvloop.install()
        print("Flushing messages with uvloop.")

    except ImportError:
        print("Flushing messages with default loop.")

    finally:
        asyncio.run(loop())


class ShutdownException(Exception):
    """Raised when a shard has triggered the shutdown of all shards."""


if __name__ == "__main__":
    import multiprocessing as mp

    import time

    from tortoise import run_async

    import httpx

    import redis.asyncio as aredis

    from payphone.db import setup, models
    from payphone.core import config

    redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "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",
        headers={"Authorization": f"Bot {config.TOKEN}"},
    )

    info = response.json()

    shard_count = info["shards"]

    manager = mp.Manager()
    shutdown = manager.Value("b", False)
    scheduled_shutdown = manager.Value("b", False)

    print(f"Starting {shard_count} shard{"s" if shard_count != 1 else ""}.")

    ps: list[mp.Process] = [
        mp.Process(
            target=shard,
            args=(i, shard_count, shutdown, scheduled_shutdown),
        )
        for i in range(shard_count)
    ]
    flushing_p = mp.Process(target=flush_process)

    for i, p in enumerate(ps):
        p.start()

        print(
            f"Shard {i}'s process started."
            + (
                " Waiting 5 seconds start initialize next shard..."
                if i != len(ps) - 1
                else " All shards were started."
            )
        )

        if i != len(ps) - 1:
            time.sleep(5)

    print("Starting messages flushing process...")

    flushing_p.start()

    print("Messages flushing process started.")

    try:
        while True:
            if shutdown.value:
                raise ShutdownException()

            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,
                            shutdown,
                            scheduled_shutdown,
                        ),
                        name=f"payphone-shard-{i}",
                    )
                    ps[i].start()

                    print(
                        f"Shard {i} was restarted. Waiting 5 seconds to check next shard..."
                    )
                    time.sleep(5)

            if not flushing_p.is_alive():
                print("Messages flushing process was stopped. Restarting it...")

                flushing_p = mp.Process(target=flush_process, name="payphone-flushing")
                flushing_p.start()

                print("Messages flushing process started.")

            time.sleep(5)

    except KeyboardInterrupt:
        print("Ctrl-C was pressed.")

    except ShutdownException:
        print("A shard shut down the bot.")

    finally:
        print(f"Terminating {shard_count} shards...")

        for i, p in enumerate(ps):
            p.terminate()
            print(f"Terminated shard {i}.")

        print("Stopping messages flushing process...")

        flushing_p.terminate()

        print("Checking that all processes have finished...")

        for p in ps:
            p.join()

        flushing_p.join()

        print("All processes have been terminated.")