diff --git a/archivebox/cli/archivebox_add.py b/archivebox/cli/archivebox_add.py index 44ed1ebd..596cd497 100644 --- a/archivebox/cli/archivebox_add.py +++ b/archivebox/cli/archivebox_add.py @@ -244,11 +244,22 @@ def add( assert command is not None exit_code = 0 + + def crawl_is_complete() -> bool: + crawl.refresh_from_db(fields=["status"]) + return crawl.status not in crawl.RUNNABLE_STATES + try: try: with foreground_shutdown_signals(first_signal_message=None), foreground_parent_watchdog(): while True: - standby_until_foreground_runner_needed(command, data_dir=CONSTANTS.DATA_DIR) + standby = standby_until_foreground_runner_needed( + command, + data_dir=CONSTANTS.DATA_DIR, + work_is_complete=crawl_is_complete, + ) + if standby["work_completed"]: + break exit_code = run_runner_worker( ["--crawl-id", str(crawl.id)], name=f"worker_runner_add_{os.getpid()}", diff --git a/archivebox/core/takeover_util.py b/archivebox/core/takeover_util.py index d2240cf7..fc456e40 100644 --- a/archivebox/core/takeover_util.py +++ b/archivebox/core/takeover_util.py @@ -2,6 +2,7 @@ from __future__ import annotations import time import sys +from collections.abc import Callable from pathlib import Path from django.db import IntegrityError @@ -323,12 +324,22 @@ def standby_until_runtime_stack_needed(command, *, data_dir: str | Path, interva return {"resumed": announced, "previous_owner_pid": previous_owner_pid} -def standby_until_foreground_runner_needed(command, *, data_dir: str | Path, interval: float = 2.0) -> dict[str, object]: +def standby_until_foreground_runner_needed( + command, + *, + data_dir: str | Path, + interval: float = 2.0, + work_is_complete: Callable[[], bool] | None = None, +) -> dict[str, object]: from archivebox.workers.supervisord_util import reap_foreground_supervisord_process announced = False previous_owner_pid = None - while not command_owns_foreground_runner(command, data_dir=data_dir): + while True: + if work_is_complete is not None and work_is_complete(): + return {"resumed": announced, "previous_owner_pid": previous_owner_pid, "work_completed": True} + if command_owns_foreground_runner(command, data_dir=data_dir): + break reap_foreground_supervisord_process() if not announced: owner = foreground_runner_owner(data_dir=data_dir) @@ -343,4 +354,4 @@ def standby_until_foreground_runner_needed(command, *, data_dir: str | Path, int time.sleep(interval) command.modified_at = timezone.now() command.save(update_fields=["modified_at"]) - return {"resumed": announced, "previous_owner_pid": previous_owner_pid} + return {"resumed": announced, "previous_owner_pid": previous_owner_pid, "work_completed": False} diff --git a/archivebox/tests/test_takeover_util.py b/archivebox/tests/test_takeover_util.py index 52104591..951b61a0 100644 --- a/archivebox/tests/test_takeover_util.py +++ b/archivebox/tests/test_takeover_util.py @@ -217,8 +217,9 @@ def test_behavior_foreground_add_keeps_existing_server_http_visible(tmp_path, in add = run_archivebox_cmd( [ "add", - "--depth=0", - "--plugins=wget", + "--depth=2", + "--max-urls=10", + "--plugins=wget,parse_html_urls", recursive_test_site["root_url"], ], cwd=tmp_path,