Refactored client to use persistent connections

and server to use zselect for handling concurrency

completion-client.zsh

@@ -17,24 +17,29 @@ autosuggest-ensure-server() {
 		fi
 	fi
 
+	autosuggest-server-connect
+}
+
+autosuggest-server-connect() {
+	unset ZLE_AUTOSUGGEST_CONNECTION
+
 	integer remaining_tries=10
-	# wait until the process is listening
-	while ! [[ -d $server_dir && -r $pid_file ]] ||\
-	 	! kill -0 $(<$pid_file) &> /dev/null && (( --remaining_tries )); do
+	while (( --remaining_tries )) && ! zsocket $socket_path &>/dev/null; do
 		sleep 0.3
 	done
-	ZLE_AUTOSUGGEST_SOCKET=$socket_path
+
+	[[ -z $REPLY ]] && return 1
+
+	ZLE_AUTOSUGGEST_CONNECTION=$REPLY
 }
 
 
 autosuggest-first-completion() {
-	zsocket $ZLE_AUTOSUGGEST_SOCKET &>/dev/null || return 1
-	local connection=$REPLY
-	local completion
-	print -u $connection - $1
-	while IFS= read -r -u $connection completion; do
-		print - ${completion}
+	setopt local_options noglob
+	local response
+	print -u $ZLE_AUTOSUGGEST_CONNECTION - $1 &> /dev/null || return 1
+	while IFS= read -r -u $ZLE_AUTOSUGGEST_CONNECTION response; do
+		[[ $response == '' ]] && break 
+		print - ${response}
 	done
-	# close fd
-	exec {connection}>&-
 }

completion-server.zsh

@@ -2,24 +2,6 @@
 # Based on:
 # https://github.com/Valodim/zsh-capture-completion/blob/master/capture.zsh
 
-if [[ -n $ZLE_AUTOSUGGEST_SERVER_LOG ]]; then
-	exec &>> "$HOME/.autosuggest-server.log"
-else
-	exec &> /dev/null
-fi
-exec < /dev/null
-
-zmodload zsh/zpty
-zmodload zsh/net/socket
-setopt noglob
-print "autosuggestion server started, pid: $$"
-
-# Start an interactive zsh connected to a zpty
-zpty z ZLE_DISABLE_AUTOSUGGEST=1 zsh -i
-print 'interactive shell started'
-# Source the init script
-zpty -w z "source '${0:a:h}/completion-server-init.zsh'"
-
 # read everything until a line containing the byte 0 is found
 read-to-null() {
 	while zpty -r z chunk; do
@@ -29,6 +11,73 @@ read-to-null() {
 	done
 }
 
+accept-connection() {
+	zsocket -a $server
+	fds[$REPLY]=1
+	print "connection accepted, fd: $REPLY" >&2
+}
+
+handle-request() {
+	local connection=$1 current line
+	integer read_something=0
+	print "request received from fd $connection" >&2
+	while read -u $connection prefix &> /dev/null; do
+		read_something=1
+		# send the prefix to be completed followed by a TAB to force
+		# completion
+		zpty -w -n z $prefix$'\t'
+		zpty -r z chunk &> /dev/null # read empty line before completions
+		current=''
+		# read completions one by one, storing the longest match
+		read-to-null | while IFS= read -r line; do
+			(( $#line > $#current )) && current=$line
+		done
+		# send the longest completion back to the client, strip the last
+		# non-printable character
+		if (( $#current )); then
+			print -u $connection - ${current:0:-1}
+		fi
+		# signal that we're done
+		print -u $connection ''
+		# clear input buffer
+		zpty -w z $'\n'
+		break # handle more requests/return to zselect
+	done
+	if ! (( read_something )); then
+		print "connection with fd $connection closed" >&2
+	  unset fds[$connection]
+		exec {connection}>&- # free the file descriptor
+	fi
+}
+
+
+if [[ -n $ZLE_AUTOSUGGEST_SERVER_LOG ]]; then
+	exec >> "$HOME/.autosuggest-server.log"
+else
+	exec > /dev/null
+fi
+
+if [[ -n $ZLE_AUTOSUGGEST_SERVER_LOG_ERRORS ]]; then
+	exec 2>> "$HOME/.autosuggest-server-errors.log"
+else
+	exec 2> /dev/null
+fi
+
+exec < /dev/null
+
+zmodload zsh/zpty
+zmodload zsh/zselect
+zmodload zsh/net/socket
+zmodload zsh/net/socket
+setopt noglob
+print "autosuggestion server started, pid: $$"
+
+# Start an interactive zsh connected to a zpty
+zpty z ZLE_DISABLE_AUTOSUGGEST=1 zsh -i
+print 'interactive shell started'
+# Source the init script
+zpty -w z "source '${0:a:h}/completion-server-init.zsh'"
+
 # wait for ok from shell
 read-to-null &> /dev/null
 print 'interactive shell ready'
@@ -43,6 +92,7 @@ cleanup() {
 	print 'removing socket and pid file...'
 	rm -f $socket_path $pid_file
 	print "autosuggestion server stopped, pid: $$"
+	exit
 }
 
 trap cleanup TERM INT HUP EXIT
@@ -64,31 +114,16 @@ print "server listening on '$socket_path'"
 
 print $$ > $pid_file
 
-while zsocket -a $server &> /dev/null; do
-	connection=$REPLY
-	print "connection accepted, fd: $connection"
-	# connection accepted, read the request and send response
-	while read -u $connection prefix &> /dev/null; do
-		# send the prefix to be completed followed by a TAB to force
-		# completion
-		zpty -w -n z $prefix$'\t'
-		zpty -r z chunk &> /dev/null # read empty line before completions
-		local current=''
-		# read completions one by one, storing the longest match
-		read-to-null | while IFS= read -r line; do
-			(( $#line > $#current )) && current=$line
-		done
-		# send the longest completion back to the client, strip the last
-		# non-printable character
-		if (( $#current )); then
-			print -u $connection - ${current:0:-1}
+typeset -A fds ready
+fds[$server]=1
+
+while zselect -A ready ${(k)fds}; do
+	queue=(${(k)ready})
+	for fd in $queue; do
+		if (( fd == server )); then
+			accept-connection
 		else
-			print -u $connection ''
+			handle-request $fd
 		fi
-		# close fd
-		exec {connection}>&-
-		print "connection closed, fd: $connection"
-		# clear input buffer
-		zpty -w z $'\n'
 	done
 done