gongxiangp2/lib/material-upload-storage-reservation.ts

import { statfs } from 'fs/promises'

import { MaterialUploadRequestError } from '@/lib/material-upload-request-guard'

const MEBIBYTE = 1024 * 1024
const GIBIBYTE = 1024 * MEBIBYTE

export const UPLOAD_STORAGE_FREE_RESERVE_BYTES = 2 * GIBIBYTE

type AvailableBytesReader = (directory: string) => Promise<bigint>

async function readAvailableBytes(directory: string): Promise<bigint> {
  const disk = await statfs(directory, { bigint: true })
  return disk.bavail * disk.bsize
}

/**
 * 记录尚未完成的磁盘写入额度。statfs 只能看到已经落盘的字节；如果不把
 * 在途合并/分片也扣除，多个请求会同时通过同一份剩余空间快照。
 */
export class MaterialUploadStorageReservationManager {
  private reservedBytes = BigInt(0)
  private lockTail: Promise<void> = Promise.resolve()

  constructor(
    private readonly minimumFreeBytes: number = UPLOAD_STORAGE_FREE_RESERVE_BYTES,
    private readonly getAvailableBytes: AvailableBytesReader = readAvailableBytes,
  ) {
    if (!Number.isSafeInteger(minimumFreeBytes) || minimumFreeBytes < 0) {
      throw new Error('minimumFreeBytes 必须是非负安全整数')
    }
  }

  async acquire(directory: string, requestedBytes: number): Promise<() => void> {
    if (!Number.isSafeInteger(requestedBytes) || requestedBytes <= 0) {
      throw new MaterialUploadRequestError(
        '上传写盘大小无效',
        400,
        'UPLOAD_FILE_SIZE_INVALID',
      )
    }

    return this.withLock(async () => {
      let availableBytes: bigint
      try {
        availableBytes = await this.getAvailableBytes(directory)
      } catch (error) {
        console.error('[MaterialUploadStorage] 无法确认上传盘剩余空间:', error)
        throw new MaterialUploadRequestError(
          '暂时无法确认上传存储空间，请稍后重试',
          503,
          'UPLOAD_STORAGE_UNKNOWN',
          3,
        )
      }

      const requested = BigInt(requestedBytes)
      const required = this.reservedBytes
        + requested
        + BigInt(this.minimumFreeBytes)
      if (availableBytes < required) {
        throw new MaterialUploadRequestError(
          '上传存储空间不足，请联系管理员清理后重试',
          507,
          'UPLOAD_STORAGE_LOW',
        )
      }

      this.reservedBytes += requested
      let released = false
      return () => {
        if (released) {
          return
        }
        released = true
        this.reservedBytes = this.reservedBytes >= requested
          ? this.reservedBytes - requested
          : BigInt(0)
      }
    })
  }

  snapshot(): { reservedBytes: bigint } {
    return { reservedBytes: this.reservedBytes }
  }

  private async withLock<T>(action: () => Promise<T>): Promise<T> {
    let unlock!: () => void
    const previous = this.lockTail
    this.lockTail = new Promise<void>((resolve) => {
      unlock = resolve
    })

    await previous
    try {
      return await action()
    } finally {
      unlock()
    }
  }
}

export const materialUploadStorageReservations = new MaterialUploadStorageReservationManager()