From bb1edef473a294805fa2e5fcc530b422b8b006dd Mon Sep 17 00:00:00 2001
From: marknguyen1302 <nguyenvu1302.work@gmail.com>
Date: Tue, 30 Jul 2024 18:41:19 +0700
Subject: [PATCH] chore: handle failed download (#943)

---
 .../download-manager.service.ts               | 72 ++++++++++++-------
 cortex-js/src/utils/download-progress.ts      | 10 ++-
 2 files changed, 56 insertions(+), 26 deletions(-)

diff --git a/cortex-js/src/infrastructure/services/download-manager/download-manager.service.ts b/cortex-js/src/infrastructure/services/download-manager/download-manager.service.ts
index 64af5627e..cd4a3ebd0 100644
--- a/cortex-js/src/infrastructure/services/download-manager/download-manager.service.ts
+++ b/cortex-js/src/infrastructure/services/download-manager/download-manager.service.ts
@@ -136,7 +136,7 @@ export class DownloadManagerService {
         this.httpService.get(url, {
           responseType: 'stream',
           signal: controller.signal,
-        }),
+        })
       );
 
       // check if response is success
@@ -164,6 +164,24 @@ export class DownloadManagerService {
 
       console.log('Downloading', basename(destination));
 
+      const timeout = 20000; // Timeout period for receiving new data
+      let timeoutId: NodeJS.Timeout;
+      const resetTimeout = () => {
+        if (timeoutId) clearTimeout(timeoutId);
+        timeoutId = setTimeout(() => {
+          try{
+          this.handleError(
+            new Error('Download timeout'),
+            downloadId,
+            destination,
+          )
+        } finally {
+          bar.stop();
+          resolve();
+        }
+        }, timeout);
+      };
+
       let transferredBytes = 0;
       const bar = new SingleBar({}, Presets.shades_classic);
       bar.start(100, 0);
@@ -190,31 +208,9 @@ export class DownloadManagerService {
           resolve();
         }
       });
-
       writer.on('error', (error) => {
         try {
-          delete this.abortControllers[downloadId][destination];
-          const currentDownloadState = this.allDownloadStates.find(
-            (downloadState) => downloadState.id === downloadId,
-          );
-          if (!currentDownloadState) return;
-
-          const downloadItem = currentDownloadState?.children.find(
-            (downloadItem) => downloadItem.id === destination,
-          );
-          if (downloadItem) {
-            downloadItem.status = DownloadStatus.Error;
-            downloadItem.error = error.message;
-          }
-
-          currentDownloadState.status = DownloadStatus.Error;
-          currentDownloadState.error = error.message;
-
-          // remove download state if all children is downloaded
-          this.allDownloadStates = this.allDownloadStates.filter(
-            (downloadState) => downloadState.id !== downloadId,
-          );
-          this.eventEmitter.emit('download.event', this.allDownloadStates);
+          this.handleError(error, downloadId, destination);
         } finally {
           bar.stop();
           resolve();
@@ -222,6 +218,7 @@ export class DownloadManagerService {
       });
 
       response.data.on('data', (chunk: any) => {
+        resetTimeout();
         transferredBytes += chunk.length;
 
         const currentDownloadState = this.allDownloadStates.find(
@@ -266,4 +263,31 @@ export class DownloadManagerService {
   getDownloadStates() {
     return this.allDownloadStates;
   }
+
+  private handleError(error: Error, downloadId: string, destination: string) {
+    console.log(this.allDownloadStates, downloadId, destination)
+    delete this.abortControllers[downloadId][destination];
+          const currentDownloadState = this.allDownloadStates.find(
+            (downloadState) => downloadState.id === downloadId,
+          );
+          if (!currentDownloadState) return;
+
+          const downloadItem = currentDownloadState?.children.find(
+            (downloadItem) => downloadItem.id === destination,
+          );
+          if (downloadItem) {
+            downloadItem.status = DownloadStatus.Error;
+            downloadItem.error = error.message;
+          }
+
+          currentDownloadState.status = DownloadStatus.Error;
+          currentDownloadState.error = error.message;
+
+          // remove download state if all children is downloaded
+          this.allDownloadStates = this.allDownloadStates.filter(
+            (downloadState) => downloadState.id !== downloadId,
+          );
+          this.eventEmitter.emit('download.event', [currentDownloadState]);
+          this.eventEmitter.emit('download.event', this.allDownloadStates);
+  }
 }
diff --git a/cortex-js/src/utils/download-progress.ts b/cortex-js/src/utils/download-progress.ts
index 2341ac843..dd803f84f 100644
--- a/cortex-js/src/utils/download-progress.ts
+++ b/cortex-js/src/utils/download-progress.ts
@@ -1,7 +1,7 @@
 import { Presets, SingleBar } from "cli-progress";
 import { Cortex } from "@cortexso/cortex.js";
 import { exit, stdin, stdout } from 'node:process';
-import { DownloadState, DownloadType } from "@/domain/models/download.interface";
+import { DownloadState, DownloadStatus, DownloadType } from "@/domain/models/download.interface";
 
 export const downloadProgress = async (cortex: Cortex, downloadId?: string, downloadType?: DownloadType) => {
     const response = await cortex.events.downloadEvent();
@@ -31,7 +31,13 @@ export const downloadProgress = async (cortex: Cortex, downloadId?: string, down
         if (!data) continue;
         if (downloadType && data.type !== downloadType) continue;
 
-        if (data.status === 'downloaded') break;
+        if (data.status === DownloadStatus.Downloaded) break;
+        if(data.status === DownloadStatus.Error) {
+          rl.close();
+          progressBar.stop();
+          console.log('\n Download failed: ', data.error);
+          exit(1);
+        }
 
         let totalBytes = 0;
         let totalTransferred = 0;