From 9edc756d297fbcb2afbb6c18de146fd01a6d19e7 Mon Sep 17 00:00:00 2001
From: Ram Prasad Agarwal <work.ramprasad@gmail.com>
Date: Thu, 30 Jan 2025 10:25:48 +0530
Subject: [PATCH] [ui-storagebrowser] fixes chunk upload for new storage
 browser (#3977)

---
 .../hooks/useFileUpload/useChunkUpload.ts     | 89 ++++++++++++-------
 1 file changed, 55 insertions(+), 34 deletions(-)

diff --git a/desktop/core/src/desktop/js/utils/hooks/useFileUpload/useChunkUpload.ts b/desktop/core/src/desktop/js/utils/hooks/useFileUpload/useChunkUpload.ts
index 1a4b478258d..704c9b576eb 100644
--- a/desktop/core/src/desktop/js/utils/hooks/useFileUpload/useChunkUpload.ts
+++ b/desktop/core/src/desktop/js/utils/hooks/useFileUpload/useChunkUpload.ts
@@ -57,21 +57,34 @@ const useChunkUpload = ({
 }: ChunkUploadOptions): UseUploadQueueResponse => {
   const config = getLastKnownConfig();
   const chunkSize = config?.storage_browser?.file_upload_chunk_size ?? DEFAULT_CHUNK_SIZE;
-  const [pendingItems, setPendingItems] = useState<UploadItem[]>([]);
+  const [processingItem, setProcessingItem] = useState<UploadItem>();
+  const [pendingUploadItems, setPendingUploadItems] = useState<UploadItem[]>([]);
+  const [awaitingStatusItems, setAwaitingStatusItems] = useState<UploadItem[]>([]);
+
+  const onError = () => {
+    if (processingItem) {
+      onStatusUpdate(processingItem, FileUploadStatus.Failed);
+      setProcessingItem(undefined);
+    }
+  };
+
+  const onSuccess = (item: UploadItem) => () => {
+    setAwaitingStatusItems(prev => [...prev, item]);
+    setProcessingItem(undefined);
+  };
 
   const { save } = useSaveData(undefined, {
     postOptions: {
       qsEncodeData: false,
-      headers: {
-        'Content-Type': 'multipart/form-data'
-      }
-    }
+      headers: { 'Content-Type': 'multipart/form-data' }
+    },
+    onError
   });
 
   const updateItemStatus = (serverResponse: TaskServerResponse[]) => {
     const statusMap = getStatusHashMap(serverResponse);
 
-    const remainingItems = pendingItems.filter(item => {
+    const remainingItems = awaitingStatusItems.filter(item => {
       const status = statusMap[item.uuid];
       if (status === TaskStatus.Success || status === TaskStatus.Failure) {
         const ItemStatus =
@@ -84,14 +97,14 @@ const useChunkUpload = ({
     if (remainingItems.length === 0) {
       onComplete();
     }
-    setPendingItems(remainingItems);
+    setAwaitingStatusItems(remainingItems);
   };
 
   const { data: tasksStatus } = useLoadData<TaskServerResponse[]>(
     '/desktop/api2/taskserver/get_taskserver_tasks/',
     {
-      pollInterval: pendingItems.length ? 5000 : undefined,
-      skip: !pendingItems.length
+      pollInterval: awaitingStatusItems.length ? 5000 : undefined,
+      skip: !awaitingStatusItems.length
     }
   );
 
@@ -101,16 +114,14 @@ const useChunkUpload = ({
     }
   }, [tasksStatus]);
 
-  const addItemToPending = (item: UploadItem) => {
-    setPendingItems(prev => [...prev, item]);
-  };
-
-  const onChunksUploadComplete = async (item: UploadItem) => {
-    const { url, payload } = getChunksCompletePayload(item, chunkSize);
-    return save(payload, {
-      url,
-      onSuccess: () => addItemToPending(item)
-    });
+  const onChunksUploadComplete = async () => {
+    if (processingItem) {
+      const { url, payload } = getChunksCompletePayload(processingItem, chunkSize);
+      return save(payload, {
+        url,
+        onSuccess: onSuccess(processingItem)
+      });
+    }
   };
 
   const uploadChunk = async (chunkItem: UploadChunkItem) => {
@@ -118,21 +129,21 @@ const useChunkUpload = ({
     return save(payload, { url });
   };
 
-  const { enqueueAsync } = useQueueProcessor<UploadChunkItem>(uploadChunk, {
-    concurrentProcess
+  const { enqueue } = useQueueProcessor<UploadChunkItem>(uploadChunk, {
+    concurrentProcess,
+    onSuccess: onChunksUploadComplete
   });
 
-  const uploadItemInChunks = async (item: UploadItem) => {
+  const uploadItemInChunks = (item: UploadItem) => {
     const chunks = createChunks(item, chunkSize);
-    await enqueueAsync(chunks);
-    return onChunksUploadComplete(item);
+    return enqueue(chunks);
   };
 
-  const uploadItemInSingleChunk = (item: UploadItem) => {
+  const uploadItemInSingleChunk = async (item: UploadItem) => {
     const { url, payload } = getChunkSinglePayload(item, chunkSize);
     return save(payload, {
       url,
-      onSuccess: () => addItemToPending(item)
+      onSuccess: onSuccess(item)
     });
   };
 
@@ -158,15 +169,25 @@ const useChunkUpload = ({
     return uploadItemInChunks(item);
   };
 
-  const {
-    enqueue: addFiles,
-    dequeue: removeFile,
-    isLoading
-  } = useQueueProcessor<UploadItem>(uploadItem, {
-    concurrentProcess: 1 // This value must be 1 always
-  });
+  const addFiles = (newItems: UploadItem[]) => {
+    setPendingUploadItems(prev => [...prev, ...newItems]);
+  };
+
+  const removeFile = (item: UploadItem) => {
+    setPendingUploadItems(prev => prev.filter(i => i.uuid !== item.uuid));
+  };
+
+  useEffect(() => {
+    // Ensures one file is broken down in chunks and uploaded to the server
+    if (!processingItem && pendingUploadItems.length) {
+      const item = pendingUploadItems[0];
+      setProcessingItem(item);
+      setPendingUploadItems(prev => prev.slice(1));
+      uploadItem(item);
+    }
+  }, [pendingUploadItems, processingItem]);
 
-  return { addFiles, removeFile, isLoading };
+  return { addFiles, removeFile, isLoading: !!processingItem || !!pendingUploadItems.length };
 };
 
 export default useChunkUpload;