Skip to content

Commit 287a99f

Browse files
geyslangeyslan
committed
feat: enhance btfgen with parallel processing
- Add support for specifying the number of jobs with the -j option. - Improve error handling for job execution and extraction failures. - Implement background job tracking for graceful cleanup on interruption. - Refactor code for better readability and maintainability.
1 parent f44f82a commit 287a99f

File tree

1 file changed

+217
-31
lines changed

1 file changed

+217
-31
lines changed

‎tools/btfgen.sh‎

Lines changed: 217 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,62 @@
11
#!/bin/bash
22

3-
usage() { echo "Usage: $0 [-a <x86_64|arm64> -o <file01.bpf.o> -o <file02.bpf.o>]" 1>&2; exit 1; }
3+
usage() {
4+
echo "Usage: $0 [-a <x86_64|arm64> -o <file01.bpf.o> -o <file02.bpf.o>] [-j <num_jobs>]" 1>&2
5+
exit 1
6+
}
7+
8+
# Set default jobs to 1
9+
j=1
10+
11+
# Get number of CPUs for validation when user sets -j
12+
if command -v nproc > /dev/null 2>&1; then
13+
max_jobs=$(nproc)
14+
elif [ -f /proc/cpuinfo ]; then
15+
max_jobs=$(grep -c ^processor /proc/cpuinfo)
16+
else
17+
max_jobs=1
18+
fi
19+
# Fallback to 1 if detection failed or result is empty/non-numeric
20+
if ! [ "${max_jobs}" -ge 1 ] 2> /dev/null; then
21+
max_jobs=1
22+
fi
423

5-
on=0
24+
o=()
625

7-
while getopts ":a:o:" opt; do
26+
while getopts ":a:o:j:" opt; do
827
case "${opt}" in
928
a)
10-
a=${OPTARG}
29+
a="${OPTARG}"
1130
[[ "${a}" != "x86_64" && "${a}" != "arm64" ]] && usage
12-
;;
31+
;;
1332
o)
14-
[[ ! -f ${OPTARG} ]] && { echo "error: could not find bpf object: ${OPTARG}"; usage; }
33+
[[ ! -f "${OPTARG}" ]] && {
34+
echo "error: could not find bpf object: ${OPTARG}"
35+
usage
36+
}
1537
o+=("${OPTARG}")
16-
;;
38+
;;
39+
j)
40+
j="${OPTARG}"
41+
# Validate it's a positive integer
42+
if ! [ "${j}" -ge 1 ] 2> /dev/null; then
43+
echo "error: -j must be a positive integer (got: ${j})"
44+
usage
45+
fi
46+
# Cap at CPU count if exceeded
47+
if [ "${j}" -gt "${max_jobs}" ]; then
48+
echo "warning: -j ${j} exceeds CPU count (${max_jobs}), using ${max_jobs}"
49+
j="${max_jobs}"
50+
fi
51+
;;
1752
*)
1853
usage
19-
;;
54+
;;
2055
esac
2156
done
22-
shift $((OPTIND-1))
57+
shift $((OPTIND - 1))
2358

24-
if [ -z "${a}" ] || [ -z "${o}" ]; then
59+
if [ -z "${a}" ] || [ ${#o[@]} -eq 0 ]; then
2560
usage
2661
fi
2762

@@ -30,19 +65,19 @@ for ofile in "${o[@]}"; do
3065
obj_cmdline+="${ofile} "
3166
done
3267

33-
basedir=$(dirname ${0})/..
68+
basedir=$(dirname "${0}")/..
3469
if [ "${basedir}" == "." ]; then
35-
basedir=$(pwd)/..
70+
basedir="$(pwd)/.."
3671
fi
3772

38-
if [ ! -d ${basedir}/archive ]; then
73+
if [ ! -d "${basedir}/archive" ]; then
3974
echo "error: could not find archive directory"
4075
exit 1
4176
fi
4277

43-
cd ${basedir}
78+
cd "${basedir}" || exit 1
4479

45-
btfgen=$(which bpftool)
80+
btfgen="$(which bpftool)"
4681
if [ -z "${btfgen}" ]; then
4782
btfgen=/usr/sbin/bpftool
4883
fi
@@ -52,10 +87,30 @@ if [ ! -x "${btfgen}" ]; then
5287
exit 1
5388
fi
5489

55-
function ctrlc ()
56-
{
90+
# Track background jobs for cleanup
91+
declare -a background_pids=()
92+
93+
function ctrlc() {
5794
echo "Exiting due to ctrl-c..."
58-
rm ${basedir}/*.btf
95+
96+
# Kill all background jobs
97+
for pid in "${background_pids[@]}"; do
98+
if kill -0 "${pid}" 2> /dev/null; then
99+
kill "${pid}" 2> /dev/null
100+
fi
101+
done
102+
103+
# Wait a bit for graceful shutdown
104+
sleep 1
105+
106+
# Force kill if still running
107+
for pid in "${background_pids[@]}"; do
108+
if kill -0 "${pid}" 2> /dev/null; then
109+
kill -9 "${pid}" 2> /dev/null
110+
fi
111+
done
112+
113+
rm -f "${basedir}"/*.btf
59114

60115
exit 2
61116
}
@@ -66,21 +121,152 @@ trap ctrlc SIGTERM
66121
# clean custom-archive directory
67122
find ./custom-archive -mindepth 1 -maxdepth 1 -type d -exec rm -rf {} \;
68123

69-
for dir in $(find ./archive/ -iregex ".*${a}.*" -type d | sed 's:\.\/archive\/::g'| sort -u); do
70-
# uncompress and process each existing input BTF .tar.xz file
71-
for file in $(find ./archive/${dir} -name *.tar.xz); do
72-
dir=$(dirname $file)
73-
base=$(basename $file)
74-
extracted=$(tar xvfJ $dir/$base); ret=$?
124+
# Function to process a single BTF file
125+
process_btf_file() {
126+
local file="$1"
127+
local btfgen="$2"
128+
local obj_cmdline="$3"
129+
130+
local dir
131+
local extracted
132+
local temp_dir
133+
local original_dir
134+
135+
dir="$(dirname "${file}")"
136+
137+
# Create a temporary directory for this job to avoid conflicts
138+
temp_dir="$(mktemp -d)"
139+
original_dir="$(pwd)"
140+
141+
# Extract in temp directory
142+
cd "${temp_dir}" || return 1
143+
extracted="$(tar xvfJ "${original_dir}/${file}" 2> /dev/null)"
144+
local ret=$?
145+
146+
if [[ ${ret} -eq 0 && -f "${extracted}" ]]; then
147+
cd "${original_dir}" || return 1
148+
149+
# Prepare output directory
150+
dir=${dir/\.\/archive\//}
151+
local out_dir="./custom-archive/${dir}"
152+
mkdir -p "${out_dir}"
153+
154+
# Move extracted file to working directory and process
155+
mv "${temp_dir}/${extracted}" "./${extracted}"
156+
157+
# Generate minimized BTF file
158+
# shellcheck disable=SC2086
159+
"${btfgen}" gen min_core_btf "${extracted}" "${out_dir}/${extracted}" ${obj_cmdline}
160+
local btfgen_ret=$?
75161

76-
dir=${dir/\.\/archive\/}
77-
out_dir="./custom-archive/${dir}"
78-
[[ ! -d ${out_dir} ]] && mkdir -p ${out_dir}
162+
# Cleanup
163+
rm -f "./${extracted}"
79164

80-
echo "Processing ${extracted}..."
165+
if [[ ${btfgen_ret} -eq 0 ]]; then
166+
printf "[SUCCESS] %s\n" "${extracted}"
167+
# Cleanup temp directory
168+
rm -rf "${temp_dir}"
169+
return 0
170+
else
171+
printf "[FAIL] %s\n" "${extracted}"
172+
# Cleanup temp directory
173+
rm -rf "${temp_dir}"
174+
return 1
175+
fi
176+
else
177+
cd "${original_dir}" || return 1
178+
printf "[FAIL] %s (extraction failed)\n" "$(basename "${file}")"
179+
# Cleanup temp directory
180+
rm -rf "${temp_dir}"
181+
return 1
182+
fi
183+
}
184+
185+
# Export the function so it can be used by background processes
186+
export -f process_btf_file
187+
188+
echo "Using ${j} parallel jobs for BTF processing..."
189+
190+
# Ensure output is line-buffered for better real-time display
191+
stty -icanon min 1 time 0 2> /dev/null || true
192+
193+
# Initialize job control variables
194+
job_count=0
195+
failed_jobs=0
196+
completed_jobs=0
197+
start_time="$(date +%s)"
198+
199+
# Collect all BTF files to process
200+
btf_files=()
201+
for dir in $(find ./archive/ -iregex ".*${a}.*" -type d | sed 's:\.\/archive\/::g' | sort -u); do
202+
while IFS= read -r -d '' file; do
203+
btf_files+=("${file}")
204+
done < <(find "./archive/${dir}" -name "*.tar.xz" -print0)
205+
done
206+
207+
total_files=${#btf_files[@]}
208+
echo "Found ${total_files} BTF files to process"
209+
210+
if [[ ${total_files} -eq 0 ]]; then
211+
echo "No BTF files found for architecture ${a}"
212+
exit 0
213+
fi
214+
215+
# Show system info
216+
echo "System: $(nproc) CPU cores, $(free -h | awk '/^Mem:/ {print $2}') RAM"
217+
echo "Started at: $(date)"
218+
echo
81219

82-
# generate one output BTF file to each input BTF file given
83-
$btfgen gen min_core_btf ${extracted} ${out_dir}/${extracted} ${obj_cmdline}
84-
[[ $ret -eq 0 ]] && [[ -f ./${extracted} ]] && rm -f ./${extracted}
220+
# Process files in parallel with job control
221+
for file in "${btf_files[@]}"; do
222+
# Wait if we've reached the maximum number of parallel jobs
223+
while [[ ${job_count} -ge ${j} ]]; do
224+
# Wait for any background job to complete
225+
wait -n
226+
exit_code=$?
227+
((job_count--))
228+
((completed_jobs++))
229+
230+
# Track failed jobs
231+
if [[ ${exit_code} -ne 0 ]]; then
232+
((failed_jobs++))
233+
fi
85234
done
235+
236+
# Start new background job
237+
process_btf_file "${file}" "${btfgen}" "${obj_cmdline}" &
238+
pid=$!
239+
background_pids+=("${pid}")
240+
((job_count++))
241+
done
242+
243+
# Wait for all remaining background jobs to complete
244+
echo -e "\nWaiting for remaining ${job_count} jobs to complete..."
245+
while [[ ${job_count} -gt 0 ]]; do
246+
wait -n
247+
exit_code=$?
248+
((job_count--))
249+
((completed_jobs++))
250+
251+
if [[ ${exit_code} -ne 0 ]]; then
252+
((failed_jobs++))
253+
fi
86254
done
255+
256+
# Final summary with timing
257+
end_time="$(date +%s)"
258+
total_elapsed=$((end_time - start_time))
259+
average_rate=$((total_files * 60 / (total_elapsed + 1)))
260+
261+
echo -e "\n\n🎉 BTF processing completed!"
262+
echo "📊 Total files: ${total_files}"
263+
echo "✅ Completed: ${completed_jobs}"
264+
echo "❌ Failed jobs: ${failed_jobs}"
265+
echo "⚙️ Parallel jobs: ${j}"
266+
echo "⏱️ Total time: $((total_elapsed / 60))m $((total_elapsed % 60))s"
267+
echo "🚀 Average rate: ${average_rate} files/min"
268+
269+
if [[ ${failed_jobs} -gt 0 ]]; then
270+
echo "error: Some BTF files failed to process"
271+
exit 1
272+
fi

0 commit comments

Comments
 (0)