rustfs-console/lib/upload-task.ts

import {
  CompleteMultipartUploadCommand,
  CreateMultipartUploadCommand,
  PutObjectCommand,
  S3Client,
  UploadPartCommand,
} from "@aws-sdk/client-s3"
import type { ManagedTask, TaskHandler, TaskLifecycleStatus } from "./task-manager"
import { formatBytes } from "./functions"
import { createTaskId } from "./task-id"
import { getUploadContentType } from "./upload-content-type"
import { logger } from "./logger"

export type UploadStatus = "pending" | "running" | "completed" | "failed" | "canceled"

export interface UploadTask extends ManagedTask<UploadStatus> {
  kind: "upload"
  file: File
  bucketName: string
  prefix?: string
  key: string
  actionLabel: string
  displayName: string
  subInfo: string
  uploadId?: string
  completedParts?: { ETag: string; PartNumber: number }[]
}

export interface UploadTaskConfig {
  chunkSize?: number // MB
  maxRetries?: number
  retryDelay?: number
}

export interface UploadTaskHelpers {
  handler: TaskHandler<UploadTask>
  createTasks: (items: { file: File; key: string }[], bucketName: string) => UploadTask[]
}

const lifecycle: TaskLifecycleStatus<UploadStatus> = {
  pending: "pending",
  running: "running",
  completed: "completed",
  failed: "failed",
  canceled: "canceled",
}

export function createUploadTaskHelpers(s3Client: S3Client, config: UploadTaskConfig = {}): UploadTaskHelpers {
  const chunkSize = (config.chunkSize ?? 5) * 1024 * 1024
  const maxRetries = config.maxRetries ?? 3
  const retryDelay = config.retryDelay ?? 1000

  const perform: TaskHandler<UploadTask>["perform"] = async (task) => {
    task.progress = 0
    if (task.file.size > chunkSize) {
      await multipartUpload(task, s3Client, chunkSize)
    } else {
      await putObject(task, s3Client)
    }
  }

  const shouldRetry: TaskHandler<UploadTask>["shouldRetry"] = (task, error) => {
    if (task.status === lifecycle.canceled) return false
    if ((task.retryCount || 0) >= maxRetries) return false
    const err = error as { $metadata?: { httpStatusCode?: number }; message?: string }
    const statusCode = err?.$metadata?.httpStatusCode
    if (statusCode && statusCode >= 400 && statusCode < 500 && statusCode !== 408 && statusCode !== 429) return false
    const errorMessage = (err?.message || "").toLowerCase()
    const nonRetryableErrors = [
      "access denied",
      "forbidden",
      "invalid credentials",
      "bucket not found",
      "file not found",
    ]
    return !nonRetryableErrors.some((msg) => errorMessage.includes(msg))
  }

  const handler: TaskHandler<UploadTask> = {
    lifecycle,
    perform,
    shouldRetry,
    isCanceledError: (error) =>
      (error as Error)?.name === "AbortError" || (error as Error)?.message?.includes("canceled"),
    maxRetries,
    retryDelay,
  }

  const createTasks = (items: { file: File; key: string }[], bucketName: string): UploadTask[] => {
    const existKeys = new Set<string>()
    return items
      .filter((item) => !existKeys.has(`${bucketName}/${item.key}`))
      .map((item) => {
        existKeys.add(`${bucketName}/${item.key}`)
        const displayName = item.key.split("/").pop() ?? item.key
        return {
          id: createTaskId(item.key),
          kind: "upload" as const,
          file: item.file,
          key: item.key,
          bucketName,
          status: lifecycle.pending,
          progress: 0,
          actionLabel: "Upload",
          displayName,
          subInfo: formatBytes(item.file.size),
          retryCount: 0,
        }
      })
  }

  return { handler, createTasks }
}

async function putObject(task: UploadTask, s3Client: S3Client) {
  const { file, bucketName, key } = task
  const abortController = new AbortController()
  task.abortController = abortController
  const contentType = await getUploadContentType(file, key)

  await s3Client.send(
    new PutObjectCommand({
      Bucket: bucketName,
      Key: key,
      Body: file,
      ContentType: contentType,
    }),
    { abortSignal: abortController.signal },
  )
  task.progress = 100
}

async function multipartUpload(task: UploadTask, s3Client: S3Client, chunkSize: number) {
  const { file, bucketName, key } = task
  const abortController = new AbortController()
  task.abortController = abortController

  let uploadId: string | undefined = task.uploadId
  const completedParts: { ETag: string; PartNumber: number }[] = task.completedParts || []
  const contentType = await getUploadContentType(file, key)

  try {
    if (!uploadId) {
      const createResponse = await s3Client.send(
        new CreateMultipartUploadCommand({
          Bucket: bucketName,
          Key: key,
          ContentType: contentType,
        }),
        { abortSignal: abortController.signal },
      )
      uploadId = createResponse.UploadId
      task.uploadId = uploadId
      task.completedParts = []
    }

    if (!uploadId) throw new Error("Failed to create multipart upload")

    const totalChunks = Math.ceil(file.size / chunkSize)
    for (let partNumber = 1; partNumber <= totalChunks; partNumber++) {
      if (abortController.signal.aborted) throw new DOMException("Upload canceled", "AbortError")
      if (completedParts.some((p) => p.PartNumber === partNumber)) continue

      const start = (partNumber - 1) * chunkSize
      const end = Math.min(start + chunkSize, file.size)
      const chunk = file.slice(start, end)

      const { ETag } = await s3Client.send(
        new UploadPartCommand({
          Bucket: bucketName,
          Key: key,
          UploadId: uploadId,
          PartNumber: partNumber,
          Body: chunk,
        }),
        { abortSignal: abortController.signal },
      )
      if (!ETag) throw new Error(`Failed to upload part ${partNumber}`)
      completedParts.push({ ETag, PartNumber: partNumber })
      task.completedParts = completedParts
      task.progress = Math.round((completedParts.length / totalChunks) * 100)
    }

    await s3Client.send(
      new CompleteMultipartUploadCommand({
        Bucket: bucketName,
        Key: key,
        UploadId: uploadId,
        MultipartUpload: { Parts: completedParts },
      }),
      { abortSignal: abortController.signal },
    )
    task.uploadId = undefined
    task.completedParts = undefined
  } catch (error) {
    if (uploadId) {
      try {
        const { AbortMultipartUploadCommand } = await import("@aws-sdk/client-s3")
        await s3Client.send(new AbortMultipartUploadCommand({ Bucket: bucketName, Key: key, UploadId: uploadId }))
      } catch (cleanupError) {
        logger.warn("Failed to abort multipart upload", cleanupError)
      } finally {
        task.uploadId = undefined
        task.completedParts = undefined
      }
    }
    throw error
  }
}