gongxiangp2/lib/material-upload-merge.ts

import { randomUUID } from 'crypto'
import { lstat, open, rename, rm } from 'fs/promises'
import path from 'path'

import { materialUploadStorageReservations } from '@/lib/material-upload-storage-reservation'

const MEBIBYTE = 1024 * 1024
const COPY_BUFFER_BYTES = MEBIBYTE

export interface UploadChunkToMerge {
  path: string
  expectedBytes: number
}

/**
 * 分片已经占据 fileSize 空间；合并前还需要再写一份 fileSize，并始终留下
 * 2GiB 安全余量。检查失败时必须保守拒绝，不能在未知磁盘状态下继续写入。
 */
export async function acquireUploadMergeDiskReservation(
  directory: string,
  fileSize: number,
): Promise<() => void> {
  return materialUploadStorageReservations.acquire(directory, fileSize)
}

/**
 * 使用固定 1MiB 缓冲区顺序复制到唯一临时文件，完整校验并 fsync 后才原子
 * rename 为正式 merged 文件。任何异常只清理本次临时文件，不会留下半成品，
 * 也不会覆盖已经验证成功的 merged 文件。
 */
export async function mergeUploadChunksAtomically(input: {
  chunks: UploadChunkToMerge[]
  expectedTotalBytes: number
  outputPath: string
  assertLockHeld?: () => void
}): Promise<string> {
  const { chunks, expectedTotalBytes, outputPath } = input
  if (
    chunks.length === 0
    || !Number.isSafeInteger(expectedTotalBytes)
    || expectedTotalBytes <= 0
  ) {
    throw new Error('上传分片合并参数无效')
  }

  let validatedTotalBytes = 0
  for (const chunk of chunks) {
    if (!Number.isSafeInteger(chunk.expectedBytes) || chunk.expectedBytes <= 0) {
      throw new Error(`上传分片大小无效: ${path.basename(chunk.path)}`)
    }
    const stats = await lstat(chunk.path)
    if (
      stats.isSymbolicLink()
      || !stats.isFile()
      || stats.size !== chunk.expectedBytes
    ) {
      throw new Error(`上传分片大小不匹配: ${path.basename(chunk.path)}`)
    }
    validatedTotalBytes += stats.size
    if (
      !Number.isSafeInteger(validatedTotalBytes)
      || validatedTotalBytes > expectedTotalBytes
    ) {
      throw new Error('上传分片总大小超过会话声明')
    }
  }
  if (validatedTotalBytes !== expectedTotalBytes) {
    throw new Error('上传分片总大小与会话声明不一致')
  }

  const partialPath = `${outputPath}.partial-${randomUUID()}`
  let outputHandle: Awaited<ReturnType<typeof open>> | null = null

  try {
    outputHandle = await open(partialPath, 'wx')
    const copyBuffer = Buffer.allocUnsafe(COPY_BUFFER_BYTES)
    let copiedTotalBytes = 0

    for (const chunk of chunks) {
      const chunkHandle = await open(chunk.path, 'r')
      let copiedChunkBytes = 0
      try {
        const openedStats = await chunkHandle.stat()
        if (!openedStats.isFile() || openedStats.size !== chunk.expectedBytes) {
          throw new Error(`上传分片在合并期间发生变化: ${path.basename(chunk.path)}`)
        }

        while (true) {
          const { bytesRead } = await chunkHandle.read(
            copyBuffer,
            0,
            copyBuffer.length,
            null,
          )
          if (bytesRead === 0) {
            break
          }

          copiedChunkBytes += bytesRead
          copiedTotalBytes += bytesRead
          if (
            copiedChunkBytes > chunk.expectedBytes
            || copiedTotalBytes > expectedTotalBytes
          ) {
            throw new Error(`上传分片在合并期间变大: ${path.basename(chunk.path)}`)
          }

          let writtenOffset = 0
          while (writtenOffset < bytesRead) {
            const { bytesWritten } = await outputHandle.write(
              copyBuffer,
              writtenOffset,
              bytesRead - writtenOffset,
              null,
            )
            if (bytesWritten <= 0) {
              throw new Error('上传文件合并写入失败')
            }
            writtenOffset += bytesWritten
          }
        }
      } finally {
        await chunkHandle.close()
      }

      if (copiedChunkBytes !== chunk.expectedBytes) {
        throw new Error(`上传分片在合并期间变小: ${path.basename(chunk.path)}`)
      }
    }

    if (copiedTotalBytes !== expectedTotalBytes) {
      throw new Error('上传文件合并大小与会话声明不一致')
    }

    await outputHandle.sync()
    await outputHandle.close()
    outputHandle = null

    const partialStats = await lstat(partialPath)
    if (
      partialStats.isSymbolicLink()
      || !partialStats.isFile()
      || partialStats.size !== expectedTotalBytes
    ) {
      throw new Error('上传文件合并结果校验失败')
    }

    input.assertLockHeld?.()
    await rename(partialPath, outputPath)
    return outputPath
  } catch (error) {
    if (outputHandle) {
      await outputHandle.close().catch(() => undefined)
    }
    await rm(partialPath, { force: true }).catch(() => undefined)
    throw error
  }
}