feat: github cache service v2 support

Author: LouisHaftmann

lib/db/index.ts

@@ -247,12 +247,12 @@ export async function pruneKeys(db: DB, keys?: Selectable<CacheKeysTable>[]) {
   }
 }
 
-export async function uploadExists(db: DB, { key, version }: { key: string; version: string }) {
+export async function getUpload(db: DB, { key, version }: { key: string; version: string }) {
   const row = await db
     .selectFrom('uploads')
     .select('id')
     .where('key', '=', key)
     .where('version', '=', version)
     .executeTakeFirst()
-  return !!row
+  return row
 }

lib/storage/index.ts

@@ -7,10 +7,10 @@ import consola from 'consola'
 import {
   findKeyMatch,
   findStaleKeys,
+  getUpload,
   pruneKeys,
   touchKey,
   updateOrCreateKey,
-  uploadExists,
   useDB,
 } from '~/lib/db'
 import { ENV } from '~/lib/env'
@@ -24,7 +24,7 @@ export interface Storage {
     keys: string[],
     version: string,
   ) => Promise<{
-    cacheKey?: string
+    cacheKey: string
     archiveLocation: string
   } | null>
   download: (objectName: string) => Promise<ReadableStream | Readable>
@@ -34,7 +34,7 @@ export interface Storage {
     chunkStart: number,
     chunkEnd: number,
   ) => Promise<void>
-  commitCache: (uploadId: number, size: number) => Promise<void>
+  commitCache: (uploadId: number | string, size: number) => Promise<void>
   reserveCache: (
     key: string,
     version: string,
@@ -71,7 +71,7 @@ export async function initializeStorage() {
         async reserveCache(key, version, totalSize) {
           logger.debug('Reserve:', { key, version })
 
-          if (await uploadExists(db, { key, version })) {
+          if (await getUpload(db, { key, version })) {
             logger.debug(`Reserve: Already reserved. Ignoring...`, { key, version })
             return {
               cacheId: null,

routes/twirp/github.actions.results.api.v1.CacheService/CreateCacheEntry.post.ts

@@ -0,0 +1,32 @@
+import { z } from 'zod'
+import { ENV } from '~/lib/env'
+import { useStorageAdapter } from '~/lib/storage'
+
+const bodySchema = z.object({
+  key: z.string(),
+  version: z.string(),
+})
+
+export default defineEventHandler(async (event) => {
+  const body = (await readBody(event)) as unknown
+  const parsedBody = bodySchema.safeParse(body)
+  if (!parsedBody.success)
+    throw createError({
+      statusCode: 400,
+      statusMessage: `Invalid body: ${parsedBody.error.message}`,
+    })
+
+  const { key, version } = parsedBody.data
+
+  const adapter = await useStorageAdapter()
+  const reservation = await adapter.reserveCache(key, version)
+  if (!reservation.cacheId)
+    return {
+      ok: false,
+    }
+
+  return {
+    ok: true,
+    signed_upload_url: `${ENV.API_BASE_URL}/upload/${reservation.cacheId}`,
+  }
+})

routes/twirp/github.actions.results.api.v1.CacheService/FinalizeCacheEntryUpload.ts

@@ -0,0 +1,36 @@
+import { z } from 'zod'
+import { getUpload, useDB } from '~/lib/db'
+import { useStorageAdapter } from '~/lib/storage'
+
+const bodySchema = z.object({
+  key: z.string(),
+  size_bytes: z.coerce.number(),
+  version: z.string(),
+})
+
+export default defineEventHandler(async (event) => {
+  const parsedBody = bodySchema.safeParse(await readBody(event))
+  if (!parsedBody.success)
+    throw createError({
+      statusCode: 400,
+      statusMessage: `Invalid body: ${parsedBody.error.message}`,
+    })
+
+  const { key, size_bytes, version } = parsedBody.data
+
+  const db = await useDB()
+  const adapter = await useStorageAdapter()
+  const upload = await getUpload(db, { key, version })
+  if (!upload)
+    throw createError({
+      statusCode: 404,
+      statusMessage: 'Upload not found',
+    })
+
+  await adapter.commitCache(upload.id, size_bytes)
+
+  return {
+    ok: true,
+    entry_id: upload.id,
+  }
+})

routes/twirp/github.actions.results.api.v1.CacheService/GetCacheEntryDownloadURL.post.ts

@@ -0,0 +1,32 @@
+import { z } from 'zod'
+import { useStorageAdapter } from '~/lib/storage'
+
+const bodySchema = z.object({
+  key: z.string(),
+  restore_keys: z.array(z.string()).nullish().optional(),
+  version: z.string(),
+})
+
+export default defineEventHandler(async (event) => {
+  const parsedBody = bodySchema.safeParse(await readBody(event))
+  if (!parsedBody.success)
+    throw createError({
+      statusCode: 400,
+      statusMessage: `Invalid body: ${parsedBody.error.message}`,
+    })
+
+  const { key, restore_keys, version } = parsedBody.data
+  const adapter = await useStorageAdapter()
+  const storageEntry = await adapter.getCacheEntry([key, ...(restore_keys ?? [])], version)
+
+  if (!storageEntry)
+    return {
+      ok: false,
+    }
+
+  return {
+    ok: true,
+    signed_download_url: storageEntry.archiveLocation,
+    matched_key: storageEntry.cacheKey,
+  }
+})

routes/upload/[cacheId].put.ts

@@ -0,0 +1,76 @@
+import { Buffer } from 'node:buffer'
+
+import { z } from 'zod'
+import { logger } from '~/lib/logger'
+
+import { useStorageAdapter } from '~/lib/storage'
+
+// https://github.com/actions/toolkit/blob/340a6b15b5879eefe1412ee6c8606978b091d3e8/packages/cache/src/cache.ts#L470
+const chunkSize = 64 * 1024 * 1024
+
+const pathParamsSchema = z.object({
+  cacheId: z.coerce.number(),
+})
+
+const sizeByBlockId = new Map<string, number>()
+
+export default defineEventHandler(async (event) => {
+  const parsedPathParams = pathParamsSchema.safeParse(event.context.params)
+  if (!parsedPathParams.success)
+    throw createError({
+      statusCode: 400,
+      statusMessage: `Invalid path parameters: ${parsedPathParams.error.message}`,
+    })
+
+  if (getQuery(event).comp === 'blocklist') {
+    setResponseStatus(event, 201)
+    return 'ok'
+  }
+
+  const blockId = getQuery(event)?.blockid as string
+  // if no block id, upload smaller than chunk size
+  const chunkIndex = blockId ? getChunkIndexFromBlockId(blockId) : 0
+  if (chunkIndex === undefined)
+    throw createError({
+      statusCode: 400,
+      statusMessage: `Invalid block id: ${blockId}`,
+    })
+
+  const { cacheId } = parsedPathParams.data
+
+  const stream = getRequestWebStream(event)
+  if (!stream) {
+    logger.debug('Upload: Request body is not a stream')
+    throw createError({ statusCode: 400, statusMessage: 'Request body must be a stream' })
+  }
+
+  const contentLengthHeader = getHeader(event, 'content-length')
+  const contentLength = Number.parseInt(contentLengthHeader ?? '')
+  if (!contentLengthHeader || Number.isNaN(contentLength)) {
+    logger.debug("Upload: 'content-length' header not found")
+    throw createError({ statusCode: 400, statusMessage: "'content-length' header is required" })
+  }
+
+  sizeByBlockId.set(blockId, contentLength)
+  const start = chunkIndex * chunkSize
+  const end = start + contentLength - 1
+
+  const adapter = await useStorageAdapter()
+  await adapter.uploadChunk(cacheId, stream as ReadableStream<Buffer>, start, end)
+
+  setResponseStatus(event, 201)
+})
+
+/**
+ * Format (base64 decoded): 06a9ffa8-2e62-4e96-8e5b-15f24c117f1f000000000006
+ */
+function getChunkIndexFromBlockId(blockId: string) {
+  const decoded = Buffer.from(blockId, 'base64').toString('utf8')
+  if (decoded.length !== 48) return
+
+  // slice off uuid and convert to number
+  const index = Number.parseInt(decoded.slice(36))
+  if (Number.isNaN(index)) return
+
+  return index
+}

tests/.env.base

@@ -2,3 +2,4 @@ API_BASE_URL=http://localhost:3000
 NODE_ENV=development
 RUNNER_TEMP=tests/temp/runner
 ACTIONS_CACHE_URL=http://localhost:3000/
+ACTIONS_RESULTS_URL=http://localhost:3000/

tests/e2e.test.ts

@@ -3,29 +3,44 @@ import fs from 'node:fs/promises'
 import path from 'node:path'
 
 import { restoreCache, saveCache } from '@actions/cache'
-import { describe, expect, test } from 'vitest'
+import { afterAll, beforeAll, describe, expect, test } from 'vitest'
 
 const TEST_TEMP_DIR = path.join(import.meta.dirname, 'temp')
 await fs.mkdir(TEST_TEMP_DIR, { recursive: true })
 const testFilePath = path.join(TEST_TEMP_DIR, 'test.bin')
 
 const MB = 1024 * 1024
 
-describe('save and restore cache with @actions/cache package', () => {
-  for (const size of [5 * MB, 50 * MB, 500 * MB, 1024 * MB])
-    test(`${size} Bytes`, { timeout: 90_000 }, async () => {
-      // save
-      const expectedContents = crypto.randomBytes(size)
-      await fs.writeFile(testFilePath, expectedContents)
-      await saveCache([testFilePath], 'cache-key')
-      await fs.rm(testFilePath)
-
-      // restore
-      const cacheHitKey = await restoreCache([testFilePath], 'cache-key')
-      expect(cacheHitKey).toBe('cache-key')
-
-      // check contents
-      const restoredContents = await fs.readFile(testFilePath)
-      expect(restoredContents.compare(expectedContents)).toBe(0)
+const versions = ['v2', 'v1'] as const
+
+for (const version of versions) {
+  describe(`save and restore cache with @actions/cache package with api ${version}`, () => {
+    beforeAll(() => {
+      if (version !== 'v2') return
+
+      process.env.ACTIONS_CACHE_SERVICE_V2 = 'true'
+      process.env.ACTIONS_RUNTIME_TOKEN = 'mock-runtime-token'
+    })
+    afterAll(() => {
+      delete process.env.ACTIONS_CACHE_SERVICE_V2
+      delete process.env.ACTIONS_RUNTIME_TOKEN
     })
-})
+
+    for (const size of [5 * MB, 50 * MB, 500 * MB, 1024 * MB])
+      test(`${size} Bytes`, { timeout: 90_000 }, async () => {
+        // save
+        const expectedContents = crypto.randomBytes(size)
+        await fs.writeFile(testFilePath, expectedContents)
+        await saveCache([testFilePath], 'cache-key')
+        await fs.rm(testFilePath)
+
+        // restore
+        const cacheHitKey = await restoreCache([testFilePath], 'cache-key')
+        expect(cacheHitKey).toBe('cache-key')
+
+        // check contents
+        const restoredContents = await fs.readFile(testFilePath)
+        expect(restoredContents.compare(expectedContents)).toBe(0)
+      })
+  })
+}