fix(firewall): keep retry budget after a non-retryable miss

Author: cursoragent

dist/main.js

@@ -258,37 +258,39 @@ export async function downloadFirewall({ edition = 'free', ...inputs }) {
  */
 export async function downloadToolWithRetry(urls) {
   let lastError
+  const remaining = urls.slice()
+  let attempt = 0
 
-  for (
-    let attempt = 0;
-    attempt <= DOWNLOAD_RETRY_DELAYS_SECONDS.length;
-    attempt += 1
-  ) {
-    const url = urls[attempt % urls.length]
+  while (remaining.length > 0) {
+    const url = remaining[attempt % remaining.length]
 
     try {
       return await downloadTool(url)
     } catch (error) {
       lastError = error
 
+      if (!isRetryableDownloadError(error)) {
+        remaining.splice(remaining.indexOf(url), 1)
+        if (remaining.length === 0) {
+          break
+        }
+        warning(
+          `Socket Firewall binary download from ${url} failed: ${errorMessage(error)}. Trying ${remaining[attempt % remaining.length]}.`,
+        )
+        continue
+      }
+
       const seconds = DOWNLOAD_RETRY_DELAYS_SECONDS[attempt]
 
       if (seconds === undefined) {
         break
       }
 
-      if (isRetryableDownloadError(error)) {
-        warning(
-          `Socket Firewall binary download from ${url} failed (attempt ${attempt + 1} of ${DOWNLOAD_RETRY_DELAYS_SECONDS.length + 1}): ${errorMessage(error)}. Retrying in ${seconds}s.`,
-        )
-        await setTimeout(seconds * 1000)
-      } else if (attempt < urls.length - 1) {
-        warning(
-          `Socket Firewall binary download from ${url} failed: ${errorMessage(error)}. Trying ${urls[(attempt + 1) % urls.length]}.`,
-        )
-      } else {
-        break
-      }
+      warning(
+        `Socket Firewall binary download from ${url} failed (attempt ${attempt + 1} of ${DOWNLOAD_RETRY_DELAYS_SECONDS.length + 1}): ${errorMessage(error)}. Retrying in ${seconds}s.`,
+      )
+      await setTimeout(seconds * 1000)
+      attempt += 1
     }
   }
 

src/tools/firewall.js

@@ -215,6 +215,54 @@ describe('downloadToolWithRetry', () => {
     expect(mockDownloadTool).toHaveBeenCalledTimes(2)
     expect(sleepDelays).toEqual([])
   })
+
+  it('keeps the healthy origin on the delay table after a miss', async () => {
+    mockDownloadTool
+      .mockRejectedValueOnce(httpError(504))
+      .mockRejectedValueOnce(httpError(404))
+      .mockResolvedValueOnce('/tmp/sfw')
+
+    await expect(downloadToolWithRetry([GITHUB, MIRROR])).resolves.toBe(
+      '/tmp/sfw',
+    )
+    expect(mockDownloadTool.mock.calls).toEqual([[GITHUB], [MIRROR], [GITHUB]])
+    expect(sleepDelays).toEqual([30_000])
+  })
+
+  it('does not spend the delay table on a miss before a flaky origin', async () => {
+    mockDownloadTool
+      .mockRejectedValueOnce(httpError(404))
+      .mockRejectedValue(httpError(504))
+
+    await expect(downloadToolWithRetry([MIRROR, GITHUB])).rejects.toThrow(
+      'Unexpected HTTP response: 504',
+    )
+    expect(mockDownloadTool.mock.calls).toEqual([
+      [MIRROR],
+      [GITHUB],
+      [GITHUB],
+      [GITHUB],
+    ])
+    expect(sleepDelays).toEqual([30_000, 60_000])
+  })
+
+  it('rethrows the flaky origin after a later miss burns no retries', async () => {
+    mockDownloadTool
+      .mockRejectedValueOnce(httpError(504))
+      .mockRejectedValueOnce(httpError(404))
+      .mockRejectedValue(httpError(504))
+
+    await expect(downloadToolWithRetry([GITHUB, MIRROR])).rejects.toThrow(
+      'Unexpected HTTP response: 504',
+    )
+    expect(mockDownloadTool.mock.calls).toEqual([
+      [GITHUB],
+      [MIRROR],
+      [GITHUB],
+      [GITHUB],
+    ])
+    expect(sleepDelays).toEqual([30_000, 60_000])
+  })
 })
 
 describe('firewallDownloadUrls', () => {