131 lines
4 KiB
Python
131 lines
4 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import sys
|
|
from multiprocessing import Process, get_context
|
|
from pathlib import Path
|
|
from collections.abc import Callable
|
|
from typing import BinaryIO
|
|
|
|
|
|
LOCK_PATH = Path(__file__).resolve().with_name(".bot.lock")
|
|
SHUTDOWN_TIMEOUT = 10.0
|
|
|
|
|
|
def acquire_instance_lock() -> BinaryIO:
|
|
"""Hold a process lock for this checkout until the last worker exits."""
|
|
lock_file = LOCK_PATH.open("a+b")
|
|
try:
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
lock_file.seek(0)
|
|
if lock_file.read(1) == b"":
|
|
lock_file.write(b"\0")
|
|
lock_file.flush()
|
|
lock_file.seek(0)
|
|
msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
|
|
else:
|
|
import fcntl
|
|
|
|
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except (BlockingIOError, OSError):
|
|
lock_file.close()
|
|
raise RuntimeError(f"another bot instance already holds {LOCK_PATH}") from None
|
|
|
|
lock_file.seek(0)
|
|
lock_file.truncate()
|
|
lock_file.write(f"{os.getpid()}\n".encode())
|
|
lock_file.flush()
|
|
return lock_file
|
|
|
|
|
|
def stop_processes(processes: list[Process]) -> None:
|
|
started = [process for process in processes if process.pid is not None]
|
|
for process in started:
|
|
if process.is_alive():
|
|
process.terminate()
|
|
for process in started:
|
|
process.join(SHUTDOWN_TIMEOUT)
|
|
for process in started:
|
|
if process.is_alive():
|
|
process.kill()
|
|
process.join()
|
|
|
|
|
|
def run_worker(target: Callable[[], None], lock_fd: int) -> None:
|
|
# Keep an explicit duplicate alive in the worker. A Linux flock belongs to
|
|
# the shared open-file description, so it remains held if the parent dies.
|
|
with os.fdopen(os.dup(lock_fd), "rb"):
|
|
target()
|
|
|
|
|
|
def run() -> int:
|
|
# Import only after taking the lock so a rejected duplicate has no bot/event
|
|
# side effects. Linux workers explicitly retain the lock descriptor: if the
|
|
# supervisor dies unexpectedly, an orphan cannot poll alongside a replacement.
|
|
instance_lock = acquire_instance_lock()
|
|
import bot
|
|
import event
|
|
import models.logger
|
|
|
|
context = get_context("spawn" if sys.platform == "win32" else "fork")
|
|
if sys.platform == "win32":
|
|
processes = [
|
|
context.Process(target=bot.run, name="bot"),
|
|
context.Process(target=event.run, name="events"),
|
|
]
|
|
else:
|
|
processes = [
|
|
context.Process(
|
|
target=run_worker,
|
|
args=(bot.run, instance_lock.fileno()),
|
|
name="bot",
|
|
),
|
|
context.Process(
|
|
target=run_worker,
|
|
args=(event.run, instance_lock.fileno()),
|
|
name="events",
|
|
),
|
|
]
|
|
|
|
def request_shutdown(_signum, _frame) -> None:
|
|
raise KeyboardInterrupt
|
|
|
|
signal.signal(signal.SIGINT, request_shutdown)
|
|
signal.signal(signal.SIGTERM, request_shutdown)
|
|
|
|
os.system("cls" if sys.platform == "win32" else "clear")
|
|
models.logger.log("Main", "Startup...")
|
|
try:
|
|
for process in processes:
|
|
process.start()
|
|
|
|
while all(process.is_alive() for process in processes):
|
|
for process in processes:
|
|
process.join(timeout=0.5)
|
|
|
|
stopped = next(process for process in processes if not process.is_alive())
|
|
models.logger.error(
|
|
"Main",
|
|
f"{stopped.name} process exited unexpectedly with code {stopped.exitcode}",
|
|
)
|
|
return 1
|
|
except KeyboardInterrupt:
|
|
models.logger.error("Main", "Stopping...")
|
|
return 0
|
|
except Exception as error:
|
|
models.logger.error("Main", f"{error}")
|
|
return 1
|
|
finally:
|
|
stop_processes(processes)
|
|
instance_lock.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(run())
|
|
except RuntimeError as error:
|
|
print(f"Bot startup refused: {error}", file=sys.stderr)
|
|
raise SystemExit(1)
|