From 42b22ecf7debe2e5ccafed0c1a282c0c188fe338 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Yhuel?= <loic.yhuel@gmail.com>
Date: Mon, 25 Mar 2024 23:19:08 +0100
Subject: [PATCH] fix(async): avoid blocking the shell while waiting

Due to "builtin echo $handler", _omz_async_callback was called immediately, as there was data to
read on the pipe. Then the shell was blocked until the end of the async prompt handler.
We now use associative array to get the handler name from the fd.

For an unknown reason, the async child can block when launching a subshell, especially if there are
several handlers registered.
Using a function to set the return code avoids the issue, and is much faster.

Fixes #12290
---
 lib/async_prompt.zsh | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/lib/async_prompt.zsh b/lib/async_prompt.zsh
index 384e49d33..f13d6d544 100644
--- a/lib/async_prompt.zsh
+++ b/lib/async_prompt.zsh
@@ -44,7 +44,7 @@ function _omz_register_handler {
 # Set up async handlers and callbacks
 function _omz_async_request {
   local -i ret=$?
-  typeset -gA _OMZ_ASYNC_FDS _OMZ_ASYNC_PIDS _OMZ_ASYNC_OUTPUT
+  typeset -gA _OMZ_ASYNC_FDS _OMZ_ASYNC_PIDS _OMZ_ASYNC_OUTPUT _OMZ_ASYNC_HANDLERS
 
   # executor runs a subshell for all async requests based on key
   local handler
@@ -82,16 +82,16 @@ function _omz_async_request {
     exec {fd}< <(
       # Tell parent process our PID
       builtin echo ${sysparams[pid]}
-      # Store handler name for callback
-      builtin echo $handler
       # Set exit code for the handler if used
-      (exit $ret)
+      () { return $ret }
       # Run the async function handler
       $handler
     )
 
     # Save FD for handler
     _OMZ_ASYNC_FDS[$handler]=$fd
+    # Save handler name for callback
+    _OMZ_ASYNC_HANDLERS[$fd]=$handler
 
     # There's a weird bug here where ^C stops working unless we force a fork
     # See https://github.com/zsh-users/zsh-autosuggestions/issues/364
@@ -114,9 +114,8 @@ function _omz_async_callback() {
   local err=$2  # Second arg will be passed in case of error
 
   if [[ -z "$err" || "$err" == "hup" ]]; then
-    # Get handler name from first line
-    local handler
-    read handler <&$fd
+    # Get handler name from fd
+    local handler=${_OMZ_ASYNC_HANDLERS[$fd]}
 
     # Store old output which is supposed to be already printed
     local old_output="${_OMZ_ASYNC_OUTPUT[$handler]}"
@@ -137,6 +136,9 @@ function _omz_async_callback() {
   # Always remove the handler
   zle -F "$fd"
 
+  # Remove the fd => handle name association
+  unset '_OMZ_ASYNC_HANDLERS[$fd]'
+
   # Unset global FD variable to prevent closing user created FDs in the precmd hook
   _OMZ_ASYNC_FDS[$handler]=-1
   _OMZ_ASYNC_PIDS[$handler]=-1