diff --git a/archivebox/cli/__init__.py b/archivebox/cli/__init__.py index 40316c02..036c049f 100644 --- a/archivebox/cli/__init__.py +++ b/archivebox/cli/__init__.py @@ -5,10 +5,12 @@ import sys from importlib import import_module import rich_click as click -from rich import print +from rich.console import Console from archivebox.config.version import VERSION +STDERR = Console(stderr=True) + if "--debug" in sys.argv: os.environ["DEBUG"] = "True" @@ -104,9 +106,8 @@ class ArchiveBoxGroup(click.Group): # handle renamed commands if cmd_name in self.renamed_commands: new_name = self.renamed_commands[cmd_name] - print( + STDERR.print( f" [violet]Hint:[/violet] `archivebox {cmd_name}` has been renamed to `archivebox {new_name}`", - file=sys.stderr, ) cmd_name = new_name ctx.invoked_subcommand = cmd_name @@ -175,7 +176,7 @@ def cli(ctx, help=False): if subcommand != "update": check_migrations(auto_apply=True) except Exception as e: - print(f"[red][X] Error setting up Django or checking data folder: {e}[/red]", file=sys.stderr) + STDERR.print(f"[red][X] Error setting up Django or checking data folder: {e}[/red]") if subcommand not in ("manage", "shell"): # not all management commands need django to be setup beforehand raise @@ -189,10 +190,29 @@ def main(args=None, prog_name=None, stdin=None): # stdin param allows passing input data from caller (used by __main__.py) # currently not used by click-based CLI, but kept for backwards compatibility + previous_unraisablehook = sys.unraisablehook + + def ignore_shutdown_unraisable(unraisable): + if isinstance(unraisable.exc_value, (KeyboardInterrupt, SystemExit)): + return + previous_unraisablehook(unraisable) + + sys.unraisablehook = ignore_shutdown_unraisable try: - cli(args=args, prog_name=prog_name) + cli(args=args, prog_name=prog_name, standalone_mode=False) + except click.Abort: + STDERR.print("\n[red][X] Got CTRL+C. Exiting...[/red]") + raise SystemExit(130) from None + except click.ClickException as err: + err.show() + raise SystemExit(err.exit_code) from None + except click.exceptions.Exit as err: + raise SystemExit(err.exit_code) from None except KeyboardInterrupt: - print("\n\n[red][X] Got CTRL+C. Exiting...[/red]") + STDERR.print("\n[red][X] Got CTRL+C. Exiting...[/red]") + raise SystemExit(130) from None + finally: + sys.unraisablehook = previous_unraisablehook if __name__ == "__main__": diff --git a/archivebox/config/django.py b/archivebox/config/django.py index 4841edbb..12354e0d 100644 --- a/archivebox/config/django.py +++ b/archivebox/config/django.py @@ -162,6 +162,6 @@ def setup_django(check_db=False, in_memory_db=False) -> None: # logfire.info(f'Started ArchiveBox v{CONSTANTS.VERSION}', argv=sys.argv) except KeyboardInterrupt: - raise SystemExit(2) + raise DJANGO_SET_UP = True diff --git a/archivebox/core/models.py b/archivebox/core/models.py index 3c2fefcc..76b640d6 100755 --- a/archivebox/core/models.py +++ b/archivebox/core/models.py @@ -27,7 +27,7 @@ from django.core.exceptions import ObjectDoesNotExist, ValidationError from django.utils.safestring import mark_safe from archivebox.config import CONSTANTS -from archivebox.config.common import get_config +from archivebox.config.common import get_config, rprint from archivebox.misc.system import atomic_write from archivebox.misc.util import ( MAX_URL_LENGTH, @@ -2207,7 +2207,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW deleted_count = empty_ars.count() if deleted_count > 0: empty_ars.delete() - print(f"[yellow]🗑️ Deleted {deleted_count} empty ArchiveResults for {self.url}[/yellow]") + rprint(f"[yellow]🗑️ Deleted {deleted_count} empty ArchiveResults for {self.url}[/yellow]") def to_json(self) -> dict: """ @@ -2315,7 +2315,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW record_crawl_id = record.get("crawl_id") if record_crawl_id and crawl and str(crawl.id) != str(record_crawl_id): - print( + rprint( f"[yellow]⚠️ Snapshot.from_json crawl mismatch: record has crawl_id={record_crawl_id}, overrides has crawl={crawl.id}[/yellow]", file=sys.stderr, ) @@ -2341,7 +2341,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW label=f"auto-created for {url[:50]}", created_by_id=created_by_id, ) - print(f"[red]⚠️ Snapshot.from_json auto-created new crawl {crawl.id} for url={url}[/red]", file=sys.stderr) + rprint(f"[red]⚠️ Snapshot.from_json auto-created new crawl {crawl.id} for url={url}[/red]", file=sys.stderr) # Parse tags (accept either a list ["tag1", "tag2"] or a comma-separated string "tag1,tag2") tags_raw = record.get("tags", "") diff --git a/archivebox/crawls/models.py b/archivebox/crawls/models.py index 6af663ee..713c2981 100755 --- a/archivebox/crawls/models.py +++ b/archivebox/crawls/models.py @@ -21,7 +21,7 @@ from django.conf import settings from django.urls import reverse_lazy from django.utils import timezone from statemachine import State, registry -from rich import print +from archivebox.config.common import rprint as print from archivebox.base_models.models import ( ModelWithUUID, diff --git a/archivebox/machine/models.py b/archivebox/machine/models.py index e93470ec..d6e8dce1 100755 --- a/archivebox/machine/models.py +++ b/archivebox/machine/models.py @@ -19,6 +19,7 @@ from django.db.models import Q, QuerySet from django.utils import timezone from django.utils.functional import cached_property +from archivebox.config.common import rprint from archivebox.base_models.models import ModelWithDeleteAfter, ModelWithHealthStats from archivebox.workers.models import BaseStateMachine, ModelWithStateMachine from .detect import get_host_guid, get_os_info, get_vm_info, get_host_network, get_host_stats @@ -2435,10 +2436,10 @@ class Process(ModelWithDeleteAfter, models.Model): if result.returncode == 0: killed += int(result.stdout.strip()) if killed > 0: - print(f"[yellow]🧹 Cleaned up {killed} orphaned Chrome processes[/yellow]") + rprint(f"[yellow]🧹 Cleaned up {killed} orphaned Chrome processes[/yellow]") return killed except (subprocess.TimeoutExpired, ValueError, FileNotFoundError) as e: - print(f"[red]Failed to cleanup orphaned Chrome: {e}[/red]") + rprint(f"[red]Failed to cleanup orphaned Chrome: {e}[/red]") return 0 @@ -2498,7 +2499,7 @@ class Process(ModelWithDeleteAfter, models.Model): cleaned += 1 if cleaned: - print(f"[yellow]🧹 Cleaned up {cleaned} orphaned worker/hook process record(s)[/yellow]") + rprint(f"[yellow]🧹 Cleaned up {cleaned} orphaned worker/hook process record(s)[/yellow]") return cleaned @@ -2553,7 +2554,7 @@ class BinaryMachine(BaseStateMachine): """Called during queued→installed transition. Runs installation synchronously.""" import sys - print(f"[cyan] 🔄 BinaryMachine.on_install() - installing {self.binary.name}[/cyan]", file=sys.stderr) + rprint(f"[cyan] 🔄 BinaryMachine.on_install() - installing {self.binary.name}[/cyan]", file=sys.stderr) # Run installation hooks (synchronous, updates abspath/version/sha256 and sets status) self.binary.run() @@ -2564,7 +2565,7 @@ class BinaryMachine(BaseStateMachine): if self.binary.status != Binary.StatusChoices.INSTALLED: # Installation failed - abort transition, stay in queued - print(f"[red] ❌ BinaryMachine - {self.binary.name} installation failed, retrying later[/red]", file=sys.stderr) + rprint(f"[red] ❌ BinaryMachine - {self.binary.name} installation failed, retrying later[/red]", file=sys.stderr) # Bump retry_at to try again later self.binary.update_and_requeue( @@ -2578,7 +2579,7 @@ class BinaryMachine(BaseStateMachine): # Abort the transition - this will raise an exception and keep us in queued raise Exception(f"Binary {self.binary.name} installation failed") - print(f"[cyan] ✅ BinaryMachine - {self.binary.name} installed successfully[/cyan]", file=sys.stderr) + rprint(f"[cyan] ✅ BinaryMachine - {self.binary.name} installed successfully[/cyan]", file=sys.stderr) @installed.enter def enter_installed(self): diff --git a/archivebox/services/supervision_service.py b/archivebox/services/supervision_service.py index ada118f3..9f2d7e17 100644 --- a/archivebox/services/supervision_service.py +++ b/archivebox/services/supervision_service.py @@ -4,7 +4,7 @@ import time from pathlib import Path from django.utils import timezone -from rich import print +from archivebox.config.common import rprint def runtime_stack_owner_types(): @@ -109,7 +109,7 @@ def ensure_daemon_stack(*, reason: str = ""): return worker if reason: - print(f"[yellow][*] Starting daemon stack for {reason}...[/yellow]") + rprint(f"[yellow][*] Starting daemon stack for {reason}...[/yellow]") return start_worker(supervisor, sonic_worker) @@ -143,7 +143,7 @@ def standby_until_leader_needed(command, *, process_type: str, data_dir: str | P if not announced: leader = newest_live_process(process_type=process_type, data_dir=data_dir, url=url) leader_pid = leader.pid if leader else "unknown" - print(f"[yellow][*] Standing by; newer ArchiveBox parent pid={leader_pid} is running the orchestrator and server.[/yellow]") + rprint(f"[yellow][*] Standing by; newer ArchiveBox parent pid={leader_pid} is running the orchestrator and server.[/yellow]") announced = True command.heartbeat() time.sleep(interval) @@ -161,7 +161,7 @@ def standby_until_runtime_stack_needed(command, *, data_dir: str | Path, interva owner = runtime_stack_owner(data_dir=data_dir) owner_pid = owner.pid if owner else "unknown" owner_type = owner.process_type if owner else "unknown" - print(f"[yellow][*] Standing by; ArchiveBox {owner_type} pid={owner_pid} owns the runtime stack.[/yellow]") + rprint(f"[yellow][*] Standing by; ArchiveBox {owner_type} pid={owner_pid} owns the runtime stack.[/yellow]") announced = True command.heartbeat() time.sleep(interval) diff --git a/archivebox/workers/supervisord_util.py b/archivebox/workers/supervisord_util.py index 8b9429e7..56ee7e7e 100644 --- a/archivebox/workers/supervisord_util.py +++ b/archivebox/workers/supervisord_util.py @@ -13,11 +13,11 @@ from typing import cast from pathlib import Path from functools import cache -from rich import print from supervisor.xmlrpc import SupervisorTransport from xmlrpc.client import Fault, ServerProxy from archivebox.config import CONSTANTS +from archivebox.config.common import rprint as print from archivebox.config.paths import get_or_create_working_tmp_dir from archivebox.config.permissions import ARCHIVEBOX_USER from archivebox.core.shutdown_util import ( diff --git a/bin/fuzz_test.sh b/bin/fuzz_test.sh index 75a9ba96..b14e5c98 100755 --- a/bin/fuzz_test.sh +++ b/bin/fuzz_test.sh @@ -12,6 +12,10 @@ FUZZ_ROUNDS="${FUZZ_ROUNDS:-8}" FUZZ_PARALLEL="${FUZZ_PARALLEL:-5}" FUZZ_KILL_MIN_SECONDS="${FUZZ_KILL_MIN_SECONDS:-60}" FUZZ_KILL_MAX_SECONDS="${FUZZ_KILL_MAX_SECONDS:-120}" +FUZZ_CTRL_C_CHANCE="${FUZZ_CTRL_C_CHANCE:-60}" +FUZZ_CTRL_C_MAX_SIGNALS="${FUZZ_CTRL_C_MAX_SIGNALS:-4}" +FUZZ_CTRL_C_MIN_SECONDS="${FUZZ_CTRL_C_MIN_SECONDS:-2}" +FUZZ_CTRL_C_MAX_SECONDS="${FUZZ_CTRL_C_MAX_SECONDS:-20}" SERVER_BASE_PORT="${FUZZ_SERVER_BASE_PORT:-8700}" SLEEP_BETWEEN_JOBS_MAX="${FUZZ_SLEEP_BETWEEN_JOBS_MAX:-5}" @@ -21,6 +25,12 @@ fi if [[ "$FUZZ_KILL_MAX_SECONDS" -lt "$FUZZ_KILL_MIN_SECONDS" ]]; then FUZZ_KILL_MAX_SECONDS="$FUZZ_KILL_MIN_SECONDS" fi +if [[ "$FUZZ_CTRL_C_MAX_SECONDS" -lt "$FUZZ_CTRL_C_MIN_SECONDS" ]]; then + FUZZ_CTRL_C_MAX_SECONDS="$FUZZ_CTRL_C_MIN_SECONDS" +fi +if [[ "$FUZZ_CTRL_C_MAX_SIGNALS" -lt 1 ]]; then + FUZZ_CTRL_C_MAX_SIGNALS=1 +fi if [[ -n "${ARCHIVEBOX_CMD:-}" ]]; then read -r -a ABX <<< "$ARCHIVEBOX_CMD" @@ -73,15 +83,76 @@ random_start_delay() { random_between 0 "$SLEEP_BETWEEN_JOBS_MAX" } -kill_tree() { - local pid="$1" +random_ctrl_c_delay() { + random_between "$FUZZ_CTRL_C_MIN_SECONDS" "$FUZZ_CTRL_C_MAX_SECONDS" +} + +random_subsecond_delay() { + printf '0.%03d\n' "$((100 + RANDOM % 400))" +} + +signal_tree() { + local signal="$1" + local pid="$2" local child if command -v pgrep >/dev/null 2>&1; then for child in $(pgrep -P "$pid" 2>/dev/null || true); do - kill_tree "$child" + signal_tree "$signal" "$child" done fi - kill "$pid" >/dev/null 2>&1 || true + kill "-$signal" "$pid" >/dev/null 2>&1 || true +} + +is_uv_wrapper_without_child() { + local pid="$1" + local comm + comm="$(ps -p "$pid" -o comm= 2>/dev/null | xargs basename 2>/dev/null || true)" + [[ "$comm" == "uv" ]] && ! pgrep -P "$pid" >/dev/null 2>&1 +} + +kill_tree() { + local pid="$1" + signal_tree TERM "$pid" +} + +start_ctrl_c_injector() { + local label="$1" + local child="$2" + local logfile="$3" + local signals idx delay burst_gap + + if [[ "$FUZZ_CTRL_C_CHANCE" -le 0 || $((RANDOM % 100)) -ge "$FUZZ_CTRL_C_CHANCE" ]]; then + return 0 + fi + + signals=$((1 + RANDOM % FUZZ_CTRL_C_MAX_SIGNALS)) + ( + for idx in $(seq 1 "$signals"); do + if [[ "$idx" -eq 1 || $((RANDOM % 2)) -eq 0 ]]; then + delay="$(random_ctrl_c_delay)" + else + delay="$(random_subsecond_delay)" + fi + sleep "$delay" + if ! kill -0 "$child" >/dev/null 2>&1; then + exit 0 + fi + if is_uv_wrapper_without_child "$child"; then + exit 0 + fi + echo "[$(ts)] CTRL_C label=$label pid=$child signal=$idx/$signals delay=${delay}s" >> "$logfile" + signal_tree INT "$child" + if [[ $((RANDOM % 3)) -eq 0 ]]; then + burst_gap="$(random_subsecond_delay)" + sleep "$burst_gap" + if kill -0 "$child" >/dev/null 2>&1 && ! is_uv_wrapper_without_child "$child"; then + echo "[$(ts)] CTRL_C_BURST label=$label pid=$child gap=${burst_gap}s" >> "$logfile" + signal_tree INT "$child" + fi + fi + done + ) & + echo "$!" } cleanup() { @@ -112,30 +183,39 @@ run_with_timeout() { echo "[$(ts)] START label=$label shell=$$ data=$DATA_DIR" echo "[$(ts)] CMD DATA_DIR=$DATA_DIR $*" echo "[$(ts)] CHAOS kill_after=${timeout}s" + echo "[$(ts)] CTRL_C chance=${FUZZ_CTRL_C_CHANCE}% max_signals=${FUZZ_CTRL_C_MAX_SIGNALS}" } | tee -a "$logfile" ( DATA_DIR="$DATA_DIR" "$@" ) >> "$logfile" 2>&1 & local child=$! + local interrupter + interrupter="$(start_ctrl_c_injector "$label" "$child" "$logfile")" ( sleep "$timeout" if kill -0 "$child" >/dev/null 2>&1; then echo "[$(ts)] TIMEOUT label=$label pid=$child after=${timeout}s" >> "$logfile" - kill "$child" >/dev/null 2>&1 || true + signal_tree TERM "$child" sleep 5 - kill -9 "$child" >/dev/null 2>&1 || true + signal_tree KILL "$child" fi ) & local watchdog=$! - trap '[[ -n "${child:-}" ]] && kill_tree "$child"; [[ -n "${watchdog:-}" ]] && kill "$watchdog" >/dev/null 2>&1 || true' INT TERM + trap '[[ -n "${child:-}" ]] && kill_tree "$child"; [[ -n "${watchdog:-}" ]] && kill "$watchdog" >/dev/null 2>&1 || true; [[ -n "${interrupter:-}" ]] && kill "$interrupter" >/dev/null 2>&1 || true' INT TERM wait "$child" local code=$? kill "$watchdog" >/dev/null 2>&1 || true + if [[ -n "$interrupter" ]]; then + kill "$interrupter" >/dev/null 2>&1 || true + fi wait "$watchdog" >/dev/null 2>&1 || true + if [[ -n "$interrupter" ]]; then + wait "$interrupter" >/dev/null 2>&1 || true + fi trap - INT TERM echo "[$(ts)] END label=$label pid=$child exit=$code log=$logfile" | tee -a "$logfile" @@ -157,22 +237,29 @@ run_server_for_a_bit() { echo "[$(ts)] START label=$label shell=$$ data=$DATA_DIR" echo "[$(ts)] CMD DATA_DIR=$DATA_DIR ${ABX[*]} server $debug_flag 127.0.0.1:$port" echo "[$(ts)] CHAOS kill_after=${hold}s" + echo "[$(ts)] CTRL_C chance=${FUZZ_CTRL_C_CHANCE}% max_signals=${FUZZ_CTRL_C_MAX_SIGNALS}" } | tee -a "$logfile" ( DATA_DIR="$DATA_DIR" "${ABX[@]}" server "${server_extra[@]}" "127.0.0.1:$port" ) >> "$logfile" 2>&1 & local child=$! + local interrupter + interrupter="$(start_ctrl_c_injector "$label" "$child" "$logfile")" - trap '[[ -n "${child:-}" ]] && kill_tree "$child"' INT TERM + trap '[[ -n "${child:-}" ]] && kill_tree "$child"; [[ -n "${interrupter:-}" ]] && kill "$interrupter" >/dev/null 2>&1 || true' INT TERM sleep "$hold" echo "[$(ts)] STOP label=$label pid=$child after=${hold}s" | tee -a "$logfile" kill_tree "$child" sleep 5 - kill -9 "$child" >/dev/null 2>&1 || true + signal_tree KILL "$child" wait "$child" >/dev/null 2>&1 local code=$? + if [[ -n "$interrupter" ]]; then + kill "$interrupter" >/dev/null 2>&1 || true + wait "$interrupter" >/dev/null 2>&1 || true + fi trap - INT TERM echo "[$(ts)] END label=$label pid=$child exit=$code log=$logfile" | tee -a "$logfile" @@ -298,6 +385,7 @@ main() { echo " rounds: $FUZZ_ROUNDS" echo " parallel: $FUZZ_PARALLEL" echo " kill after: ${FUZZ_KILL_MIN_SECONDS}s-${FUZZ_KILL_MAX_SECONDS}s" + echo " ctrl+c: ${FUZZ_CTRL_C_CHANCE}% chance, ${FUZZ_CTRL_C_MAX_SIGNALS} max, ${FUZZ_CTRL_C_MIN_SECONDS}s-${FUZZ_CTRL_C_MAX_SECONDS}s plus bursts" echo " urls: ${URLS[*]}" echo