Skip to content

Commit 67e6e89

Browse files
committed
feat: Implement download rate limiting for incremental updates
Implement download rate limiting for incremental downloads and updates using the trickle tool.
1 parent d85de2f commit 67e6e89

1 file changed

Lines changed: 76 additions & 4 deletions

File tree

src/internal/system/apt/apt.go

Lines changed: 76 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"fmt"
1010
"os"
1111
"os/exec"
12+
"regexp"
1213
"strconv"
1314
"strings"
1415
"syscall"
@@ -41,6 +42,79 @@ func (p *APTSystem) FindCMD(id string) *system.Command {
4142
return p.CmdSet[id]
4243
}
4344

45+
// extractDownloadLimit extracts the value of Acquire::http::Dl-Limit from the command arguments.
46+
// Example input: []string{"package1", "package2", "-o", "Acquire::http::Dl-Limit=300", "-o", "Dir::Etc::SourceList=/dev/null"}
47+
// Returns: The extracted limit (KB/s), or 0 if not found.
48+
func extractDownloadLimit(cmdArgs []string) int {
49+
// Use a regular expression to match Acquire::http::Dl-Limit=<number>
50+
dlLimitRegex := regexp.MustCompile(`^Acquire::http::Dl-Limit=(\d+)$`)
51+
52+
for i, arg := range cmdArgs {
53+
// Look for the "-o" option
54+
if arg == "-o" && i+1 < len(cmdArgs) {
55+
nextArg := cmdArgs[i+1]
56+
// Check if the next argument matches the Dl-Limit pattern
57+
if matches := dlLimitRegex.FindStringSubmatch(nextArg); len(matches) > 1 {
58+
if limit, err := strconv.Atoi(matches[1]); err == nil {
59+
return limit
60+
}
61+
}
62+
}
63+
}
64+
65+
return 0
66+
}
67+
68+
// createIncrementalCommandLine creates the command line for incremental update/download operations.
69+
// It supports extracting download limit from cmdArgs and uses the trickle tool for bandwidth control if possible.
70+
func createIncrementalCommandLine(cmdType string, cmdArgs []string) *exec.Cmd {
71+
// Extract download limit setting
72+
downloadLimit := extractDownloadLimit(cmdArgs)
73+
logger.Debugf("Download limit extracted: %d KB/s", downloadLimit)
74+
75+
// Build base command arguments
76+
baseArgs := []string{"upgrade", "--status-fd", "3"}
77+
if cmdType == system.IncrementalDownloadJobType {
78+
baseArgs = append(baseArgs, "--download-only")
79+
}
80+
81+
// Apply download limit with trickle if set and available
82+
if downloadLimit > 0 {
83+
if cmd := buildTrickleCommand(downloadLimit, system.DeepinImmutableCtlPath, baseArgs); cmd != nil {
84+
return cmd
85+
}
86+
}
87+
88+
// Default command (no bandwidth restriction)
89+
return exec.Command(system.DeepinImmutableCtlPath, baseArgs...)
90+
}
91+
92+
// buildTrickleCommand constructs a command using trickle for bandwidth control.
93+
// Returns nil if trickle is unavailable, so that the caller can fall back to the default command.
94+
func buildTrickleCommand(downloadLimit int, targetCmd string, targetArgs []string) *exec.Cmd {
95+
tricklePath, err := exec.LookPath("trickle")
96+
if err != nil {
97+
logger.Warningf("trickle binary not found, will not apply download limit: %v", err)
98+
return nil
99+
}
100+
101+
// Run a small test: trickle -s -d 100 echo.
102+
testCmd := exec.Command(tricklePath, "-s", "-d", "100", "echo")
103+
err = testCmd.Run()
104+
if err != nil {
105+
logger.Warningf("trickle test failed, will not apply download limit: %v", err)
106+
return nil
107+
}
108+
109+
// Build the trickle command: trickle -s -d <limit> <target_cmd> <target_args...>
110+
trickleArgs := []string{"-s", "-d", strconv.Itoa(downloadLimit), targetCmd}
111+
trickleArgs = append(trickleArgs, targetArgs...)
112+
trickleArgs = append(trickleArgs, "--limit")
113+
114+
logger.Debugf("Applying download limit with trickle: %d KB/s", downloadLimit)
115+
return exec.Command(tricklePath, trickleArgs...)
116+
}
117+
44118
func createCommandLine(cmdType string, cmdArgs []string) *exec.Cmd {
45119
var args = []string{"-y"}
46120

@@ -87,10 +161,8 @@ func createCommandLine(cmdType string, cmdArgs []string) *exec.Cmd {
87161
return exec.Command("/usr/bin/lastore-apt-clean")
88162
case system.BackupJobType:
89163
return exec.Command(system.DeepinImmutableCtlPath, "admin", "deploy", "--backup", "-j")
90-
case system.IncrementalDownloadJobType:
91-
return exec.Command(system.DeepinImmutableCtlPath, "upgrade", "--download-only", "--status-fd", "3")
92-
case system.IncrementalUpdateJobType:
93-
return exec.Command(system.DeepinImmutableCtlPath, "upgrade", "--status-fd", "3")
164+
case system.IncrementalDownloadJobType, system.IncrementalUpdateJobType:
165+
return createIncrementalCommandLine(cmdType, cmdArgs)
94166
case system.FixErrorJobType:
95167
var errType system.JobErrorType
96168
if len(cmdArgs) >= 1 {

0 commit comments

Comments
 (0)