From befcadd260059af88092391051fe27056888c44d Mon Sep 17 00:00:00 2001 From: Jason Ross Date: Thu, 4 Jun 2026 17:10:57 -0500 Subject: [PATCH] adding aria2 for faster downloads --- net.go | 26 ++++++++++- net_test.go | 121 +++++++++++++++++++++++++++++++++++++++++++++++++ packages.go | 2 +- system.go | 32 +++++++++++++ system_test.go | 33 ++++++++++++++ 5 files changed, 212 insertions(+), 2 deletions(-) diff --git a/net.go b/net.go index 70ac52b..d3cb0ab 100644 --- a/net.go +++ b/net.go @@ -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)) diff --git a/net_test.go b/net_test.go index 0719d52..f2b5c3c 100644 --- a/net_test.go +++ b/net_test.go @@ -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)) + } +} + diff --git a/packages.go b/packages.go index 284a0f2..ad27091 100644 --- a/packages.go +++ b/packages.go @@ -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", diff --git a/system.go b/system.go index 4705ae0..d3af1d3 100644 --- a/system.go +++ b/system.go @@ -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 { diff --git a/system_test.go b/system_test.go index e132449..4b9a5ed 100644 --- a/system_test.go +++ b/system_test.go @@ -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) + } +} +