From 0e023fbf0329edd9abd0057730e2faf33b06dbc8 Mon Sep 17 00:00:00 2001
From: heavyscientist <github@mainloot.com>
Date: Mon, 17 Feb 2025 18:45:27 -0800
Subject: [PATCH] retry file downloads 5 times, exit 1 on failure

---
 cmd/rustmaps/generate.go |  3 +-
 pkg/rustmaps/download.go | 88 ++++++++++++++++++++++++++++------------
 2 files changed, 64 insertions(+), 27 deletions(-)

diff --git a/cmd/rustmaps/generate.go b/cmd/rustmaps/generate.go
index 5782b12..1ea2815 100644
--- a/cmd/rustmaps/generate.go
+++ b/cmd/rustmaps/generate.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 	"time"
 
@@ -42,7 +43,7 @@ var generateCmd = &cobra.Command{
 			version := now.Format("2006-01-02_15-04-05")
 			if err := generator.Download(logger, version); err != nil {
 				fmt.Printf("Error downloading maps: %v\n", err)
-				return
+				os.Exit(1)
 			}
 
 			fmt.Printf("Maps downloaded to %s\n", filepath.Join(generator.GetDownloadsDir(), version))
diff --git a/pkg/rustmaps/download.go b/pkg/rustmaps/download.go
index 948a317..3ae463d 100644
--- a/pkg/rustmaps/download.go
+++ b/pkg/rustmaps/download.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -29,43 +30,78 @@ func (g *Generator) OverrideDownloadsDir(log *zap.Logger, dir string) {
 
 // DownloadFile downloads a file using net/http
 func (g *Generator) DownloadFile(log *zap.Logger, url, target string) error {
+	maxRetries := 3
+	backoff := 5 * time.Second
+
 	client := &http.Client{
 		Timeout: 10 * time.Second,
 	}
 
-	req, err := http.NewRequest("GET", url, nil)
-	if err != nil {
-		log.Error("Error creating request", zap.Error(err))
-		return err
-	}
+	var lastErr error
+	for attempt := 0; attempt <= maxRetries; attempt++ {
+		if attempt > 0 {
+			sleepDuration := backoff * time.Duration(math.Pow(2, float64(attempt-1)))
+			log.Info("Retrying download",
+				zap.Int("attempt", attempt),
+				zap.Duration("backoff", sleepDuration))
+			time.Sleep(sleepDuration)
+		}
 
-	resp, err := client.Do(req)
-	if err != nil {
-		log.Error("Error downloading file", zap.Error(err))
-		return err
-	}
-	defer resp.Body.Close()
+		req, err := http.NewRequest("GET", url, nil)
+		if err != nil {
+			lastErr = err
+			log.Error("Error creating request", zap.Error(err))
+			continue
+		}
 
-	if resp.StatusCode != http.StatusOK {
-		log.Error("Error downloading file", zap.String("status", resp.Status))
-		return fmt.Errorf("error downloading file: %s", resp.Status)
-	}
+		resp, err := client.Do(req)
+		if err != nil {
+			lastErr = err
+			log.Error("Error downloading file",
+				zap.Error(err),
+				zap.Int("attempt", attempt))
+			continue
+		}
 
-	file, err := os.Create(target)
-	if err != nil {
-		log.Error("Error creating file", zap.Error(err))
-		return err
-	}
-	defer file.Close()
+		// Always close response body, but keep error for checking
+		defer resp.Body.Close()
 
-	_, err = io.Copy(file, resp.Body)
-	if err != nil {
-		log.Error("Error writing file", zap.Error(err))
-		return err
+		if resp.StatusCode != http.StatusOK {
+			lastErr = fmt.Errorf("error downloading file: %s", resp.Status)
+			log.Error("Error downloading file",
+				zap.String("status", resp.Status),
+				zap.Int("attempt", attempt))
+			continue
+		}
+
+		file, err := os.Create(target)
+		if err != nil {
+			lastErr = err
+			log.Error("Error creating file", zap.Error(err))
+			return err // Don't retry file creation errors
+		}
+		defer file.Close()
+
+		_, err = io.Copy(file, resp.Body)
+		if err != nil {
+			lastErr = err
+			log.Error("Error writing file",
+				zap.Error(err),
+				zap.Int("attempt", attempt))
+			continue
+		}
+
+		// If we get here, the download was successful
+		log.Info("File downloaded successfully",
+			zap.String("url", url),
+			zap.String("target", target),
+			zap.Int("attempts", attempt+1))
+		return nil
 	}
 
-	return nil
+	return fmt.Errorf("failed after %d attempts, last error: %v", maxRetries, lastErr)
 }
+
 func (g *Generator) Download(log *zap.Logger, version string) error {
 	if len(g.maps) == 0 {
 		log.Warn("No maps loaded")