fix(agent-updates): allow bounded slow artifact downloads

services/agent-updates/publisher/release_helpers/downloads.py

@@ -19,7 +19,7 @@
 import zipfile
 
 from .bundle import BundleVerifier
-from .context import need, sha, decode
+from .context import OperatorError, need, sha, decode
 
 
 SCHEMA_BINDINGS = "agent-release-bindings-v1"
@@ -36,6 +36,7 @@
 MAX_ZIP = MAX_FILE = 512 * 1024**2
 MAX_TOTAL = 1024**3
 CHUNK = 1024**2
+TRANSFER_TIMEOUT_SECONDS = 900
 HEX = re.compile(r"[0-9a-f]{64}\Z")
 
 
@@ -90,11 +91,16 @@ def get(self, path):
 
     def download(self, artifact_id, output, maximum, timeout):
         need(type(artifact_id) is int and 0 < artifact_id < 2**64, "invalid artifact ID")
-        result = subprocess.run(
-            [str(self.ctx.tool("gh")), "api", "--hostname", "github.com", "--method", "GET",
-             f"repos/{self.ctx.core.REPOSITORY}/actions/artifacts/{artifact_id}/zip"],
-            stdin=subprocess.DEVNULL, stdout=output, stderr=subprocess.DEVNULL,
-            timeout=timeout, check=False, preexec_fn=cap(maximum))
+        try:
+            result = subprocess.run(
+                [str(self.ctx.tool("gh")), "api", "--hostname", "github.com", "--method", "GET",
+                 f"repos/{self.ctx.core.REPOSITORY}/actions/artifacts/{artifact_id}/zip"],
+                stdin=subprocess.DEVNULL, stdout=output, stderr=subprocess.DEVNULL,
+                timeout=timeout, check=False, preexec_fn=cap(maximum))
+        except subprocess.TimeoutExpired:
+            raise OperatorError(
+                f"GitHub GET exceeded the {timeout}-second transfer timeout; partial archive retained"
+            ) from None
         need(result.returncode == 0, "GitHub GET failed or exceeded its bound")
 
 
@@ -218,7 +224,7 @@ def run(self, args):
              and not directory.exists() and not directory.is_symlink(), "producer destination already exists")
         fd = os.open(archive, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
         with os.fdopen(fd, "wb") as output:
-            self.download(artifact_id, output, size + 1, 180)
+            self.download(artifact_id, output, size + 1, TRANSFER_TIMEOUT_SECONDS)
             output.flush()
             os.fsync(output.fileno())
         actual, _ = read_record(archive, MAX_ZIP)
@@ -455,7 +461,7 @@ def run(self, args):
         archive = output / "artifact.zip"
         fd = os.open(archive, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
         with os.fdopen(fd, "wb") as stream:
-            self.download(artifact["id"], stream, args.zip_size + 1, 900)
+            self.download(artifact["id"], stream, args.zip_size + 1, TRANSFER_TIMEOUT_SECONDS)
             stream.flush()
             os.fsync(stream.fileno())
         actual, _ = read_record(archive, MAX_ZIP)

services/agent-updates/publisher/test_download_helpers.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 import stat
 import struct
+import subprocess
 import tempfile
 from types import SimpleNamespace
 import unittest
@@ -318,6 +319,57 @@ def test_producer_download_binds_exact_artifact_zip_and_tooling(self):
         self.assertEqual(receipt["bindings_sha256"], self.ctx.bindings_sha256)
         self.assertTrue((self.root / "producer-linux-download.json").is_file())
 
+    def test_slow_producer_transfer_can_finish_before_the_shared_bounded_timeout(self):
+        zip_bytes, files, responses = self.producer_remote()
+        self.ctx.gh_get = Mock(side_effect=responses)
+
+        def transfer(command, **options):
+            simulated_seconds = 300
+            if simulated_seconds > options["timeout"]:
+                raise subprocess.TimeoutExpired(command, options["timeout"])
+            options["stdout"].write(zip_bytes)
+            return subprocess.CompletedProcess(command, 0)
+
+        with patch.object(d.subprocess, "run", side_effect=transfer) as request:
+            receipt = d.ProducerDownloader(self.ctx).run(SimpleNamespace(network=True, platform="linux"))
+        request.assert_called_once()
+        self.assertLessEqual(request.call_args.kwargs["timeout"], 900)
+        self.assertEqual(receipt["files"], self.records(files))
+        self.assertEqual(receipt["actual_zip_sha256"], d.sha(zip_bytes))
+        self.assertTrue((self.root / "producer-linux-download.json").is_file())
+
+    def test_producer_timeout_retains_partial_without_retry_extraction_or_success_receipt(self):
+        _, _, responses = self.producer_remote()
+        self.ctx.gh_get = Mock(side_effect=responses)
+        partial = b"incomplete archive"
+
+        def transfer(command, **options):
+            simulated_seconds = 901
+            if simulated_seconds > options["timeout"]:
+                options["stdout"].write(partial)
+                raise subprocess.TimeoutExpired(command, options["timeout"])
+            return subprocess.CompletedProcess(command, 0)
+
+        downloader = d.ProducerDownloader(self.ctx)
+        args = SimpleNamespace(network=True, platform="linux")
+        selected = self.ctx.bindings["artifacts"]["linux"]
+        archive = self.root / f"producer-linux-{selected['id']}.zip"
+        with patch.object(d.subprocess, "run", side_effect=transfer) as request, \
+             patch.object(downloader, "extract") as extract:
+            with self.assertRaisesRegex(OperatorError, "transfer timeout; partial archive retained"):
+                downloader.run(args)
+            self.assertEqual(archive.read_bytes(), partial)
+            self.assertEqual(stat.S_IMODE(archive.stat().st_mode), 0o600)
+            self.assertFalse((self.root / f"packages-linux-{selected['id']}").exists())
+            self.assertFalse((self.root / "producer-linux-download.json").exists())
+            self.ctx.gh_get = Mock(side_effect=responses[:2])
+            with self.assertRaisesRegex(OperatorError, "producer destination already exists"):
+                downloader.run(args)
+            request.assert_called_once()
+            extract.assert_not_called()
+            self.assertLessEqual(request.call_args.kwargs["timeout"], 900)
+        self.assertEqual(archive.read_bytes(), partial)
+
     def test_producer_attempt_change_during_download_fails_without_receipt(self):
         zip_bytes, _, responses = self.producer_remote()
         responses[-1]["run_attempt"] = 3