fix(base): Publish command output at most every two seconds

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:

agent/builder.py

@@ -10,7 +10,6 @@
 import time
 import warnings
 from dataclasses import dataclass, field
-from datetime import datetime
 from pathlib import Path
 from subprocess import Popen
 from typing import TYPE_CHECKING, Dict, List, TypedDict
@@ -19,7 +18,7 @@
 import jinja2
 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
@@ -563,7 +562,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
@@ -657,8 +656,8 @@ def _get_build_environment(self) -> dict:
     def _publish_docker_build_output(self, result):
         for line in result:
             self.output["build"].append(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"""
@@ -689,13 +688,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
@@ -714,19 +713,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}"
@@ -753,7 +740,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})
@@ -793,7 +780,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}"
@@ -826,7 +813,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):
@@ -875,21 +862,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(
@@ -913,9 +888,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):

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)