From 54fdabfaa77f81632ffe6b5c01283ee30ce6fd58 Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Mon, 11 Mar 2024 17:25:52 +0100
Subject: [PATCH] Improve performance by running container checks in parallel

---
 README.md               |   1 +
 check_container_upgrade | 104 ++++++++++++++++++++++++++++------------
 2 files changed, 75 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index 4fde042..ec4e798 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,7 @@ service nagios-nrpe-server reload
 Usage : check_container_upgrade [-d] [-E /path/to/engine] [container1,...]
 	-E [path]               Force a specific engine (possible values: auto docker podman, default: auto)
 	-x [container]		Exclude specified container (could be repeat)
+	-M [integer]		Max number of container checks to run in parallel (default: 4, 0=no limit)
 	-d			Debug mode
 	-X			Enable bash tracing (=set -x)
 	-h 			Show this message
diff --git a/check_container_upgrade b/check_container_upgrade
index 5457415..595e0b0 100755
--- a/check_container_upgrade
+++ b/check_container_upgrade
@@ -9,6 +9,8 @@
 ENGINE="auto"
 POSSIBLE_ENGINES=( "auto" "docker" "podman" )
 DEBUG=0
+MAX_PARALLEL_CHECKS=4
+ONLY_CONTAINERS=()
 EXCLUDED_CONTAINERS=( buildx_buildkit_default )
 declare -rA CHECK_PLUGINS=(
 	["/usr/lib/nagios/plugins/check_apt"]="/usr/lib/nagios/plugins/check_apt -u -U -t 60 -l"
@@ -18,7 +20,7 @@ declare -rA CHECK_PLUGINS=(
 function debug() {
 	if [ $DEBUG -eq 1 ]
 	then
-		>&2 echo -e "[DEBUG] $1"
+		>&2 echo -e "[DEBUG] $@"
 	fi
 }
 
@@ -41,9 +43,10 @@ function usage() {
 	error="$1"
 	[ -n "$error" ] && echo "$error"
 	cat << EOF
-Usage : $0 [-d] [-E /path/to/engine] [container1,...]
+Usage : $(basename $0) [-d] [-E /path/to/engine] [container1,...]
 	-E [path]               Force a specific engine (possible values: ${POSSIBLE_ENGINES[@]}, default: $ENGINE)
 	-x [container]		Exclude specified container (could be repeat)
+	-M [integer]		Max number of container checks to run in parallel (default: $MAX_PARALLEL_CHECKS, 0=no limit)
 	-d			Debug mode
 	-X			Enable bash tracing (=set -x)
 	-h 			Show this message
@@ -52,34 +55,43 @@ EOF
 	exit 0
 }
 
-while getopts "hdXE:x:" OPTION
+idx=1
+while [ $idx -le $# ]
 do
-	case $OPTION in
-		d)
+        OPT=${!idx}
+        case $OPT in
+		-d)
 			DEBUG=1
 		;;
-		h)
+		-h)
 			usage
 		;;
-		E)
-			ENGINE=$OPTARG
+		-E)
+			((idx++))
+			ENGINE=${!idx}
 			if [ ! -x "$ENGINE" ]
 			then
-				in_array $ENGINE $POSSIBLE_ENGINES || usage "Invalid engine $ENGINE"
+				in_array $ENGINE ${POSSIBLE_ENGINES[@]} || usage "Invalid engine $ENGINE"
 			fi
 		;;
-		x)
-			EXCLUDED_CONTAINERS+=( $OPTARG )
+		-x)
+			((idx++))
+			EXCLUDED_CONTAINERS+=( ${!idx} )
 		;;
-		X)
+		-M)
+			((idx++))
+			MAX_PARALLEL_CHECKS=${!idx}
+		;;
+		-X)
 			set -x
 		;;
-		*)
-			usage "Unknown option $OPTION"
+                *)
+			ONLY_CONTAINERS+=( $OPT )
 		;;
-	esac
+        esac
+        ((idx++))
 done
-ONLY_CONTAINERS=( "${@:$OPTIND}" )
+
 ! is_empty $ONLY_CONTAINERS && debug "Only containers: ${ONLY_CONTAINERS[@]}"
 
 if [ "$ENGINE" == "auto" ]
@@ -106,6 +118,8 @@ then
 fi
 
 EXIT_CODE=0
+declare -A CONTAINER_STATUS_FILE
+declare -A CONTAINER_PID
 declare -A UPTODATE
 declare -A ERRORS
 CHECKED_CONTAINERS=( )
@@ -114,19 +128,11 @@ debug "List running containers..."
 RUNNING_CONTAINERS=$($ENGINE ps --format '{{.Names}}' | tr '\n' ' ')
 debug "Running containers: $RUNNING_CONTAINERS"
 
-for container in $RUNNING_CONTAINERS
-do
-	if ! is_empty $ONLY_CONTAINERS && ! in_array $container $ONLY_CONTAINERS
-	then
-		debug "$container - Ignored"
-		continue
-	fi
-	if in_array $container $EXCLUDED_CONTAINERS
-	then
-		debug "$container - Excluded"
-		continue
-	fi
-	CHECKED_CONTAINERS+=( "$container" )
+# Implement check inside a function to allow running it in parallel
+# Parameters : [container] [output file]
+function check_container() {
+	container="$1"
+	output_file="$2"
 	STATUS=""
 	for check_plugin in ${CHECK_PLUGINS[@]}
 	do
@@ -149,6 +155,44 @@ do
 		STATUS="UNKNOWN - No check plugin available"
 		ex=3
 	fi
+	echo $STATUS > $output_file
+	return $ex
+}
+
+debug "Trigger check of all selected containers..."
+for container in $RUNNING_CONTAINERS
+do
+	if ! is_empty $ONLY_CONTAINERS && ! in_array $container ${ONLY_CONTAINERS[@]}
+	then
+		debug "$container - Ignored"
+		continue
+	fi
+	if in_array $container ${EXCLUDED_CONTAINERS[@]}
+	then
+		debug "$container - Excluded"
+		continue
+	fi
+	if [ $MAX_PARALLEL_CHECKS -gt 0 -a "$(jobs | wc -l)" -ge $MAX_PARALLEL_CHECKS ]
+	then
+		debug "Max parallel checks count reached. Waiting some check ending"
+		wait -n
+		debug "Some check ended, continue"
+	fi
+	CHECKED_CONTAINERS+=( "$container" )
+	CONTAINER_STATUS_FILE+=( ["$container"]=$( mktemp ) )
+	check_container $container ${CONTAINER_STATUS_FILE[$container]} & CONTAINER_PID+=( ["$container"]=$! )
+done
+
+debug "Wait for each individual container check and handle their result..."
+for container in ${!CONTAINER_PID[@]}
+do
+	pid=${CONTAINER_PID[$container]}
+	debug "$container - Waiting for PID ${pid}..."
+	wait $pid
+	ex=$?
+	debug "$container - Check return ${ex}"
+	STATUS=$( cat ${CONTAINER_STATUS_FILE[$container]} )
+	rm -f ${CONTAINER_STATUS_FILE[$container]}
 	if [ $ex -eq 0 ]
 	then
 		UPTODATE+=( ["$container"]=$STATUS )
@@ -164,7 +208,7 @@ if ! is_empty $ONLY_CONTAINERS
 then
 	for container in ${ONLY_CONTAINERS[@]}
 	do
-		if ! in_array $container $CHECKED_CONTAINERS
+		if ! in_array $container ${CHECKED_CONTAINERS[@]}
 		then
 			debug "$container - Not found"
 			ERRORS+=( ["$container"]="Not found" )