test: stop proxy before cleaning up its query engine

This commit is contained in:
Yuneng Jiang 2026-09-14 04:04:21 -07:00
parent d7c13625e9
commit 16d7e695c0
No known key found for this signature in database
2 changed files with 23 additions and 9 deletions

View file

@ -14,7 +14,8 @@ guard6_installed=false
cleanup() {
original_status=$?
trap - EXIT INT TERM
sudo .venv/bin/python .circleci/scripts/stop_integration_processes.py "$integration_identity" "$(id -u)" \
sudo .venv/bin/python .circleci/scripts/stop_integration_processes.py \
"$integration_identity" "$(id -u)" "$proxy_pid" "$upstream_pid" \
> "$results/process-cleanup.txt" 2>&1 || original_status=1
for owned_pid in "$proxy_pid" "$upstream_pid"; do
if [ -n "$owned_pid" ]; then

View file

@ -6,10 +6,10 @@ import psutil
def owned_processes(identity: str, owner_uid: int) -> tuple[psutil.Process, ...]:
owned: Final[list[psutil.Process]] = []
for process in psutil.process_iter(["uids"]):
if process.info["uids"].real != owner_uid:
continue
for process in psutil.process_iter():
try:
if process.uids().real != owner_uid:
continue
if process.environ().get("INTEGRATION_RUN_ID") == identity:
owned.append(process)
except psutil.NoSuchProcess:
@ -17,15 +17,23 @@ def owned_processes(identity: str, owner_uid: int) -> tuple[psutil.Process, ...]
return tuple(owned)
def main(identity: str, owner_uid: int) -> int:
def main(identity: str, owner_uid: int, root_pids: tuple[int, ...]) -> int:
assert owner_uid > 0, "The integration process owner must be a non-root UID"
owned: Final = owned_processes(identity, owner_uid)
for process in owned:
roots: Final = tuple(process for process in owned if process.pid in root_pids)
for process in roots:
try:
process.terminate()
except psutil.NoSuchProcess:
continue
psutil.wait_procs(owned, timeout=8)
psutil.wait_procs(roots, timeout=30)
residual: Final = owned_processes(identity, owner_uid)
for process in residual:
try:
process.terminate()
except psutil.NoSuchProcess:
continue
psutil.wait_procs(residual, timeout=10)
remaining: Final = owned_processes(identity, owner_uid)
for process in remaining:
try:
@ -34,9 +42,14 @@ def main(identity: str, owner_uid: int) -> int:
continue
psutil.wait_procs(remaining, timeout=2)
survivors: Final = owned_processes(identity, owner_uid)
print(f"Owned integration processes: {len(owned)}, forced: {len(remaining)}, remaining: {len(survivors)}")
print(
f"Owned integration processes: {len(owned)}, roots: {len(roots)}, "
f"residual: {len(residual)}, forced: {len(remaining)}, remaining: {len(survivors)}"
)
for process in remaining:
print(f"Forced cleanup was required for PID {process.pid}")
return 1 if remaining or survivors else 0
if __name__ == "__main__":
raise SystemExit(main(sys.argv[1], int(sys.argv[2])))
raise SystemExit(main(sys.argv[1], int(sys.argv[2]), tuple(int(value) for value in sys.argv[3:] if value)))