storage.ts

/* eslint-disable no-shadow */
/* eslint-disable ts/method-signature-style */
import type { Kysely } from 'kysely'
import type { ReadableStream } from 'node:stream/web'
import type { Database, StorageLocation } from './db'
import type { Env } from './schemas'
import { randomUUID } from 'node:crypto'
import { once } from 'node:events'
import { createReadStream, createWriteStream } from 'node:fs'
import fs from 'node:fs/promises'
import { Agent } from 'node:https'
import path from 'node:path'
import { PassThrough, Readable } from 'node:stream'
import { pipeline } from 'node:stream/promises'
import { createSingletonPromise } from '@antfu/utils'
import {
  DeleteObjectsCommand,
  GetObjectCommand,
  HeadBucketCommand,
  HeadObjectCommand,
  ListObjectsV2Command,
  S3Client,
} from '@aws-sdk/client-s3'
import { Upload as S3Upload } from '@aws-sdk/lib-storage'
import { getSignedUrl } from '@aws-sdk/s3-request-presigner'
import { Storage as GcsClient } from '@google-cloud/storage'
import { NodeHttpHandler } from '@smithy/node-http-handler'
import { sql } from 'kysely'
import { chunk } from 'remeda'
import { match } from 'ts-pattern'
import { getDatabase } from './db'
import { env } from './env'
import { generateNumberId } from './helpers'
import { logger } from './logger'

function escapeLikePattern(value: string) {
  return value
    .replaceAll('\\', '\\\\')
    .replaceAll('%', String.raw`\%`)
    .replaceAll('_', String.raw`\_`)
}

export class ObjectNotFoundError extends Error {
  constructor(objectName: string) {
    super(`Object not found in storage: ${objectName}`)
    this.name = 'ObjectNotFoundError'
  }
}

export const PARTS_DELETE_GRACE_MS = 15 * 60 * 1000

export class Storage {
  adapter
  private db
  private mergeStreamPromises = new Set<Promise<void>>()

  private constructor({ db, adapter }: { adapter: StorageAdapter; db: Kysely<Database> }) {
    this.adapter = adapter
    this.db = db
  }

  static async getAdapterFromEnv() {
    return await match(env)
      .with({ STORAGE_DRIVER: 's3' }, S3Adapter.fromEnv)
      .with({ STORAGE_DRIVER: 'filesystem' }, FileSystemAdapter.fromEnv)
      .with({ STORAGE_DRIVER: 'gcs' }, GcsAdapter.fromEnv)
      .exhaustive()
  }

  static async fromEnv() {
    return new Storage({
      adapter: await Storage.getAdapterFromEnv(),
      db: await getDatabase(),
    })
  }

  waitForOngoingMerges() {
    return Promise.all(this.mergeStreamPromises)
  }

  async uploadPart(uploadId: number, partIndex: number, stream: ReadableStream) {
    const upload = await this.db
      .selectFrom('uploads')
      .where('id', '=', uploadId)
      .select(['folderName'])
      .executeTakeFirst()
    if (!upload) return

    await this.db
      .updateTable('uploads')
      .set({
        startedPartUploadCount: sql`${sql.ref('startedPartUploadCount')} + 1`,
      })
      .where('id', '=', uploadId)
      .execute()

    await this.adapter.uploadStream(
      `${upload.folderName}/parts/${partIndex}`,
      Readable.fromWeb(stream),
    )

    await this.db
      .updateTable('uploads')
      .set({
        lastPartUploadedAt: Date.now(),
        finishedPartUploadCount: sql`${sql.ref('finishedPartUploadCount')} + 1`,
      })
      .where('id', '=', uploadId)
      .execute()
  }

  async completeUpload({
    key,
    version,
    scope,
    repoId,
  }: {
    key: string
    version: string
    scope: string
    repoId: string
  }) {
    const upload = await this.db
      .selectFrom('uploads')
      .where('key', '=', key)
      .where('version', '=', version)
      .where('scope', '=', scope)
      .where('repoId', '=', repoId)
      .selectAll()
      .executeTakeFirst()
    if (!upload) return

    if (upload.finishedPartUploadCount === 0) {
      await this.db.deleteFrom('uploads').where('id', '=', upload.id).execute()
      throw new Error('No parts have been uploaded')
    }

    if (upload.startedPartUploadCount !== upload.finishedPartUploadCount) {
      await this.db.deleteFrom('uploads').where('id', '=', upload.id).execute()
      throw new Error(
        `Not all parts have been uploaded (only ${upload.finishedPartUploadCount} of ${upload.startedPartUploadCount} parts uploaded)`,
      )
    }

    const partCount = await this.adapter.countFilesInFolder(`${upload.folderName}/parts`)
    if (partCount !== upload.finishedPartUploadCount) {
      await this.db.deleteFrom('uploads').where('id', '=', upload.id).execute()
      throw new Error(
        `Uploaded part count does not match actual part count in storage (expected ${upload.finishedPartUploadCount} but found ${partCount})`,
      )
    }

    await this.db.transaction().execute(async (tx) => {
      const locationId = randomUUID()
      await tx
        .insertInto('storage_locations')
        .values({
          id: locationId,
          folderName: upload.folderName,
          partCount,
          mergedAt: null,
          mergeStartedAt: null,
          partsDeletedAt: null,
          lastDownloadedAt: null,
        })
        .execute()

      const existingCacheEntry = await tx
        .selectFrom('cache_entries')
        .where('key', '=', key)
        .where('version', '=', version)
        .where('scope', '=', scope)
        .where('repoId', '=', repoId)
        .innerJoin('storage_locations', 'storage_locations.id', 'cache_entries.locationId')
        .select(['cache_entries.id', 'cache_entries.locationId', 'storage_locations.folderName'])
        .executeTakeFirst()

      if (existingCacheEntry) {
        await tx
          .updateTable('cache_entries')
          .set({
            updatedAt: Date.now(),
            locationId,
          })
          .where('id', '=', existingCacheEntry.id)
          .execute()
        await tx
          .deleteFrom('storage_locations')
          .where('id', '=', existingCacheEntry.locationId)
          .execute()
        await this.adapter.deleteFolder(existingCacheEntry.folderName)
      } else
        await tx
          .insertInto('cache_entries')
          .values({
            key: upload.key,
            version: upload.version,
            id: randomUUID(),
            updatedAt: Date.now(),
            locationId,
            scope,
            repoId,
          })
          .execute()

      await tx.deleteFrom('uploads').where('id', '=', upload.id).execute()
    })

    return upload
  }

  async download(cacheEntryId: string): Promise<Readable | undefined> {
    const storageLocation = await this.db
      .selectFrom('storage_locations')
      .innerJoin('cache_entries', 'cache_entries.locationId', 'storage_locations.id')
      .where('cache_entries.id', '=', cacheEntryId)
      .selectAll('storage_locations')
      .executeTakeFirst()
    if (!storageLocation) return

    void this.db
      .updateTable('storage_locations')
      .set({
        lastDownloadedAt: Date.now(),
      })
      .where('id', '=', storageLocation.id)
      .execute()

    try {
      if (storageLocation.mergedAt || storageLocation.mergeStartedAt)
        return await this.downloadFromCacheEntryLocation(storageLocation)

      await this.ensurePartsExist(storageLocation)

      await this.db
        .updateTable('storage_locations')
        .set({
          mergeStartedAt: Date.now(),
        })
        .where('id', '=', storageLocation.id)
        .execute()

      const responseStream = new PassThrough()
      const mergerStream = new PassThrough()

      const mergePromise = this.adapter
        .uploadStream(`${storageLocation.folderName}/merged`, mergerStream)
        .then(async () => {
          await this.db
            .updateTable('storage_locations')
            .set({
              mergedAt: Date.now(),
            })
            .where('id', '=', storageLocation.id)
            .execute()
        })
        .catch(async () => {
          await this.db
            .updateTable('storage_locations')
            .set({
              mergedAt: null,
              mergeStartedAt: null,
            })
            .where('id', '=', storageLocation.id)
            .execute()
          mergerStream.destroy()
        })
      this.mergeStreamPromises.add(mergePromise)
      mergePromise.finally(() => this.mergeStreamPromises.delete(mergePromise))

      this.pumpPartsToStreams(storageLocation, responseStream, mergerStream).catch((err) => {
        responseStream.destroy(err)
        mergerStream.destroy(err)
        if (err instanceof ObjectNotFoundError) {
          logger.warn(`Stale cache entry ${cacheEntryId}: ${err.message}`)
          void this.deleteStaleCacheEntry(cacheEntryId).catch((deleteErr) => {
            logger.error(`Failed to delete stale cache entry ${cacheEntryId}`, deleteErr)
          })
        }
      })

      return responseStream
    } catch (err) {
      if (err instanceof ObjectNotFoundError) {
        logger.warn(`Stale cache entry ${cacheEntryId}: ${err.message}`)
        await this.deleteStaleCacheEntry(cacheEntryId)
        return
      }
      throw err
    }
  }

  private async ensurePartsExist(location: StorageLocation) {
    const partsFolder = `${location.folderName}/parts`
    const partExists = await Promise.all(
      Array.from({ length: location.partCount }, (_, i) =>
        this.adapter.fileExists(`${partsFolder}/${i}`),
      ),
    )
    const missingPartIndex = partExists.findIndex((exists) => !exists)
    if (missingPartIndex !== -1) throw new ObjectNotFoundError(`${partsFolder}/${missingPartIndex}`)
  }

  private async deleteStaleCacheEntry(cacheEntryId: string) {
    await this.db.deleteFrom('cache_entries').where('id', '=', cacheEntryId).execute()
  }

  private async downloadFromCacheEntryLocation(location: StorageLocation) {
    if (location.mergedAt) return this.adapter.createDownloadStream(`${location.folderName}/merged`)

    await this.ensurePartsExist(location)
    return Readable.from(this.streamParts(location))
  }

  private async pumpPartsToStreams(
    location: StorageLocation,
    responseStream: PassThrough,
    mergerStream: PassThrough,
  ) {
    if (location.partsDeletedAt) throw new Error('No parts to feed')

    for await (const chunk of this.streamParts(location)) {
      const responseWantsMore = responseStream.write(chunk)
      const mergerWantsMore = mergerStream.write(chunk)

      if (!responseWantsMore) await once(responseStream, 'drain')
      if (!mergerWantsMore) await once(mergerStream, 'drain')
    }

    responseStream.end()
    mergerStream.end()

    await globalThis.gc?.()
  }

  private async *streamParts(location: StorageLocation) {
    if (location.partsDeletedAt) throw new Error('No parts to feed for location with deleted parts')

    for (let i = 0; i < location.partCount; i++) {
      const partStream = await this.adapter.createDownloadStream(
        `${location.folderName}/parts/${i}`,
      )

      for await (const chunk of partStream) yield chunk

      await globalThis.gc?.()
    }
  }

  async createUpload({
    key,
    version,
    scope,
    repoId,
  }: {
    key: string
    version: string
    scope: string
    repoId: string
  }) {
    const existingUpload = await this.db
      .selectFrom('uploads')
      .where('key', '=', key)
      .where('version', '=', version)
      .where('scope', '=', scope)
      .where('repoId', '=', repoId)
      .select('id')
      .executeTakeFirst()
    if (existingUpload) return

    const uploadId = generateNumberId()
    await this.db
      .insertInto('uploads')
      .values({
        id: uploadId,
        folderName: uploadId.toString(),
        createdAt: Date.now(),
        key,
        version,
        scope,
        repoId,
        lastPartUploadedAt: null,
        finishedPartUploadCount: 0,
        startedPartUploadCount: 0,
      })
      .execute()

    return { id: uploadId }
  }

  async matchCacheEntry({
    keys: [primaryKey, ...restoreKeys],
    version,
    scopes,
    repoId,
  }: {
    keys: [string, ...string[]]
    version: string
    scopes: string[]
    repoId: string
  }) {
    for (const scope of scopes) {
      const exactPrimaryMatch = await this.db
        .selectFrom('cache_entries')
        .where('key', '=', primaryKey)
        .where('version', '=', version)
        .where('scope', '=', scope)
        .where('repoId', '=', repoId)
        .selectAll()
        .executeTakeFirst()
      if (exactPrimaryMatch)
        return {
          match: exactPrimaryMatch,
          type: 'exact-primary' as const,
        }

      const prefixedPrimaryMatch = await this.db
        .selectFrom('cache_entries')
        .where(
          sql<boolean>`${sql.ref('key')} like ${`${escapeLikePattern(primaryKey)}%`} escape ${'\\'}`,
        )
        .where('version', '=', version)
        .where('scope', '=', scope)
        .where('repoId', '=', repoId)
        .orderBy('cache_entries.updatedAt', 'desc')
        .selectAll()
        .executeTakeFirst()

      if (prefixedPrimaryMatch)
        return {
          match: prefixedPrimaryMatch,
          type: 'prefixed-primary' as const,
        }

      if (restoreKeys.length === 0) continue

      for (const key of restoreKeys) {
        const exactMatch = await this.db
          .selectFrom('cache_entries')
          .where('key', '=', key)
          .where('version', '=', version)
          .where('scope', '=', scope)
          .where('repoId', '=', repoId)
          .orderBy('updatedAt', 'desc')
          .selectAll()
          .executeTakeFirst()
        if (exactMatch)
          return {
            match: exactMatch,
            type: 'exact-restore' as const,
          }

        const prefixedMatch = await this.db
          .selectFrom('cache_entries')
          .where(
            sql<boolean>`${sql.ref('key')} like ${`${escapeLikePattern(key)}%`} escape ${'\\'}`,
          )
          .where('version', '=', version)
          .where('scope', '=', scope)
          .where('repoId', '=', repoId)
          .orderBy('updatedAt', 'desc')
          .selectAll()
          .executeTakeFirst()

        if (prefixedMatch)
          return {
            match: prefixedMatch,
            type: 'prefixed-restore' as const,
          }
      }
    }
  }

  async getCacheEntryWithDownloadUrl(args: Parameters<typeof this.matchCacheEntry>[0]) {
    const cacheEntry = await this.matchCacheEntry(args)
    if (!cacheEntry) return

    const defaultUrl = `${env.API_BASE_URL}/download/${cacheEntry.match.id}`

    if (!env.ENABLE_DIRECT_DOWNLOADS || !this.adapter.createDownloadUrl)
      return {
        downloadUrl: defaultUrl,
        cacheEntry: cacheEntry.match,
      }

    const location = await this.db
      .selectFrom('storage_locations')
      .where('id', '=', cacheEntry.match.locationId)
      .select(['folderName', 'mergedAt'])
      .executeTakeFirst()
    if (!location) throw new Error('Storage location not found')

    const downloadUrl = location.mergedAt
      ? await this.adapter.createDownloadUrl(`${location.folderName}/merged`)
      : defaultUrl

    return {
      downloadUrl,
      cacheEntry: cacheEntry.match,
    }
  }
}

export const getStorage = createSingletonPromise(async () => Storage.fromEnv())

interface StorageAdapter {
  createDownloadStream(objectName: string): Promise<Readable>
  uploadStream(objectName: string, stream: Readable): Promise<void>
  deleteFolder(folderName: string): Promise<void>
  fileExists(objectName: string): Promise<boolean>
  countFilesInFolder(folderName: string): Promise<number>
  createDownloadUrl?(objectName: string): Promise<string>
  clear(): Promise<void>
}

class S3Adapter implements StorageAdapter {
  private s3
  private bucket
  private keyPrefix = 'gh-actions-cache'

  constructor({ bucket, s3 }: { s3: S3Client; bucket: string }) {
    this.s3 = s3
    this.bucket = bucket
  }

  static async fromEnv(env: Extract<Env, { STORAGE_DRIVER: 's3' }>) {
    const bucket = env.STORAGE_S3_BUCKET
    const agent = new Agent({
      keepAlive: true,
      maxSockets: 50,
      keepAliveMsecs: 1000,
    })
    const s3 = new S3Client({
      forcePathStyle: true,
      region: env.AWS_REGION,
      requestHandler: new NodeHttpHandler({
        httpsAgent: agent,
        socketTimeout: 3000,
      }),
    })

    try {
      await s3.send(
        new HeadBucketCommand({
          Bucket: bucket,
        }),
      )
    } catch (err: any) {
      if (err.name === 'NotFound') {
        throw new Error(`Bucket ${bucket} does not exist`)
      }
      throw err
    }

    return new S3Adapter({ s3, bucket })
  }

  async createDownloadStream(objectName: string) {
    try {
      const response = await this.s3.send(
        new GetObjectCommand({
          Bucket: this.bucket,
          Key: `${this.keyPrefix}/${objectName}`,
        }),
      )
      if (!response.Body) throw new Error('No body in S3 get object response')

      return response.Body as Readable
    } catch (err: any) {
      if (err.name === 'NoSuchKey') throw new ObjectNotFoundError(objectName)
      throw err
    }
  }

  async fileExists(objectName: string) {
    try {
      await this.s3.send(
        new HeadObjectCommand({
          Bucket: this.bucket,
          Key: `${this.keyPrefix}/${objectName}`,
        }),
      )
      return true
    } catch (err: any) {
      if (
        err.name === 'NoSuchKey' ||
        err.name === 'NotFound' ||
        err.$metadata?.httpStatusCode === 404
      )
        return false
      throw err
    }
  }

  async deleteFolder(folderName: string) {
    return this.deleteByPrefix(`${this.keyPrefix}/${folderName}/`)
  }

  async clear() {
    return this.deleteByPrefix(this.keyPrefix)
  }

  private async deleteByPrefix(prefix: string) {
    const listResponse = await this.s3.send(
      new ListObjectsV2Command({
        Bucket: this.bucket,
        Prefix: prefix,
      }),
    )

    if (!listResponse.Contents || listResponse.Contents.length === 0) return

    await Promise.all(
      chunk(
        listResponse.Contents.filter((obj): obj is { Key: string } => !!obj.Key),
        1000,
      ).map((chunkedObjects) =>
        this.s3.send(
          new DeleteObjectsCommand({
            Bucket: this.bucket,
            Delete: {
              Objects: chunkedObjects.map((obj) => ({
                Key: obj.Key,
              })),
              Quiet: true,
            },
          }),
        ),
      ),
    )
  }

  async uploadStream(objectName: string, iterator: AsyncIterable<Uint8Array>) {
    await new S3Upload({
      client: this.s3,
      params: {
        Bucket: this.bucket,
        Key: `${this.keyPrefix}/${objectName}`,
        Body: iterator as Readable,
      },
      queueSize: 1,
      partSize: 5 * 1024 * 1024, // 5MB
      leavePartsOnError: false,
    }).done()
  }

  async countFilesInFolder(folderName: string) {
    const listResponse = await this.s3.send(
      new ListObjectsV2Command({
        Bucket: this.bucket,
        Prefix: `${this.keyPrefix}/${folderName}/`,
      }),
    )

    return listResponse.KeyCount ?? 0
  }

  async createDownloadUrl(objectName: string) {
    return getSignedUrl(
      this.s3,
      new GetObjectCommand({
        Bucket: this.bucket,
        Key: `${this.keyPrefix}/${objectName}`,
      }),
      {
        expiresIn: 10 * 60 * 1000, // 10min
      },
    )
  }
}

class FileSystemAdapter implements StorageAdapter {
  private rootFolder

  constructor({ rootFolder }: { rootFolder: string }) {
    this.rootFolder = path.resolve(rootFolder)
  }

  private safePath(name: string) {
    const resolved = path.resolve(this.rootFolder, name)
    if (!resolved.startsWith(this.rootFolder + path.sep) && resolved !== this.rootFolder)
      throw new Error(`Invalid object name`)
    return resolved
  }

  static async fromEnv(env: Extract<Env, { STORAGE_DRIVER: 'filesystem' }>) {
    const rootFolder = env.STORAGE_FILESYSTEM_PATH
    await fs.mkdir(rootFolder, {
      recursive: true,
    })

    return new FileSystemAdapter({
      rootFolder,
    })
  }

  async createDownloadStream(objectName: string) {
    const filePath = this.safePath(objectName)
    try {
      await fs.access(filePath)
    } catch {
      throw new ObjectNotFoundError(objectName)
    }
    return createReadStream(filePath)
  }

  async fileExists(objectName: string) {
    try {
      await fs.access(this.safePath(objectName))
      return true
    } catch (err: any) {
      if (err.code === 'ENOENT') return false
      throw err
    }
  }

  async deleteFolder(folderName: string) {
    await fs.rm(this.safePath(folderName), {
      recursive: true,
      force: true,
    })
  }

  async clear() {
    await fs.rm(this.rootFolder, {
      recursive: true,
      force: true,
    })
    await fs.mkdir(this.rootFolder, {
      recursive: true,
    })
  }

  async uploadStream(objectName: string, stream: Readable) {
    const filePath = this.safePath(objectName)
    await fs.mkdir(path.dirname(filePath), { recursive: true })
    await pipeline(stream, createWriteStream(filePath))
  }

  async countFilesInFolder(folderName: string) {
    try {
      const dir = await fs.readdir(this.safePath(folderName), {
        withFileTypes: true,
      })
      return dir.filter((item) => item.isFile()).length
    } catch (err: any) {
      if (err.code === 'ENOENT') return 0
      throw err
    }
  }
}

class GcsAdapter implements StorageAdapter {
  private bucket
  private keyPrefix = 'gh-actions-cache'

  constructor({ bucket, gcs }: { bucket: string; gcs: GcsClient }) {
    this.bucket = gcs.bucket(bucket)
  }

  static async fromEnv(env: Extract<Env, { STORAGE_DRIVER: 'gcs' }>) {
    const bucketName = env.STORAGE_GCS_BUCKET

    const gcs = new GcsClient({
      keyFilename: env.STORAGE_GCS_SERVICE_ACCOUNT_KEY,
      apiEndpoint: env.STORAGE_GCS_ENDPOINT,
    })
    const bucket = gcs.bucket(bucketName)

    await bucket.getMetadata()

    return new GcsAdapter({
      bucket: bucketName,
      gcs,
    })
  }

  async createDownloadStream(objectName: string) {
    const file = this.bucket.file(`${this.keyPrefix}/${objectName}`)
    const [exists] = await file.exists()
    if (!exists) throw new ObjectNotFoundError(objectName)
    return file.createReadStream()
  }

  async fileExists(objectName: string) {
    const [exists] = await this.bucket.file(`${this.keyPrefix}/${objectName}`).exists()
    return exists
  }

  async deleteFolder(folderName: string) {
    await this.bucket.deleteFiles({
      prefix: `${this.keyPrefix}/${folderName}/`,
    })
  }

  async clear() {
    await this.bucket.deleteFiles({
      prefix: this.keyPrefix,
    })
  }

  async uploadStream(objectName: string, iterator: AsyncIterable<Uint8Array>) {
    const file = this.bucket.file(`${this.keyPrefix}/${objectName}`)

    await pipeline(
      iterator,
      file.createWriteStream({
        resumable: false,
        validation: false,
      }),
    )
  }

  async countFilesInFolder(folderName: string) {
    return this.bucket
      .getFiles({
        prefix: `${this.keyPrefix}/${folderName}/`,
        autoPaginate: true,
      })
      .then((res) => res[0].length)
  }

  async createDownloadUrl(objectName: string) {
    return this.bucket
      .file(`${this.keyPrefix}/${objectName}`)
      .getSignedUrl({
        action: 'read',
        expires: Date.now() + 10 * 60 * 1000, // 10min
      })
      .then((res) => res[0])
  }
}