diff --git a/agent/base.py b/agent/base.py index 574a51b7..9f4749e7 100644 --- a/agent/base.py +++ b/agent/base.py @@ -5,6 +5,7 @@ import shutil import subprocess import tempfile +import time import traceback from contextlib import suppress from datetime import datetime @@ -19,10 +20,30 @@ from agent.utils import get_execution_result if TYPE_CHECKING: - from typing import Any + from typing import Any, Callable from agent.job import Job, Step +PUBLISH_INTERVAL_SECONDS = 2 + + +class Throttle: + """Calls function at most once every `seconds`, dropping calls made in between. + + A class, not a closure, so RQ can pickle its owner. + """ + + def __init__(self, function: Callable, seconds: float): + self.function = function + self.seconds = seconds + self.last_call = -seconds + + def __call__(self, *args): + if time.monotonic() - self.last_call < self.seconds: + return + self.last_call = time.monotonic() + self.function(*args) + class Base: if TYPE_CHECKING: @@ -135,6 +156,8 @@ def parse_output(self, process) -> str: line = b"" lines = [] prev_char = None + # Each publish rewrites the whole output in Redis and its AOF + publish_lines = Throttle(self.publish_lines, PUBLISH_INTERVAL_SECONDS) # This is equivalent of remove_crs # Make sure output matches what'll be shown in the terminal # This won't work for top, htop etc, but good enough to handle progress bars @@ -147,9 +170,9 @@ def parse_output(self, process) -> str: if char == b"\n": lines.append(line.decode(errors="replace")) line = b"" - self.publish_lines(lines) + publish_lines(lines) elif prev_char == b"\r": - self.publish_lines([*lines, line.decode(errors="replace")]) + publish_lines([*lines, line.decode(errors="replace")]) line = b"" line += char else: diff --git a/agent/builder.py b/agent/builder.py index aeea8aeb..beb63728 100644 --- a/agent/builder.py +++ b/agent/builder.py @@ -21,7 +21,7 @@ import requests import semantic_version as sv -from agent.base import Base +from agent.base import PUBLISH_INTERVAL_SECONDS, Base, Throttle from agent.build_utils.validations import check_python_syntax, get_package_manager_files from agent.exceptions import AgentException, RegistryDownException from agent.job import Job, Step, job, step @@ -600,7 +600,7 @@ def __init__( self.no_cache = no_cache self.no_push = no_push - self.last_published = datetime.now() + self.publish_output = Throttle(self.publish_data, PUBLISH_INTERVAL_SECONDS) self.build_failed = False self.build_token = build_token self.secret_path = None @@ -731,8 +731,8 @@ def _get_build_environment(self) -> dict: def _publish_docker_build_output(self, result): for line in result: self.output["build"].append(stamp_step_end(line)) - self._publish_throttled_output(False) - self._publish_throttled_output(True) + self.publish_output(self.output) + self.publish_data(self.output) def _wait_for_registry_recovery(self): """Wait for registry to recover after restart""" @@ -763,13 +763,13 @@ def _push_docker_image(self): except RegistryDownException as e: if attempt == max_retries - 1: - self._publish_throttled_output(True) + self.publish_data(self.output) raise Exception("Failed to push image after multiple attempts") from e self._wait_for_registry_recovery() except Exception: - self._publish_throttled_output(True) + self.publish_data(self.output) raise return None @@ -788,19 +788,7 @@ def _push_image(self, client): auth_config=auth_config, ): self.output["push"].append(line) - self._publish_throttled_output(False) - - def _publish_throttled_output(self, flush: bool): - if flush: - self.publish_data(self.output) - return - - now = datetime.now() - if (now - self.last_published).total_seconds() <= 1: - return - - self.last_published = now - self.publish_data(self.output) + self.publish_output(self.output) def _get_image_name(self): return f"{self.image_repository}:{self.image_tag}" @@ -827,7 +815,7 @@ def _run( input_file.close() return_code = process.wait() - self._publish_throttled_output(True) + self.publish_data(self.output) self.build_failed = return_code != 0 self.data.update({"build_failed": self.build_failed}) @@ -867,7 +855,7 @@ def __init__( self.patch_build_app_instructions = patch_build_app_instructions self.container_name = f"patch-build-{build_name}" self.output: Output = {"build": [], "push": []} - self.last_published = datetime.now() + self.publish_output = Throttle(self.publish_data, PUBLISH_INTERVAL_SECONDS) def _get_image_name(self) -> str: return f"{self.image_repository}:{self.image_tag}" @@ -900,7 +888,7 @@ def _start_base_container(self): def _pull_app_updates(self): for patch_build_app_info in self.patch_build_app_instructions: self._pull_app(patch_build_app_info) - self._publish_throttled_output(True) + self.publish_data(self.output) return self.output["build"] def _pull_app(self, patch_build_app_info: PatchBuildAppInfo): @@ -949,21 +937,9 @@ def _docker_exec(self, command: str, publish: bool = True) -> str: output = result.get("output", "") if isinstance(result, dict) else "" if publish: self.output["build"].append(output) - self._publish_throttled_output(False) + self.publish_output(self.output) return output - def _publish_throttled_output(self, flush: bool) -> None: - if flush: - self.publish_data(self.output) - return - - now = datetime.now() - if (now - self.last_published).total_seconds() <= 1: - return - - self.last_published = now - self.publish_data(self.output) - @step("Commit Image") def _commit_patch_image(self): self.execute( @@ -987,9 +963,9 @@ def _push_patch_image(self): auth_config=auth_config, ): self.output["push"].append(line) - self._publish_throttled_output(False) + self.publish_output(self.output) - self._publish_throttled_output(True) + self.publish_data(self.output) return self.output["push"] def _cleanup_container(self): diff --git a/agent/tests/test_base.py b/agent/tests/test_base.py new file mode 100644 index 00000000..351f3457 --- /dev/null +++ b/agent/tests/test_base.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import io +import pickle +import unittest +from unittest.mock import MagicMock, patch + +from agent.base import PUBLISH_INTERVAL_SECONDS, Base, Throttle + + +class FinishedProcess: + def __init__(self, output: bytes): + self.stdout = io.BytesIO(output) + + def poll(self): + return 0 + + +@patch("agent.base.time.monotonic", return_value=1000.0) +class TestParseOutput(unittest.TestCase): + def setUp(self): + self.base = Base() + self.base.update_redis = MagicMock() + + def test_lines_within_two_seconds_are_published_once_plus_the_final_output(self, _): + output = b"".join(f"line {i}\n".encode() for i in range(1000)) + + self.base.parse_output(FinishedProcess(output)) + + self.assertEqual(self.base.update_redis.call_count, 2) + self.assertEqual(self.base.data["output"], "\n".join(f"line {i}" for i in range(1000))) + + def test_progress_bar_updates_within_two_seconds_are_published_once_plus_the_final_output(self, _): + output = b"".join(f"\r{percent}%".encode() for percent in range(101)) + b"\n" + + self.base.parse_output(FinishedProcess(output)) + + self.assertEqual(self.base.update_redis.call_count, 2) + self.assertEqual(self.base.data["output"], "100%") + + def test_output_is_published_again_after_two_seconds(self, monotonic): + monotonic.side_effect = [1000.0, 1000.0, 1001.0, 1002.0, 1002.0] + + self.base.parse_output(FinishedProcess(b"a\nb\nc\n")) + + # a (first), b skipped at +1s, c at +2s, then the final output + self.assertEqual(self.base.update_redis.call_count, 3) + + +class TestThrottle(unittest.TestCase): + def test_object_holding_a_throttle_can_be_pickled_for_rq(self): + base = Base() + base.publish_output = Throttle(base.publish_data, PUBLISH_INTERVAL_SECONDS) + + restored = pickle.loads(pickle.dumps(base)) + + self.assertEqual(restored.publish_output.seconds, PUBLISH_INTERVAL_SECONDS)