Speed up version client by partial response reads (#807)

__tests__/download/manifest.test.ts

@@ -17,6 +17,7 @@ const {
   fetchManifest,
   getAllVersions,
   getArtifact,
+  getFirstMatchingVersion,
   getLatestVersion,
   parseManifest,
 } = await import("../../src/download/manifest");
@@ -33,6 +34,7 @@ function createMockResponse(
   data: string,
 ) {
   return {
+    body: null,
     ok,
     status,
     statusText,
@@ -40,6 +42,31 @@ function createMockResponse(
   };
 }
 
+function createStreamingMockResponse(
+  chunks: string[],
+  cancel: () => void | Promise<void> = () => {},
+  close = false,
+) {
+  const encoder = new TextEncoder();
+  return {
+    body: new ReadableStream<Uint8Array>({
+      cancel,
+      start(controller) {
+        for (const chunk of chunks) {
+          controller.enqueue(encoder.encode(chunk));
+        }
+        if (close) {
+          controller.close();
+        }
+      },
+    }),
+    ok: true,
+    status: 200,
+    statusText: "OK",
+    text: async () => chunks.join(""),
+  };
+}
+
 describe("manifest", () => {
   beforeEach(() => {
     clearManifestCache();
@@ -105,6 +132,75 @@ describe("manifest", () => {
         getLatestVersion("https://example.com/custom.ndjson"),
       ).resolves.toBe("0.9.26");
     });
+
+    it("stops reading the default manifest after the first record", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      const cancel = jest.fn();
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse(
+          [
+            latestVersion.slice(0, 100),
+            `${latestVersion.slice(100)}\n`,
+            "invalid trailing data\n",
+          ],
+          cancel,
+        ),
+      );
+
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+      await expect(
+        getArtifact("0.9.26", "aarch64", "apple-darwin"),
+      ).resolves.toBeDefined();
+
+      expect(cancel).toHaveBeenCalledTimes(1);
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+    });
+
+    it("does not fail when canceling the remaining response fails", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse([`${latestVersion}\n`], () =>
+          Promise.reject(new Error("cancel failed")),
+        ),
+      );
+
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+    });
+  });
+
+  describe("getFirstMatchingVersion", () => {
+    it("stops at the first matching record", async () => {
+      const cancel = jest.fn();
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse(
+          [`${sampleManifestResponse}\n`, "invalid trailing data\n"],
+          cancel,
+        ),
+      );
+
+      await expect(
+        getFirstMatchingVersion((version) => version === "0.9.25"),
+      ).resolves.toBe("0.9.25");
+
+      expect(cancel).toHaveBeenCalledTimes(1);
+    });
+
+    it("caches a fully consumed response stream", async () => {
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse(
+          [`${sampleManifestResponse}\n`],
+          undefined,
+          true,
+        ),
+      );
+
+      await expect(
+        getFirstMatchingVersion((version) => version === "0.0.1"),
+      ).resolves.toBeUndefined();
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+    });
   });
 
   describe("getArtifact", () => {
@@ -168,6 +264,60 @@ describe("manifest", () => {
 
       expect(artifact).toBeUndefined();
     });
+
+    it("does not cache records from a failed stream read", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      mockFetch
+        .mockResolvedValueOnce(
+          createStreamingMockResponse([
+            `${latestVersion}\n`,
+            "invalid manifest record\n",
+          ]),
+        )
+        .mockResolvedValueOnce(
+          createMockResponse(true, 200, "OK", sampleManifestResponse),
+        );
+
+      await expect(
+        getArtifact("0.0.1", "aarch64", "apple-darwin"),
+      ).rejects.toThrow("Failed to parse manifest data");
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+
+      expect(mockFetch).toHaveBeenCalledTimes(2);
+    });
+
+    it("does not cache records when the response stream fails", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      const encoder = new TextEncoder();
+      let sentVersion = false;
+      const body = new ReadableStream<Uint8Array>({
+        pull(controller) {
+          if (!sentVersion) {
+            sentVersion = true;
+            controller.enqueue(encoder.encode(`${latestVersion}\n`));
+            return;
+          }
+          controller.error(new Error("response stream failed"));
+        },
+      });
+      mockFetch
+        .mockResolvedValueOnce({
+          body,
+          ok: true,
+          status: 200,
+          statusText: "OK",
+        })
+        .mockResolvedValueOnce(
+          createMockResponse(true, 200, "OK", sampleManifestResponse),
+        );
+
+      await expect(
+        getArtifact("0.0.1", "aarch64", "apple-darwin"),
+      ).rejects.toThrow("response stream failed");
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+
+      expect(mockFetch).toHaveBeenCalledTimes(2);
+    });
   });
 
   describe("parseManifest", () => {