162 lines
5 KiB
Python
162 lines
5 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from multiprocessing import Process, get_context
|
|
from pathlib import Path
|
|
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 stop_subprocesses(processes: list[subprocess.Popen]) -> None:
|
|
for process in processes:
|
|
if process.poll() is None:
|
|
process.terminate()
|
|
for process in processes:
|
|
try:
|
|
process.wait(SHUTDOWN_TIMEOUT)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.wait()
|
|
|
|
|
|
def run() -> int:
|
|
# Import only after taking the lock so a rejected duplicate has no bot/event
|
|
# side effects. Linux workers are fresh interpreters which explicitly inherit
|
|
# the lock descriptor. If the supervisor dies, an orphan keeps the lock and a
|
|
# replacement cannot poll alongside it.
|
|
instance_lock = acquire_instance_lock()
|
|
import models.logger
|
|
|
|
shutdown_requested = False
|
|
|
|
def request_shutdown(_signum, _frame) -> None:
|
|
nonlocal shutdown_requested
|
|
shutdown_requested = True
|
|
|
|
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...")
|
|
|
|
if sys.platform == "win32":
|
|
import bot
|
|
import event
|
|
|
|
context = get_context("spawn")
|
|
processes: list[Process] = []
|
|
definitions = [("bot", bot.run), ("events", event.run)]
|
|
try:
|
|
for name, target in definitions:
|
|
process = context.Process(target=target, name=name)
|
|
process.start()
|
|
processes.append(process)
|
|
|
|
while not shutdown_requested and all(process.is_alive() for process in processes):
|
|
time.sleep(0.5)
|
|
|
|
if not shutdown_requested:
|
|
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
|
|
return 0
|
|
except Exception as error:
|
|
models.logger.error("Main", f"{error}")
|
|
return 1
|
|
finally:
|
|
stop_processes(processes)
|
|
instance_lock.close()
|
|
|
|
project_dir = Path(__file__).resolve().parent
|
|
definitions = [
|
|
("bot", "import bot; bot.run()"),
|
|
("events", "import event; event.run()"),
|
|
]
|
|
processes: list[tuple[str, subprocess.Popen]] = []
|
|
try:
|
|
for name, command in definitions:
|
|
process = subprocess.Popen(
|
|
[sys.executable, "-c", command],
|
|
cwd=project_dir,
|
|
pass_fds=(instance_lock.fileno(),),
|
|
)
|
|
processes.append((name, process))
|
|
|
|
while not shutdown_requested and all(process.poll() is None for _, process in processes):
|
|
time.sleep(0.5)
|
|
|
|
if not shutdown_requested:
|
|
name, stopped = next(
|
|
(name, process) for name, process in processes if process.poll() is not None
|
|
)
|
|
models.logger.error(
|
|
"Main",
|
|
f"{name} process exited unexpectedly with code {stopped.returncode}",
|
|
)
|
|
return 1
|
|
return 0
|
|
except Exception as error:
|
|
models.logger.error("Main", f"{error}")
|
|
return 1
|
|
finally:
|
|
stop_subprocesses([process for _, process in 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)
|