Skip to content

Commit 04c5fe2

Browse files
DevinwongCopilot
andauthored
feat: implement budget timeout for apt_get_install (#8379)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 54aa84a commit 04c5fe2

2 files changed

Lines changed: 188 additions & 6 deletions

File tree

parts/linux/cloud-init/artifacts/ubuntu/cse_helpers_ubuntu.sh

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -54,15 +54,55 @@ _apt_get_install() {
5454
local retries=$1
5555
local wait_sleep=$2
5656
local apt_opts=$3
57-
shift && shift && shift
57+
local maxBudget=${4:-0}
58+
shift && shift && shift && shift
59+
local packages=("${@}")
60+
61+
local hasBudget=false
62+
if [ "${maxBudget}" -gt 0 ]; then
63+
hasBudget=true
64+
fi
65+
66+
local opStartTime
67+
opStartTime=$(date +%s)
5868

5969
for i in $(seq 1 $retries); do
70+
# CSE timeout guard
71+
if ! check_cse_timeout; then
72+
echo "CSE timeout approaching, exiting apt_get_install early." >&2
73+
return 2
74+
fi
75+
76+
# Per-operation budget check
77+
if [ "$hasBudget" = true ]; then
78+
local opElapsed
79+
opElapsed=$(( $(date +%s) - opStartTime ))
80+
if [ "$opElapsed" -ge "$maxBudget" ]; then
81+
echo "apt_get_install budget of ${maxBudget}s exceeded after ${opElapsed}s, exiting early." >&2
82+
return 2
83+
fi
84+
fi
85+
6086
wait_for_apt_locks
6187
export DEBIAN_FRONTEND=noninteractive
6288
dpkg --configure -a --force-confdef
6389

64-
if apt-get install ${apt_opts} -o Dpkg::Options::="--force-confold" --no-install-recommends "${@}"; then
65-
echo "Executed apt-get install \"${packages[@]}\" $i times"
90+
# Cap per-attempt timeout to the remaining budget so a single attempt
91+
# cannot overrun the operation window.
92+
local install_ok=false
93+
if [ "$hasBudget" = true ]; then
94+
local remaining=$(( maxBudget - ( $(date +%s) - opStartTime ) ))
95+
if [ "$remaining" -lt 1 ]; then
96+
echo "apt_get_install budget of ${maxBudget}s exceeded, exiting early." >&2
97+
return 2
98+
fi
99+
timeout "$remaining" apt-get install ${apt_opts} -o Dpkg::Options::="--force-confold" --no-install-recommends "${packages[@]}" && install_ok=true
100+
else
101+
apt-get install ${apt_opts} -o Dpkg::Options::="--force-confold" --no-install-recommends "${packages[@]}" && install_ok=true
102+
fi
103+
104+
if [ "$install_ok" = true ]; then
105+
echo "Executed apt-get install \"${packages[*]}\" $i times"
66106
wait_for_apt_locks
67107
DEBIAN_FRONTEND=noninteractive apt-get clean
68108
wait_for_apt_locks
@@ -72,14 +112,31 @@ _apt_get_install() {
72112
if [ $i -eq $retries ]; then
73113
return 1
74114
else
115+
# Check budget/CSE again before sleeping
116+
if ! check_cse_timeout; then
117+
echo "CSE timeout approaching, exiting apt_get_install early." >&2
118+
return 2
119+
fi
120+
if [ "$hasBudget" = true ]; then
121+
local postElapsed=$(( $(date +%s) - opStartTime ))
122+
if [ "$postElapsed" -ge "$maxBudget" ]; then
123+
echo "apt_get_install budget of ${maxBudget}s exceeded after ${postElapsed}s, exiting early." >&2
124+
return 2
125+
fi
126+
fi
75127
sleep $wait_sleep
76128
apt_get_update
77129
fi
78130
done
79131
}
80132
apt_get_install() {
81-
retries=$1; wait_sleep=$2; timeout=$3; shift && shift && shift
82-
_apt_get_install $retries $wait_sleep "-y" "$@"
133+
local retries=$1; local wait_sleep=$2; local timeout=$3; shift && shift && shift
134+
# Only apply per-operation budget during real CSE runs; during VHD build use no cap.
135+
local maxBudget=0
136+
if [ -n "${CSE_STARTTIME_SECONDS:-}" ]; then
137+
maxBudget=$timeout
138+
fi
139+
_apt_get_install "$retries" "$wait_sleep" "-y" "$maxBudget" "$@"
83140
}
84141
apt_get_purge() {
85142
retries=$1; wait_sleep=$2; timeout=$3; shift && shift && shift
@@ -168,7 +225,8 @@ apt_get_install_from_local_repo() {
168225
# Install package from local repo using core installation function
169226
local retries=10
170227
local wait_sleep=5
171-
if ! _apt_get_install $retries $wait_sleep "${opts}" "${package_name}"; then
228+
# maxBudget=0: no per-operation time cap for local repo installs (no network download involved)
229+
if ! _apt_get_install $retries $wait_sleep "${opts}" 0 "${package_name}"; then
172230
echo "Failed to install ${package_name} from local repo"
173231
rm -f "${tmp_list}"
174232
rmdir "${tmp_dir}"
Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
#!/bin/bash
2+
3+
Describe 'apt_get_install budget timeout'
4+
apt_install_precheck() {
5+
CSE_STARTTIME_SECONDS=$(date +%s)
6+
}
7+
BeforeEach apt_install_precheck
8+
9+
Include "./parts/linux/cloud-init/artifacts/cse_helpers.sh"
10+
Include "./parts/linux/cloud-init/artifacts/ubuntu/cse_helpers_ubuntu.sh"
11+
12+
Describe '_apt_get_install with budget'
13+
# Mock apt-related commands and timeout to isolate budget logic
14+
wait_for_apt_locks() { :; }
15+
dpkg() { :; }
16+
apt_get_update() { :; }
17+
# Pass-through timeout mock: skip the timeout value, run the rest
18+
timeout() { shift; "$@"; }
19+
apt-get() {
20+
# Simulate install failure by default so retries are exercised
21+
if [ "$1" = "install" ]; then
22+
return 1
23+
fi
24+
# apt-get clean, etc.
25+
return 0
26+
}
27+
28+
It "returns 0 when install succeeds within budget"
29+
apt-get() {
30+
if [ "$1" = "install" ]; then
31+
return 0
32+
fi
33+
return 0
34+
}
35+
When call _apt_get_install 3 1 "-y" 60 fake-package
36+
The status should eq 0
37+
The stdout should include 'Executed apt-get install "fake-package"'
38+
End
39+
40+
It "logs all package names when installing multiple packages"
41+
apt-get() {
42+
if [ "$1" = "install" ]; then
43+
return 0
44+
fi
45+
return 0
46+
}
47+
When call _apt_get_install 1 0 "-y" 0 pkg-one pkg-two pkg-three
48+
The status should eq 0
49+
The stdout should include 'Executed apt-get install "pkg-one pkg-two pkg-three"'
50+
End
51+
52+
It "returns 1 when install fails and retries exhausted (no budget)"
53+
When call _apt_get_install 2 0 "-y" 0 fake-package
54+
The status should eq 1
55+
End
56+
57+
It "returns 2 when per-operation budget is exceeded"
58+
# Mock timeout to sleep so elapsed time exceeds the 1s budget
59+
timeout() {
60+
sleep 2
61+
return 1
62+
}
63+
When call _apt_get_install 5 0 "-y" 1 fake-package
64+
The status should eq 2
65+
The stderr should include "apt_get_install budget of 1s exceeded"
66+
End
67+
68+
It "returns 2 when CSE timeout is already exceeded before first attempt"
69+
CSE_STARTTIME_SECONDS=$(( $(date +%s) - 800 ))
70+
When call _apt_get_install 3 1 "-y" 600 fake-package
71+
The status should eq 2
72+
The stderr should include "CSE timeout approaching"
73+
End
74+
75+
It "does not apply budget when CSE_STARTTIME_SECONDS is unset"
76+
unset CSE_STARTTIME_SECONDS
77+
apt-get() {
78+
if [ "$1" = "install" ]; then
79+
return 0
80+
fi
81+
return 0
82+
}
83+
# maxBudget=1 but since CSE_STARTTIME_SECONDS is unset, budget is ignored by apt_get_install wrapper
84+
# Here we test _apt_get_install directly with budget=0 (what the wrapper passes when unset)
85+
When call _apt_get_install 1 0 "-y" 0 fake-package
86+
The status should eq 0
87+
The stdout should include 'Executed apt-get install "fake-package"'
88+
The stderr should include "Warning: CSE_STARTTIME_SECONDS environment variable is not set."
89+
End
90+
End
91+
92+
Describe 'apt_get_install wrapper'
93+
wait_for_apt_locks() { :; }
94+
dpkg() { :; }
95+
apt_get_update() { :; }
96+
timeout() { shift; "$@"; }
97+
apt-get() {
98+
if [ "$1" = "install" ]; then
99+
return 0
100+
fi
101+
return 0
102+
}
103+
104+
It "passes timeout as budget during CSE run"
105+
CSE_STARTTIME_SECONDS=$(date +%s)
106+
When call apt_get_install 1 0 60 fake-package
107+
The status should eq 0
108+
The stdout should include 'Executed apt-get install "fake-package"'
109+
End
110+
111+
It "does not apply budget during VHD build (CSE_STARTTIME_SECONDS unset)"
112+
unset CSE_STARTTIME_SECONDS
113+
# Override timeout mock to fail if called — proves budget was not applied
114+
timeout() {
115+
echo "ERROR: timeout should not be called during VHD build" >&2
116+
return 1
117+
}
118+
When call apt_get_install 1 0 60 fake-package
119+
The status should eq 0
120+
The stdout should include 'Executed apt-get install "fake-package"'
121+
The stderr should include "Warning: CSE_STARTTIME_SECONDS environment variable is not set."
122+
End
123+
End
124+
End

0 commit comments

Comments
 (0)