From b031be4ecda283c920c8b42b38f89100bb6b2b50 Mon Sep 17 00:00:00 2001
From: Xpl0itU <24777100+Xpl0itU@users.noreply.github.com>
Date: Fri, 29 Mar 2024 16:32:19 +0100
Subject: [PATCH] Optimize download logic

---
 downloader.go | 74 ++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 53 insertions(+), 21 deletions(-)

diff --git a/downloader.go b/downloader.go
index 113ccfb..daab52d 100644
--- a/downloader.go
+++ b/downloader.go
@@ -41,6 +41,32 @@ func calculateDownloadSpeed(downloaded int64, startTime, endTime time.Time) int6
 	return 0
 }
 
+func downloadChunk(buffer []byte, resp *http.Response, file *os.File, doRetries bool, attempt int) (int64, error) {
+	n, err := resp.Body.Read(buffer)
+	if err != nil && err != io.EOF {
+		if doRetries && attempt < maxRetries {
+			time.Sleep(retryDelay)
+			return 0, nil
+		}
+		return 0, fmt.Errorf("download error after %d attempts: %+v", attempt, err)
+	}
+
+	if n == 0 {
+		return 0, nil
+	}
+
+	_, err = file.Write(buffer[:n])
+	if err != nil {
+		if doRetries && attempt < maxRetries {
+			time.Sleep(retryDelay)
+			return 0, nil
+		}
+		return 0, fmt.Errorf("write error after %d attempts: %+v", attempt, err)
+	}
+
+	return int64(n), nil
+}
+
 func downloadFile(ctx context.Context, progressReporter ProgressReporter, client *http.Client, downloadURL, dstPath string, doRetries bool, buffer []byte) error {
 	filePath := filepath.Base(dstPath)
 
@@ -75,33 +101,39 @@ func downloadFile(ctx context.Context, progressReporter ProgressReporter, client
 		var downloaded int64
 
 		startTime = time.Now()
+		ticker := time.NewTicker(250 * time.Millisecond)
+		defer ticker.Stop()
+
+	Loop:
 		for {
-			n, err := resp.Body.Read(buffer)
-			if err != nil && err != io.EOF {
-				if doRetries && attempt < maxRetries {
-					time.Sleep(retryDelay)
-					break
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-ticker.C:
+				n, err := downloadChunk(buffer, resp, file, doRetries, attempt)
+				if err != nil {
+					return err
 				}
-				return fmt.Errorf("download error after %d attempts: %+v", attempt, err)
-			}
 
-			if n == 0 {
-				break
-			}
-
-			_, err = file.Write(buffer[:n])
-			if err != nil {
-				if doRetries && attempt < maxRetries {
-					time.Sleep(retryDelay)
-					break
+				if n == 0 {
+					break Loop
 				}
-				return fmt.Errorf("write error after %d attempts: %+v", attempt, err)
-			}
 
-			downloaded += int64(n)
-			progressReporter.UpdateDownloadProgress(downloaded, calculateDownloadSpeed(downloaded, startTime, time.Now()), filePath)
+				downloaded += n
+				progressReporter.UpdateDownloadProgress(downloaded, calculateDownloadSpeed(downloaded, startTime, time.Now()), filePath)
+			default:
+				n, err := downloadChunk(buffer, resp, file, doRetries, attempt)
+				if err != nil {
+					return err
+				}
+
+				if n == 0 {
+					break Loop
+				}
+
+				downloaded += n
+			}
 		}
-		break
 	}
 
 	return nil