import { TaskItem } from "@/pages/DataManagement/dataset.model"; import { calculateSHA256, checkIsFilesExist, streamSplitAndUpload, StreamUploadResult } from "@/utils/file.util"; import { App } from "antd"; import { useRef, useState } from "react"; export function useFileSliceUpload( { preUpload, uploadChunk, cancelUpload, }: { preUpload: (id: string, params: Record) => Promise<{ data: number }>; uploadChunk: (id: string, formData: FormData, config: Record) => Promise; cancelUpload: ((reqId: number) => Promise) | null; }, showTaskCenter = true, // 上传时是否显示任务中心 enableStreamUpload = true // 是否启用流式分割上传 ) { const { message } = App.useApp(); const [taskList, setTaskList] = useState([]); const taskListRef = useRef([]); // 用于固定任务顺序 const createTask = (detail: Record = {}) => { const { dataset } = detail; const title = `上传数据集: ${dataset.name} `; const controller = new AbortController(); const task: TaskItem = { key: dataset.id, title, percent: 0, reqId: -1, controller, size: 0, updateEvent: detail.updateEvent, hasArchive: detail.hasArchive, prefix: detail.prefix, }; taskListRef.current = [task, ...taskListRef.current]; setTaskList(taskListRef.current); // 立即显示任务中心,让用户感知上传已开始 if (showTaskCenter) { window.dispatchEvent( new CustomEvent("show:task-popover", { detail: { show: true } }) ); } return task; }; const updateTaskList = (task: TaskItem) => { taskListRef.current = taskListRef.current.map((item) => item.key === task.key ? task : item ); setTaskList(taskListRef.current); }; const removeTask = (task: TaskItem) => { const { key } = task; taskListRef.current = taskListRef.current.filter( (item) => item.key !== key ); setTaskList(taskListRef.current); if (task.isCancel && task.cancelFn) { task.cancelFn(); } if (task.updateEvent) { // 携带前缀信息,便于刷新后仍停留在当前目录 window.dispatchEvent( new CustomEvent(task.updateEvent, { detail: { prefix: task.prefix }, }) ); } if (showTaskCenter) { window.dispatchEvent( new CustomEvent("show:task-popover", { detail: { show: false } }) ); } }; async function buildFormData({ file, reqId, i, j }: { file: { slices: Blob[]; name: string; size: number }; reqId: number; i: number; j: number }) { const formData = new FormData(); const { slices, name, size } = file; const checkSum = await calculateSHA256(slices[j]); formData.append("file", slices[j]); formData.append("reqId", reqId.toString()); formData.append("fileNo", (i + 1).toString()); formData.append("chunkNo", (j + 1).toString()); formData.append("fileName", name); formData.append("fileSize", size.toString()); formData.append("totalChunkNum", slices.length.toString()); formData.append("checkSumHex", checkSum); return formData; } async function uploadSlice(task: TaskItem, fileInfo: { loaded: number; i: number; j: number; files: { slices: Blob[]; name: string; size: number }[]; totalSize: number }) { if (!task) { return; } const { reqId, key, controller } = task; const { loaded, i, j, files, totalSize } = fileInfo; // 检查是否已取消 if (controller.signal.aborted) { throw new Error("Upload cancelled"); } const formData = await buildFormData({ file: files[i], i, j, reqId, }); let newTask = { ...task }; await uploadChunk(key, formData, { signal: controller.signal, onUploadProgress: (e) => { const loadedSize = loaded + e.loaded; const curPercent = Number((loadedSize / totalSize) * 100).toFixed(2); newTask = { ...newTask, ...taskListRef.current.find((item) => item.key === key), size: loadedSize, percent: curPercent >= 100 ? 99.99 : curPercent, }; updateTaskList(newTask); }, }); } async function uploadFile({ task, files, totalSize }: { task: TaskItem; files: { slices: Blob[]; name: string; size: number; originFile: Blob }[]; totalSize: number }) { console.log('[useSliceUpload] Calling preUpload with prefix:', task.prefix); const { data: reqId } = await preUpload(task.key, { totalFileNum: files.length, totalSize, datasetId: task.key, hasArchive: task.hasArchive, prefix: task.prefix, }); console.log('[useSliceUpload] PreUpload response reqId:', reqId); const newTask: TaskItem = { ...task, reqId, isCancel: false, cancelFn: () => { // 使用 newTask 的 controller 确保一致性 newTask.controller.abort(); cancelUpload?.(reqId); if (newTask.updateEvent) window.dispatchEvent(new Event(newTask.updateEvent)); }, }; updateTaskList(newTask); // 注意:show:task-popover 事件已在 createTask 中触发,此处不再重复触发 // // 更新数据状态 if (task.updateEvent) window.dispatchEvent(new Event(task.updateEvent)); let loaded = 0; for (let i = 0; i < files.length; i++) { // 检查是否已取消 if (newTask.controller.signal.aborted) { throw new Error("Upload cancelled"); } const { slices } = files[i]; for (let j = 0; j < slices.length; j++) { // 检查是否已取消 if (newTask.controller.signal.aborted) { throw new Error("Upload cancelled"); } await uploadSlice(newTask, { loaded, i, j, files, totalSize, }); loaded += slices[j].size; } } removeTask(newTask); } const handleUpload = async ({ task, files }: { task: TaskItem; files: { slices: Blob[]; name: string; size: number; originFile: Blob }[] }) => { const isErrorFile = await checkIsFilesExist(files); if (isErrorFile) { message.error("文件被修改或删除,请重新选择文件上传"); removeTask({ ...task, isCancel: false, ...taskListRef.current.find((item) => item.key === task.key), }); return; } try { const totalSize = files.reduce((acc, file) => acc + file.size, 0); await uploadFile({ task, files, totalSize }); } catch (err) { console.error(err); message.error("文件上传失败,请稍后重试"); removeTask({ ...task, isCancel: true, ...taskListRef.current.find((item) => item.key === task.key), }); } }; /** * 流式分割上传处理 * 用于大文件按行分割并立即上传的场景 */ const handleStreamUpload = async ({ task, files }: { task: TaskItem; files: File[] }) => { try { console.log('[useSliceUpload] Starting stream upload for', files.length, 'files'); const totalSize = files.reduce((acc, file) => acc + file.size, 0); // 存储所有文件的 reqId,用于取消上传 const reqIds: number[] = []; const newTask: TaskItem = { ...task, reqId: -1, isCancel: false, cancelFn: () => { // 使用 newTask 的 controller 确保一致性 newTask.controller.abort(); // 取消所有文件的预上传请求 reqIds.forEach(id => cancelUpload?.(id)); if (newTask.updateEvent) window.dispatchEvent(new Event(newTask.updateEvent)); }, }; updateTaskList(newTask); let totalUploadedLines = 0; let totalProcessedBytes = 0; const results: StreamUploadResult[] = []; // 逐个处理文件,每个文件单独调用 preUpload for (let i = 0; i < files.length; i++) { // 检查是否已取消 if (newTask.controller.signal.aborted) { throw new Error("Upload cancelled"); } const file = files[i]; console.log(`[useSliceUpload] Processing file ${i + 1}/${files.length}: ${file.name}`); const result = await streamSplitAndUpload( file, (formData, config) => uploadChunk(task.key, formData, { ...config, signal: newTask.controller.signal, }), (currentBytes, totalBytes, uploadedLines) => { // 检查是否已取消 if (newTask.controller.signal.aborted) { return; } // 更新进度 const overallBytes = totalProcessedBytes + currentBytes; const curPercent = Number((overallBytes / totalSize) * 100).toFixed(2); const updatedTask: TaskItem = { ...newTask, ...taskListRef.current.find((item) => item.key === task.key), size: overallBytes, percent: curPercent >= 100 ? 99.99 : curPercent, streamUploadInfo: { currentFile: file.name, fileIndex: i + 1, totalFiles: files.length, uploadedLines: totalUploadedLines + uploadedLines, }, }; updateTaskList(updatedTask); }, 1024 * 1024, // 1MB chunk size { resolveReqId: async ({ totalFileNum, totalSize }) => { const { data: reqId } = await preUpload(task.key, { totalFileNum, totalSize, datasetId: task.key, hasArchive: task.hasArchive, prefix: task.prefix, }); console.log(`[useSliceUpload] File ${file.name} preUpload response reqId:`, reqId); reqIds.push(reqId); return reqId; }, hasArchive: newTask.hasArchive, prefix: newTask.prefix, signal: newTask.controller.signal, maxConcurrency: 3, } ); results.push(result); totalUploadedLines += result.uploadedCount; totalProcessedBytes += file.size; console.log(`[useSliceUpload] File ${file.name} processed, uploaded ${result.uploadedCount} lines`); } console.log('[useSliceUpload] Stream upload completed, total lines:', totalUploadedLines); removeTask(newTask); message.success(`成功上传 ${totalUploadedLines} 个文件(按行分割)`); } catch (err) { console.error('[useSliceUpload] Stream upload error:', err); if (err.message === "Upload cancelled") { message.info("上传已取消"); } else { message.error("文件上传失败,请稍后重试"); } removeTask({ ...task, isCancel: true, ...taskListRef.current.find((item) => item.key === task.key), }); } }; /** * 注册流式上传事件监听 * 返回注销函数 */ const registerStreamUploadListener = () => { if (!enableStreamUpload) return () => {}; const streamUploadHandler = async (e: Event) => { const customEvent = e as CustomEvent; const { dataset, files, updateEvent, hasArchive, prefix } = customEvent.detail; const controller = new AbortController(); const task: TaskItem = { key: dataset.id, title: `上传数据集: ${dataset.name} (按行分割)`, percent: 0, reqId: -1, controller, size: 0, updateEvent, hasArchive, prefix, }; taskListRef.current = [task, ...taskListRef.current]; setTaskList(taskListRef.current); // 显示任务中心 if (showTaskCenter) { window.dispatchEvent( new CustomEvent("show:task-popover", { detail: { show: true } }) ); } await handleStreamUpload({ task, files }); }; window.addEventListener("upload:dataset-stream", streamUploadHandler); return () => { window.removeEventListener("upload:dataset-stream", streamUploadHandler); }; }; return { taskList, createTask, removeTask, handleUpload, handleStreamUpload, registerStreamUploadListener, }; }