payphone/run.py

from multiprocessing import managers

from datetime import datetime, timezone

from payphone.core import logging


logging.setup()

log = logging.logger


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

    log.info("Running shard %s with default loop.", shard_id)

    bot.loop.set_debug(True)
    bot.run(config.TOKEN)


def flush_process():
    import asyncio

    from procs.flush_messages import loop

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

        uvloop.install()
        log.info("Flushing messages with uvloop.")

    except ImportError:
        log.info("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()

    log.info("Initializing DB and running Redis cleanup...")

    run_async(init())

    log.info("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)

    log.info("Starting %s shards.", shard_count)

    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()

        log.info("Shard %s's process started.", i)
        log.info(
            "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)

    log.info("Starting messages flushing process...")

    flushing_p.start()

    log.info("Messages flushing process started.")

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

            for i in range(shard_count):
                if not ps[i].is_alive():
                    log.warning("Shard %s's process exited. Restarting the shard...", i)

                    ps[i] = mp.Process(
                        target=shard,
                        args=(
                            i,
                            shard_count,
                            shutdown,
                            scheduled_shutdown,
                        ),
                        name=f"payphone-shard-{i}",
                    )
                    ps[i].start()

                    log.info(
                        "Shard %s was restarted. Waiting 5 seconds to check next shard...",
                        i,
                    )
                    time.sleep(5)

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

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

                log.info("Messages flushing process started.")

            time.sleep(5)

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

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

    finally:
        log.info("Terminating %s shards...", shard_count)

        for i, p in enumerate(ps):
            p.terminate()
            log.info("Terminated shard %s.", i)

        log.info("Stopping messages flushing process...")

        flushing_p.terminate()

        log.info("Checking that all processes have finished...")

        for p in ps:
            p.join()

        flushing_p.join()

        log.info("All processes have been terminated.")