bash-com/pool

#!/bin/bash


# (1) Init.
#--------------------------------------------------------#
# (1) Get current absolute dir
ROOT=$(dirname `realpath $0`);

# (2) Check argc
test $# -lt 2 && echo -e "error: too fewarguments\n\n\e[1mUSAGE\e[0m\n\tpool <key> <port>\n\n\e[1mARGUMENTS\e[0m\n\t<key>\tThe key of the pool (client must match it)\n\t<port>\tTo port to listen to\n" && exit 1;

# (3) Check @PORT range #
MIN_PORT=1024;
MAX_PORT=49151;
test "$2" -gt "$MAX_PORT" -o "$2" -lt "$MIN_PORT" && echo "error: <port> must be between $MIN_PORT and $MAX_PORT" && exit 1;

# (4) Set argument explicit names #
POOL_KEY="$1";
POOL_PORT="$2";

# (5) Create a temporary file #
TMPFILE="/dev/shm/pool_pid";


# (6) Define on-exit routine #
LOOP_PID="";
on_exit(){

	test -f /dev/shm/pool_input_pid && source /dev/shm/pool_input_pid && rm /dev/shm/pool_input_pid;

	echo "- killing oop";
	kill -INT $LOOP_PID 2>/dev/null;

	echo "- killing ${#INPUT_PID[@]} loggers";
	for logger_name in "${!INPUT_PID[@]}"; do

		# kill each logger input
		echo "  - killing $logger_name (${INPUT_PID[$logger_name]})";
		test ! -z "${INPUT_PID[$logger_name]}" && kill -INT ${INPUT_PID[$logger_name]} 2>/dev/null;

	done;

	echo "- killing main listener ($POOL_PID)";
	kill -INT $POOL_PID 2>/dev/null;

	# delete tmp file
	echo "- deleting tmp file";
	rm $TMPFILE 2>/dev/null;

	exit 1;

}




# (2) Launch pool listener
#--------------------------------------------------------#
# (1) Bind-input #
echo -e "\e[32m+\e[0m listen $POOL_PORT";
$ROOT/bind-input pool $POOL_PORT > $TMPFILE;
POOL_PID="`cat $TMPFILE`";



# (3) Launch pool manager
#--------------------------------------------------------#
# infinite listener
(
	USED_PORTS=($POOL_PORT);
	declare -A INPUT_PID; # will contain each bind-input pid

	trap "on_exit;" HUP INT KILL TERM;

	while true; do

		# Listen on port @POOL_PORT
		MSG=`$ROOT/read pool`;

		# Do nothing if empty msg or invalid format
		test -z "$MSG" && sleep .5 && continue;

		echo -e "  + received '\e[33m$MSG\e[0m'";

		# check if END msg
		test "$MSG" = "END${POOL_KEY}END" && echo -e "\e[31m+\e[0m stop listening on $POOL_PORT" && break;

		# check format
		echo -n "$MSG" | grep -vP '^([^:]+):([^:]+):([^:]+):([^:]+)$' >/dev/null && echo '  + invalid format' && sleep .5 && continue;

		# Extract ID - PORT
		KEY="`echo -ne \"$MSG\" | sed 's/^\(.\+\):\(.\+\):\(.\+\):\(.\+\)$/\1/'`";
		LOGGER_NAME="`echo -ne \"$MSG\" | sed 's/^\(.\+\):\(.\+\):\(.\+\):\(.\+\)$/\2/'`";
		LOGGER_HOST="`echo -ne \"$MSG\" | sed 's/^\(.\+\):\(.\+\):\(.\+\):\(.\+\)$/\3/'`";
		LOGGER_PORT="`echo -ne \"$MSG\" | sed 's/^\(.\+\):\(.\+\):\(.\+\):\(.\+\)$/\4/'`";

		# check port format
		echo -n "$LOGGER_PORT" | grep -vP "^\d+$" >/dev/null && echo '  + invalid format: NaN' && sleep .5 && continue;

		# if ID does not match
		test "$KEY" != "$POOL_KEY" && echo '  + wrong id' && sleep .5 && continue;

		# create new listening socket
		echo "    + new logger '$LOGGER_NAME' from $LOGGER_HOST";

		# Find an available port
		PORT="$MIN_PORT";
		while true; do

			# do not use same port as master (if on same machine)
			test "$PORT" = "$LOGGER_PORT" && PORT="`expr $PORT + 1`" && continue;
			# ignore already used ports
			for used in ${USED_PORTS[*]}; do
				test "$PORT" = "$used" && PORT="`expr $PORT + 1`" && continue;
			done

			# if port not in use -> use it
			ss -tl4 "( sport = :$PORT )" | grep "$PORT" >/dev/null 2>&1 || break;

			# else try next port (+1)
			PORT="`expr $PORT + 1`";

		done;

		# if already an input for this logger -> kill it
		if [ ! -z "${INPUT_PID[$LOGGER_NAME]}" ]; then
			echo "      + already an input -> killing it";
			kill -INT ${INPUT_PID[$LOGGER_NAME]};
		fi;

		# Bind-input for the logger at the free port found
		# Store bind-input PID
		echo -e "      \e[32m+\e[0m listen $PORT ($LOGGER_NAME)";
		$ROOT/bind-input $LOGGER_NAME $PORT > $TMPFILE;
		INPUT_PID["$LOGGER_NAME"]="`cat $TMPFILE`";
		declare -p INPUT_PID > /dev/shm/pool_input_pid;

		# add @PORT to used ports
		USED_PORTS="$USED_PORTS $PORT";


		# Notify that input is active to MASTER
		sleep 1;
		echo "      + sending '$PORT' to $LOGGER_HOST:$LOGGER_PORT";
		bash -c "exec 4<>/dev/tcp/$LOGGER_HOST/$LOGGER_PORT; echo -n \"$PORT\" >&4;";

		echo "    + background process";

	done;

)&
LOOP_PID=$!;

# (4) Set exit management
#--------------------------------------------------------#
# (1) Set trap for "exit" signals #
trap "kill -INT $LOOP_PID 2>/dev/null;" HUP INT KILL TERM;

# (2) Wait for main loop to end
wait $LOOP_PID;


# (3) Wait for each daemon input to end #
test -f /dev/shm/pool_input_pid && source /dev/shm/pool_input_pid;
echo "  + ${#INPUT_PID[@]} loggers running in background";
for logger_name in "${!INPUT_PID[@]}"; do

	while [ -e /proc/${INPUT_PID[$logger_name]} ]; do
		sleep .5
	done
	echo "  + $logger_name terminated (${INPUT_PID[$logger_name]})";

done;

on_exit;