bastebin/production.py

#!/usr/bin/env python3
"""
production.py — Manage the Bastebin Gunicorn process.

Usage:
    python production.py start   [--host HOST] [--port PORT] [--workers N]
    python production.py stop
    python production.py restart [--host HOST] [--port PORT] [--workers N]
    python production.py status
"""

import argparse
import multiprocessing
import os
import signal
import subprocess
import sys
import time

BASE_DIR  = os.path.dirname(os.path.abspath(__file__))
PID_FILE  = os.path.join(BASE_DIR, 'gunicorn.pid')
LOG_FILE  = os.path.join(BASE_DIR, 'gunicorn.log')
CONF_FILE = os.path.join(BASE_DIR, 'gunicorn.conf.py')
VENV_BIN  = os.path.join(BASE_DIR, '.venv', 'bin')


def _gunicorn_bin() -> str:
    candidate = os.path.join(VENV_BIN, 'gunicorn')
    if os.path.isfile(candidate):
        return candidate
    import shutil
    found = shutil.which('gunicorn')
    if not found:
        sys.exit('gunicorn not found — run python setup.py first.')
    return found


def _read_pid() -> int | None:
    try:
        with open(PID_FILE) as f:
            return int(f.read().strip())
    except (FileNotFoundError, ValueError):
        return None


def _process_alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
        return True
    except (ProcessLookupError, PermissionError):
        return False


def cmd_status() -> None:
    pid = _read_pid()
    if pid and _process_alive(pid):
        print(f'Running  (PID {pid})')
    else:
        if pid:
            os.remove(PID_FILE)
        print('Stopped')


def cmd_start(host: str, port: int, workers: int) -> None:
    pid = _read_pid()
    if pid and _process_alive(pid):
        print(f'Already running (PID {pid}). Use "restart" to reload.')
        return

    gunicorn = _gunicorn_bin()
    log = open(LOG_FILE, 'a')

    proc = subprocess.Popen(
        [
            gunicorn,
            '--config',  CONF_FILE,
            '--bind',    f'{host}:{port}',
            '--workers', str(workers),
            '--pid',     PID_FILE,
            'wsgi:app',
        ],
        stdout=log,
        stderr=log,
        cwd=BASE_DIR,
        start_new_session=True,
    )

    # Wait briefly to confirm the process stayed alive.
    time.sleep(1.5)
    if proc.poll() is not None:
        log.close()
        sys.exit(f'Gunicorn exited immediately. Check {LOG_FILE} for details.')

    log.close()
    # Gunicorn writes its own PID file; confirm it appeared.
    pid = _read_pid()
    print(f'Started  (PID {pid or proc.pid})  →  http://{host}:{port}')
    print(f'Logs: {LOG_FILE}')


def cmd_stop(graceful: bool = True) -> bool:
    pid = _read_pid()
    if not pid or not _process_alive(pid):
        print('Not running.')
        if pid:
            os.remove(PID_FILE)
        return False

    sig = signal.SIGTERM if graceful else signal.SIGKILL
    os.kill(pid, sig)

    # Wait up to 10 s for a clean shutdown.
    for _ in range(20):
        time.sleep(0.5)
        if not _process_alive(pid):
            break

    if _process_alive(pid):
        # Force-kill if still alive after graceful period.
        os.kill(pid, signal.SIGKILL)
        time.sleep(0.5)

    if os.path.exists(PID_FILE):
        os.remove(PID_FILE)

    print(f'Stopped  (was PID {pid})')
    return True


def cmd_restart(host: str, port: int, workers: int) -> None:
    pid = _read_pid()
    if pid and _process_alive(pid):
        # Send SIGHUP — Gunicorn performs a graceful reload without downtime.
        os.kill(pid, signal.SIGHUP)
        print(f'Reloaded (PID {pid})')
    else:
        print('Not running — starting fresh.')
        cmd_start(host, port, workers)


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog='production.py',
        description='Manage the Bastebin Gunicorn process.',
    )
    sub = parser.add_subparsers(dest='command', metavar='COMMAND')
    sub.required = True

    def _add_server_args(p: argparse.ArgumentParser) -> None:
        p.add_argument('--host',    default='0.0.0.0', metavar='HOST',
                       help='Bind address (default: 0.0.0.0)')
        p.add_argument('--port',    default=5000, type=int, metavar='PORT',
                       help='Bind port (default: 5000)')
        p.add_argument('--workers', default=None, type=int, metavar='N',
                       help='Worker processes (default: cpu_count + 1, min 2)')

    p_start   = sub.add_parser('start',   help='Start Gunicorn in the background')
    p_restart = sub.add_parser('restart', help='Gracefully reload (SIGHUP) or start if stopped')
    sub.add_parser('stop',   help='Stop Gunicorn gracefully')
    sub.add_parser('status', help='Show whether Gunicorn is running')

    _add_server_args(p_start)
    _add_server_args(p_restart)

    return parser


def main() -> None:
    parser = _build_parser()
    args   = parser.parse_args()

    default_workers = max(2, multiprocessing.cpu_count() + 1)
    workers = getattr(args, 'workers', None) or default_workers
    host    = getattr(args, 'host',    '0.0.0.0')
    port    = getattr(args, 'port',    5000)

    if args.command == 'start':
        cmd_start(host, port, workers)
    elif args.command == 'stop':
        cmd_stop()
    elif args.command == 'restart':
        cmd_restart(host, port, workers)
    elif args.command == 'status':
        cmd_status()


if __name__ == '__main__':
    main()