refactor: unify tarball resolution and optimize download latency (#468)

Author: aprendendofelipe

packages/app/e2e.test.ts

@@ -1,4 +1,3 @@
-import type { Response } from "@cloudflare/workers-types";
 import ezSpawn from "@jsdevtools/ez-spawn";
 import { simulation } from "@simulacrum/github-api-simulator";
 import fs from "node:fs/promises";
@@ -171,9 +170,15 @@ describe.sequential.each([
     expect(shaBlob.size).toBeGreaterThan(0);
 
     // Test download with ref matches SHA content
-    const refResponse = await fetchWithRedirect(
+    const refResponse = await worker.fetch(
       `/${owner}/${repo}/playground-a@${ref}`,
     );
+    expect(refResponse.status).toBe(200);
+    expect(refResponse.headers.get("x-pkg-name-key")).toBe("playground-a");
+    expect(refResponse.headers.get("x-commit-key")).toBe(
+      `${owner}:${repo}:${fullSha}`,
+    );
+
     const refBlob = await refResponse.blob();
     const shaBlobSize = await shaBlob.arrayBuffer();
     const refBlobSize = await refBlob.arrayBuffer();
@@ -197,6 +202,29 @@ describe.sequential.each([
     );
   }, 20_000);
 
+  it(`returns metadata for HEAD requests (${mode})`, async () => {
+    const [owner, repo] = payload.repository.full_name.split("/");
+    const fullSha = pr ? payload.workflow_run.head_sha : gitRevParse;
+    const sha = fullSha.substring(0, 7);
+
+    const headResponse = await worker.fetch(
+      `/${owner}/${repo}/playground-a@${sha}`,
+      { method: "HEAD" },
+    );
+
+    expect(headResponse.status).toBe(200);
+    expect(headResponse.headers.get("x-pkg-name-key")).toBe("playground-a");
+    expect(headResponse.headers.get("x-commit-key")).toBe(
+      `${owner}:${repo}:${sha}`,
+    );
+    expect(headResponse.headers.get("content-type")).toBe(
+      "application/tar+gzip",
+    );
+    expect(headResponse.headers.get("etag")).toBeDefined();
+    const lastModified = headResponse.headers.get("last-modified");
+    expect(new Date(lastModified!).toString()).not.toBe("Invalid Date");
+  });
+
   it(`serves and installs playground-b for ${mode}`, async () => {
     const [owner, repo] = payload.repository.full_name.split("/");
     const fullSha = pr ? payload.workflow_run.head_sha : gitRevParse;
@@ -231,48 +259,57 @@ describe.sequential.each([
   }, 20_000);
 });
 
-describe("URL redirects", () => {
+describe("URL resolution", () => {
   describe("standard packages", () => {
-    it("redirects full URLs correctly", async () => {
-      const response = await fetchWithRedirect("/tinylibs/tinybench@a832a55");
-      expect(response.url).toContain("/tinylibs/tinybench/tinybench@a832a55");
+    it.each([
+      ["full", "/tinylibs/tinybench/tinybench@a832a55"],
+      ["compact", "/tinybench@a832a55"],
+      ["with .tgz extension", "/tinybench@a832a55.tgz"],
+    ])("resolves %s URLs", async (_, url) => {
+      const response = await worker.fetch(url);
+
+      expect(response.headers.get("x-commit-key")).toBe(
+        "tinylibs:tinybench:a832a55",
+      );
+      expect(response.headers.get("x-pkg-name-key")).toBe("tinybench");
     });
 
-    it("redirects compact URLs correctly", async () => {
-      const response = await fetchWithRedirect("/tinybench@a832a55");
-      expect(response.url).toContain("/tinylibs/tinybench/tinybench@a832a55");
+    it("resolves URL with full Git SHA", async () => {
+      const response = await worker.fetch(
+        "/tinylibs/tinybench/tinybench@a832a55e8f50c419ed8414024899e37e69b1f999",
+      );
+
+      expect(response.headers.get("x-pkg-name-key")).toBe("tinybench");
+      expect(response.headers.get("x-commit-key")).toBe(
+        "tinylibs:tinybench:a832a55e8f50c419ed8414024899e37e69b1f999",
+      );
     });
   });
 
   describe("scoped packages", () => {
-    const expectedPath = `/stackblitz/sdk/${encodeURIComponent("@stackblitz/sdk")}@a832a55`;
+    it.each([
+      ["full", "/stackblitz/sdk/@stackblitz/sdk@a832a55"],
+      ["encoded", "/stackblitz/sdk/%40stackblitz%2Fsdk@a832a55"],
+      ["compact", "/@stackblitz/sdk@a832a55"],
+      ["compact encoded", "/%40stackblitz%2Fsdk@a832a55"],
+    ])("resolves %s URLs", async (_, url) => {
+      const response = await worker.fetch(url);
 
-    it("redirects full scoped package URLs correctly", async () => {
-      const response = await fetchWithRedirect(
-        "/stackblitz/sdk/@stackblitz/sdk@a832a55",
+      expect(response.headers.get("x-pkg-name-key")).toBe("@stackblitz:sdk");
+      expect(response.headers.get("x-commit-key")).toBe(
+        "stackblitz:sdk:a832a55",
       );
-      expect(response.url).toContain(expectedPath);
     });
 
-    it("redirects compact scoped package URLs correctly", async () => {
-      const response = await fetchWithRedirect("/@stackblitz/sdk@a832a55");
-      expect(response.url).toContain(expectedPath);
+    it("resolves URL with full Git SHA", async () => {
+      const response = await worker.fetch(
+        "/stackblitz/sdk/@stackblitz/sdk@a832a55e8f50c419ed8414024899e37e69b1f999",
+      );
+
+      expect(response.headers.get("x-pkg-name-key")).toBe("@stackblitz:sdk");
+      expect(response.headers.get("x-commit-key")).toBe(
+        "stackblitz:sdk:a832a55e8f50c419ed8414024899e37e69b1f999",
+      );
     });
   });
 });
-
-async function fetchWithRedirect(
-  url: string,
-  maxRedirects = 999,
-): Promise<Response> {
-  const response = await worker.fetch(url, { redirect: "manual" });
-
-  if (response.status >= 300 && response.status < 400 && maxRedirects > 0) {
-    const location = response.headers.get("location");
-    if (location) {
-      return fetchWithRedirect(location, maxRedirects - 1);
-    }
-  }
-
-  return response as unknown as Response;
-}

packages/app/fixtures/workflow_run.in_progress.json

@@ -6,7 +6,7 @@
       "id": 9394452824,
       "name": "Preview & Release",
       "node_id": "WFR_kwLOLiqblM8AAAACL_P5WA",
-      "head_branch": "main",
+      "head_branch": "@test/@chaotic.branch/name-@v1.@",
       "head_sha": "ded05e838c418096e5dd77a29101c8af9e73daea",
       "path": ".github/workflows/ci.yml",
       "display_title": "chore: 007 (#94)",

packages/app/nuxt.config.ts

@@ -1,6 +1,3 @@
-// import ncb from "nitro-cloudflare-dev";
-import { resolve } from "pathe";
-
 // https://nuxt.com/docs/api/configuration/nuxt-config
 export default defineNuxtConfig({
   sourcemap: true,
@@ -66,23 +63,6 @@ export default defineNuxtConfig({
     test: "",
   },
 
-  hooks: {
-    "nitro:build:before": (nitro) => {
-      // Override the server routes with the client routes so they are higher priority
-      const clientRenderer = resolve(
-        "node_modules/nuxt/dist/core/runtime/nitro/renderer",
-      );
-      nitro.options.handlers.unshift({
-        route: "/",
-        handler: clientRenderer,
-      });
-      nitro.options.handlers.unshift({
-        route: "/~/**",
-        handler: clientRenderer,
-      });
-    },
-  },
-
   icon: {
     clientBundle: {
       icons: ["mdi-github"],

packages/app/server/middleware/tarball-resolver.ts

@@ -0,0 +1,143 @@
+import {
+  extractOwnerAndRepo,
+  extractRepository,
+  isValidGitHash,
+} from "@pkg-pr-new/utils";
+import { getPackageManifest } from "query-registry";
+import { normalizeKey } from "unstorage";
+
+const RESERVED_ROUTES = new Set(["api", "~"]);
+const ALLOWED_METHODS = new Set(["GET", "HEAD"]);
+
+export default eventHandler(async (event) => {
+  let decodedPath: string;
+  try {
+    const path = event.path
+      .split("?")[0]
+      // yarn support
+      .replace(/\.tgz$/, "");
+    decodedPath = decodeURIComponent(path);
+  } catch {
+    throw createError({
+      statusCode: 400,
+      message: "Malformed URI",
+    });
+  }
+
+  let separatorIndex = -1;
+
+  for (let i = 2; i < decodedPath.length - 1; i++) {
+    if (decodedPath[i] === "@" && decodedPath[i - 1] !== "/") {
+      separatorIndex = i;
+      break;
+    }
+  }
+
+  if (separatorIndex === -1) return;
+
+  let refOrSha = decodedPath.slice(separatorIndex + 1);
+
+  const pathSegments = decodedPath
+    .slice(0, separatorIndex)
+    .split("/")
+    .filter(Boolean);
+  if (pathSegments.length === 0) return;
+
+  const rootSegment = pathSegments[0];
+  if (RESERVED_ROUTES.has(rootSegment) || !/^[a-z0-9@]/i.test(rootSegment)) {
+    return;
+  }
+  if (!ALLOWED_METHODS.has(event.method)) return;
+
+  let packageName = pathSegments.pop()!;
+
+  if (pathSegments.at(-1)?.startsWith("@")) {
+    packageName = `${pathSegments.pop()}/${packageName}`;
+  }
+
+  let owner = pathSegments.shift();
+  let repo = pathSegments.shift() ?? (owner ? packageName : undefined);
+
+  if (pathSegments.length > 0) return;
+
+  if (!repo) {
+    try {
+      const manifest = await getPackageManifest(packageName);
+
+      const repository = extractRepository(manifest);
+      if (!repository) throw new Error();
+
+      const match = extractOwnerAndRepo(repository);
+      if (!match) throw new Error();
+
+      [owner, repo] = match;
+    } catch {
+      throw createError({
+        statusCode: 404,
+        message: "Registry or repository not found",
+      });
+    }
+  }
+
+  const isFullGitHash = isValidGitHash(refOrSha);
+  if (!isFullGitHash) {
+    const cursorBucket = useCursorsBucket(event);
+    const cursorKey = `${owner}:${repo}:${refOrSha}`;
+    const currentCursor = await cursorBucket.getItem(cursorKey);
+
+    if (currentCursor) {
+      refOrSha = currentCursor.sha;
+    }
+  }
+
+  const repositoryCommitKey = `${owner}:${repo}:${refOrSha}`;
+  setResponseHeader(event, "x-commit-key", repositoryCommitKey);
+
+  const normalizedPkgName = normalizeKey(packageName);
+  setResponseHeader(event, "x-pkg-name-key", normalizedPkgName);
+
+  const prefix = `${usePackagesBucket.base}:${repositoryCommitKey}`;
+
+  const binding = useBinding(event);
+  const { objects } = await binding.list({ prefix });
+
+  const packageMetadata = objects.find(({ key }) => {
+    // bucket:package:stackblitz-labs:pkg.pr.new:ded05e838c418096e5dd77a29101c8af9e73daea:playground-b
+    if (!key.endsWith(normalizedPkgName)) return false;
+
+    // ...:playground-b
+    const remainder = key.slice(prefix.length);
+    const colonIdx = remainder.indexOf(":");
+
+    return remainder.slice(colonIdx + 1) === normalizedPkgName;
+  });
+
+  if (!packageMetadata) {
+    throw createError({
+      statusCode: 404,
+      message: "Pkg not found",
+    });
+  }
+
+  setResponseHeader(event, "content-type", "application/tar+gzip");
+  setResponseHeader(event, "etag", packageMetadata.etag);
+  setResponseHeader(
+    event,
+    "last-modified",
+    packageMetadata.uploaded.toUTCString(),
+  );
+
+  if (event.method === "HEAD") {
+    setResponseStatus(event, 200);
+    return send(event, null);
+  }
+
+  const downloadedAtBucket = useDownloadedAtBucket(event);
+  event.waitUntil(downloadedAtBucket.setItem(packageMetadata.key, Date.now()));
+
+  const object = await binding.get(packageMetadata.key);
+  const stream = object?.body;
+
+  // TODO: add HTTP caching
+  return stream;
+});