Merge pull request #25 from JMR-dev/feat-add-aria2

adding aria2 for faster downloads
2026-06-04 17:42:05 -05:00
parent 60cf195f45 befcadd260
commit 07c212d962
5 changed files with 212 additions and 2 deletions
@@ -19,7 +19,31 @@ var httpClient = &http.Client{Timeout: httpClientTimeout}

 // downloadReal streams url -> dest. Returns true on success.
 func downloadReal(url, dest string) bool {
-	fmt.Printf("  Downloading %s ...\n", filepath.Base(url))
+	taskPrintf("  Downloading %s ...\n", filepath.Base(url))
+
+	dir, file := filepath.Split(dest)
+	dir = filepath.Clean(dir)
+
+	// Try aria2c first if available
+	if hasCmd("aria2c") {
+		res := runCmd([]string{"aria2c", "-x", "16", "-s", "16", "-k", "1M", "-d", dir, "-o", file, url}, CmdOpts{Out: taskOut()})
+		if res.OK() {
+			return true
+		}
+		taskPrintf("  [WARN] aria2c download failed for %s, falling back to curl ...\n", url)
+	}
+
+	// Fallback to curl
+	if hasCmd("curl") {
+		res := runCmd([]string{"curl", "-L", "--fail", "-o", dest, url}, CmdOpts{Out: taskOut()})
+		if res.OK() {
+			return true
+		}
+		taskPrintf("  [WARN] curl download failed for %s ...\n", url)
+	}
+
+	// Final fallback: Go built-in HTTP client
+	taskPrintf("  Falling back to built-in HTTP client for %s ...\n", url)
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {
 		errLog(fmt.Sprintf("Download failed for %s: %v", url, err))
@@ -236,3 +236,124 @@ func TestNetRealErrors(t *testing.T) {
 	}
 }

+func TestDownloadRealFallbackChain(t *testing.T) {
+	defer resetMocks()
+
+	tmpDir := t.TempDir()
+	destFile := filepath.Join(tmpDir, "out.txt")
+
+	// Case 1: aria2c works
+	var aria2cCalled bool
+	var curlCalled bool
+	hasCmd = func(name string) bool {
+		if name == "aria2c" || name == "curl" {
+			return true
+		}
+		return false
+	}
+	runCmd = func(argv []string, opts CmdOpts) CmdResult {
+		if argv[0] == "aria2c" {
+			aria2cCalled = true
+			_ = os.WriteFile(destFile, []byte("aria2c content"), 0644)
+			return CmdResult{ExitCode: 0}
+		}
+		if argv[0] == "curl" {
+			curlCalled = true
+			return CmdResult{ExitCode: 0}
+		}
+		return CmdResult{ExitCode: 1}
+	}
+
+	success := downloadReal("https://example.com/file", destFile)
+	if !success {
+		t.Fatal("expected download via aria2c to succeed")
+	}
+	if !aria2cCalled {
+		t.Error("expected aria2c to be called")
+	}
+	if curlCalled {
+		t.Error("expected curl NOT to be called when aria2c succeeds")
+	}
+
+	// Case 2: aria2c fails, falls back to curl, curl succeeds
+	resetMocks()
+	aria2cCalled = false
+	curlCalled = false
+	hasCmd = func(name string) bool {
+		if name == "aria2c" || name == "curl" {
+			return true
+		}
+		return false
+	}
+	runCmd = func(argv []string, opts CmdOpts) CmdResult {
+		if argv[0] == "aria2c" {
+			aria2cCalled = true
+			return CmdResult{ExitCode: 1, Err: fmt.Errorf("aria2c simulated error")}
+		}
+		if argv[0] == "curl" {
+			curlCalled = true
+			_ = os.WriteFile(destFile, []byte("curl content"), 0644)
+			return CmdResult{ExitCode: 0}
+		}
+		return CmdResult{ExitCode: 1}
+	}
+
+	success = downloadReal("https://example.com/file", destFile)
+	if !success {
+		t.Fatal("expected download to succeed via curl fallback")
+	}
+	if !aria2cCalled {
+		t.Error("expected aria2c to be attempted")
+	}
+	if !curlCalled {
+		t.Error("expected curl to be attempted after aria2c failed")
+	}
+
+	// Case 3: aria2c fails, curl fails, falls back to Go HTTP client
+	resetMocks()
+	aria2cCalled = false
+	curlCalled = false
+	hasCmd = func(name string) bool {
+		if name == "aria2c" || name == "curl" {
+			return true
+		}
+		return false
+	}
+	runCmd = func(argv []string, opts CmdOpts) CmdResult {
+		if argv[0] == "aria2c" {
+			aria2cCalled = true
+			return CmdResult{ExitCode: 1, Err: fmt.Errorf("aria2c simulated error")}
+		}
+		if argv[0] == "curl" {
+			curlCalled = true
+			return CmdResult{ExitCode: 1, Err: fmt.Errorf("curl simulated error")}
+		}
+		return CmdResult{ExitCode: 1}
+	}
+	oldTransport := httpClient.Transport
+	defer func() { httpClient.Transport = oldTransport }()
+	httpClient.Transport = &mockTripper{
+		roundTripFunc: func(req *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(bytes.NewBufferString("go http content")),
+			}, nil
+		},
+	}
+
+	success = downloadReal("https://example.com/file", destFile)
+	if !success {
+		t.Fatal("expected download to succeed via Go http fallback")
+	}
+	if !aria2cCalled {
+		t.Error("expected aria2c to be attempted")
+	}
+	if !curlCalled {
+		t.Error("expected curl to be attempted")
+	}
+	data, _ := os.ReadFile(destFile)
+	if string(data) != "go http content" {
+		t.Errorf("expected file content to be 'go http content', got %q", string(data))
+	}
+}
+
@@ -10,10 +10,10 @@ package main
 // URL templates use the substitutions described in formatURL.

 var SystemPackages = []string{
+	"aria2",
 	"age",
 	"ansible",
 	"ansible-core",
-	"aria2",
 	"btm",
 	"build-essential",
 	"buildah",
@@ -458,6 +458,38 @@ func brewInstallMany(pkgs []string) (failed []string) {
 func installSystemPackages(regular, special []string) {
 	fmt.Println("\n=== System Packages ===")

+	// Ensure aria2 is installed first and on the system path
+	var installAria2 bool
+	var remainingRegular []string
+	for _, p := range regular {
+		if p == "aria2" {
+			installAria2 = true
+		} else {
+			remainingRegular = append(remainingRegular, p)
+		}
+	}
+
+	if installAria2 || !hasCmd("aria2c") {
+		fmt.Println("  Ensuring aria2 is installed first and on the system path ...")
+		var res CmdResult
+		switch pkgMgr {
+		case "brew":
+			res = runCmd([]string{"brew", "install", "aria2"}, CmdOpts{})
+		case "pacman":
+			res = runCmd([]string{"pacman", "-S", "--noconfirm", "--needed", "aria2"}, CmdOpts{AsSudo: true})
+		default: // dnf, apt-get
+			res = runCmd([]string{pkgMgr, "install", "-y", "aria2"}, CmdOpts{AsSudo: true})
+		}
+		if !res.OK() {
+			warn(fmt.Sprintf("Failed to install aria2: %v", res.Err))
+		} else if !hasCmd("aria2c") {
+			warn("aria2 was installed but 'aria2c' is not found on the system path")
+		} else {
+			fmt.Println("  aria2 is installed and on the system path.")
+		}
+		regular = remainingRegular
+	}
+
 	if pkgMgr == "brew" {
 		failed := pkgInstallMany(regular)
 		for _, p := range failed {
@@ -495,3 +495,36 @@ func TestSystemGoEdgeCases(t *testing.T) {
 	installSystemPackages([]string{"docker-ce"}, []string{})
 }

+func TestInstallSystemPackagesAria2First(t *testing.T) {
+	defer resetMocks()
+
+	pkgMgr = "dnf"
+	var runCmdCalls [][]string
+	runCmd = func(argv []string, opts CmdOpts) CmdResult {
+		runCmdCalls = append(runCmdCalls, argv)
+		return CmdResult{ExitCode: 0}
+	}
+	hasCmd = func(name string) bool {
+		if name == "aria2c" {
+			return false
+		}
+		return true
+	}
+
+	installSystemPackages([]string{"git", "aria2", "tmux"}, []string{})
+
+	if len(runCmdCalls) < 2 {
+		t.Fatalf("expected at least 2 command calls, got %d: %v", len(runCmdCalls), runCmdCalls)
+	}
+
+	firstCall := runCmdCalls[0]
+	if len(firstCall) < 4 || firstCall[0] != "dnf" || firstCall[1] != "install" || firstCall[3] != "aria2" {
+		t.Errorf("expected first call to be installing aria2, got: %v", firstCall)
+	}
+
+	secondCall := runCmdCalls[1]
+	if len(secondCall) < 5 || secondCall[0] != "dnf" || secondCall[1] != "install" || secondCall[3] != "git" || secondCall[4] != "tmux" {
+		t.Errorf("expected second call to install remaining packages, got: %v", secondCall)
+	}
+}
+