Skip to content

Commit 8d2d38a

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 8d2d38a

File tree

1 file changed

+182
-24
lines changed

1 file changed

+182
-24
lines changed

‎tools/btfgen.sh‎

Lines changed: 182 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,33 @@
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() { echo "Usage: $0 [-a <x86_64|arm64> -o <file01.bpf.o> -o <file02.bpf.o>] [-j <num_jobs>]" 1>&2; exit 1; }
44

5-
on=0
5+
# Default to number of processors, fallback to 1 if nproc is not available
6+
j=$(nproc 2>/dev/null || echo 1)
67

7-
while getopts ":a:o:" opt; do
8+
while getopts ":a:o:j:" opt; do
89
case "${opt}" in
910
a)
10-
a=${OPTARG}
11+
a="${OPTARG}"
1112
[[ "${a}" != "x86_64" && "${a}" != "arm64" ]] && usage
1213
;;
1314
o)
14-
[[ ! -f ${OPTARG} ]] && { echo "error: could not find bpf object: ${OPTARG}"; usage; }
15+
[[ ! -f "${OPTARG}" ]] && { echo "error: could not find bpf object: ${OPTARG}"; usage; }
1516
o+=("${OPTARG}")
1617
;;
18+
j)
19+
j="${OPTARG}"
20+
[[ ! "${j}" =~ ^[0-9]+$ ]] && { echo "error: -j must be a positive integer"; usage; }
21+
[[ "${j}" -lt 1 ]] && { echo "error: -j must be at least 1"; usage; }
22+
;;
1723
*)
1824
usage
1925
;;
2026
esac
2127
done
2228
shift $((OPTIND-1))
2329

24-
if [ -z "${a}" ] || [ -z "${o}" ]; then
30+
if [ -z "${a}" ] || [ ${#o[@]} -eq 0 ]; then
2531
usage
2632
fi
2733

@@ -30,19 +36,19 @@ for ofile in "${o[@]}"; do
3036
obj_cmdline+="${ofile} "
3137
done
3238

33-
basedir=$(dirname ${0})/..
39+
basedir=$(dirname "${0}")/..
3440
if [ "${basedir}" == "." ]; then
35-
basedir=$(pwd)/..
41+
basedir="$(pwd)/.."
3642
fi
3743

38-
if [ ! -d ${basedir}/archive ]; then
44+
if [ ! -d "${basedir}/archive" ]; then
3945
echo "error: could not find archive directory"
4046
exit 1
4147
fi
4248

43-
cd ${basedir}
49+
cd "${basedir}" || exit 1
4450

45-
btfgen=$(which bpftool)
51+
btfgen="$(which bpftool)"
4652
if [ -z "${btfgen}" ]; then
4753
btfgen=/usr/sbin/bpftool
4854
fi
@@ -52,10 +58,31 @@ if [ ! -x "${btfgen}" ]; then
5258
exit 1
5359
fi
5460

61+
# Track background jobs for cleanup
62+
declare -a background_pids=()
63+
5564
function ctrlc ()
5665
{
5766
echo "Exiting due to ctrl-c..."
58-
rm ${basedir}/*.btf
67+
68+
# Kill all background jobs
69+
for pid in "${background_pids[@]}"; do
70+
if kill -0 "${pid}" 2>/dev/null; then
71+
kill "${pid}" 2>/dev/null
72+
fi
73+
done
74+
75+
# Wait a bit for graceful shutdown
76+
sleep 1
77+
78+
# Force kill if still running
79+
for pid in "${background_pids[@]}"; do
80+
if kill -0 "${pid}" 2>/dev/null; then
81+
kill -9 "${pid}" 2>/dev/null
82+
fi
83+
done
84+
85+
rm -f "${basedir}"/*.btf
5986

6087
exit 2
6188
}
@@ -66,21 +93,152 @@ trap ctrlc SIGTERM
6693
# clean custom-archive directory
6794
find ./custom-archive -mindepth 1 -maxdepth 1 -type d -exec rm -rf {} \;
6895

96+
# Function to process a single BTF file
97+
process_btf_file() {
98+
local file="$1"
99+
local btfgen="$2"
100+
local obj_cmdline="$3"
101+
102+
local dir
103+
local extracted
104+
local temp_dir
105+
local original_dir
106+
107+
dir="$(dirname "${file}")"
108+
109+
# Create a temporary directory for this job to avoid conflicts
110+
temp_dir="$(mktemp -d)"
111+
original_dir="$(pwd)"
112+
113+
# Extract in temp directory
114+
cd "${temp_dir}" || return 1
115+
extracted="$(tar xvfJ "${original_dir}/${file}" 2>/dev/null)"
116+
local ret=$?
117+
118+
if [[ ${ret} -eq 0 && -f "${extracted}" ]]; then
119+
cd "${original_dir}" || return 1
120+
121+
# Prepare output directory
122+
dir=${dir/\.\/archive\/}
123+
local out_dir="./custom-archive/${dir}"
124+
mkdir -p "${out_dir}"
125+
126+
# Move extracted file to working directory and process
127+
mv "${temp_dir}/${extracted}" "./${extracted}"
128+
129+
# Generate minimized BTF file
130+
# shellcheck disable=SC2086
131+
"${btfgen}" gen min_core_btf "${extracted}" "${out_dir}/${extracted}" ${obj_cmdline}
132+
local btfgen_ret=$?
133+
134+
# Cleanup
135+
rm -f "./${extracted}"
136+
137+
if [[ ${btfgen_ret} -eq 0 ]]; then
138+
printf "[SUCCESS] %s\n" "${extracted}"
139+
# Cleanup temp directory
140+
rm -rf "${temp_dir}"
141+
return 0
142+
else
143+
printf "[FAIL] %s\n" "${extracted}"
144+
# Cleanup temp directory
145+
rm -rf "${temp_dir}"
146+
return 1
147+
fi
148+
else
149+
cd "${original_dir}" || return 1
150+
printf "[FAIL] %s (extraction failed)\n" "$(basename "${file}")"
151+
# Cleanup temp directory
152+
rm -rf "${temp_dir}"
153+
return 1
154+
fi
155+
}
156+
157+
# Export the function so it can be used by background processes
158+
export -f process_btf_file
159+
160+
echo "Using ${j} parallel jobs for BTF processing..."
161+
162+
# Ensure output is line-buffered for better real-time display
163+
stty -icanon min 1 time 0 2>/dev/null || true
164+
165+
# Initialize job control variables
166+
job_count=0
167+
failed_jobs=0
168+
completed_jobs=0
169+
start_time="$(date +%s)"
170+
171+
# Collect all BTF files to process
172+
btf_files=()
69173
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=$?
174+
while IFS= read -r -d '' file; do
175+
btf_files+=("${file}")
176+
done < <(find "./archive/${dir}" -name "*.tar.xz" -print0)
177+
done
75178

76-
dir=${dir/\.\/archive\/}
77-
out_dir="./custom-archive/${dir}"
78-
[[ ! -d ${out_dir} ]] && mkdir -p ${out_dir}
179+
total_files=${#btf_files[@]}
180+
echo "Found ${total_files} BTF files to process"
79181

80-
echo "Processing ${extracted}..."
182+
if [[ ${total_files} -eq 0 ]]; then
183+
echo "No BTF files found for architecture ${a}"
184+
exit 0
185+
fi
81186

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}
187+
# Show system info
188+
echo "System: $(nproc) CPU cores, $(free -h | awk '/^Mem:/ {print $2}') RAM"
189+
echo "Started at: $(date)"
190+
echo
191+
192+
# Process files in parallel with job control
193+
for file in "${btf_files[@]}"; do
194+
# Wait if we've reached the maximum number of parallel jobs
195+
while [[ ${job_count} -ge ${j} ]]; do
196+
# Wait for any background job to complete
197+
wait -n
198+
exit_code=$?
199+
((job_count--))
200+
((completed_jobs++))
201+
202+
# Track failed jobs
203+
if [[ ${exit_code} -ne 0 ]]; then
204+
((failed_jobs++))
205+
fi
85206
done
207+
208+
# Start new background job
209+
process_btf_file "${file}" "${btfgen}" "${obj_cmdline}" &
210+
pid=$!
211+
background_pids+=("${pid}")
212+
((job_count++))
86213
done
214+
215+
# Wait for all remaining background jobs to complete
216+
echo -e "\nWaiting for remaining ${job_count} jobs to complete..."
217+
while [[ ${job_count} -gt 0 ]]; do
218+
wait -n
219+
exit_code=$?
220+
((job_count--))
221+
((completed_jobs++))
222+
223+
if [[ ${exit_code} -ne 0 ]]; then
224+
((failed_jobs++))
225+
fi
226+
done
227+
228+
# Final summary with timing
229+
end_time="$(date +%s)"
230+
total_elapsed=$((end_time - start_time))
231+
average_rate=$((total_files * 60 / (total_elapsed + 1)))
232+
233+
echo -e "\n\n🎉 BTF processing completed!"
234+
echo "📊 Total files: ${total_files}"
235+
echo "✅ Completed: ${completed_jobs}"
236+
echo "❌ Failed jobs: ${failed_jobs}"
237+
echo "⚙️ Parallel jobs: ${j}"
238+
echo "⏱️ Total time: $((total_elapsed / 60))m $((total_elapsed % 60))s"
239+
echo "🚀 Average rate: ${average_rate} files/min"
240+
241+
if [[ ${failed_jobs} -gt 0 ]]; then
242+
echo "Warning: Some BTF files failed to process"
243+
exit 1
244+
fi

0 commit comments

Comments
 (0)