Skip to content

Commit

Permalink
refactor(async): Add a wrapper around zsh-async
Browse files Browse the repository at this point in the history
  • Loading branch information
denysdovhan committed Sep 23, 2022
1 parent 1b7c053 commit 772a3d8
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 48 deletions.
75 changes: 43 additions & 32 deletions lib/core.zsh
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,6 @@
# Tools for loading sections, building sections and invoking the renderer
# ------------------------------------------------------------------------------

# Unique array of async jobs
typeset -ahU SPACESHIP_JOBS=()

# Loads the sections from files and functions
# USAGE:
# spaceship::core::load_sections
Expand All @@ -32,17 +29,14 @@ spaceship::core::load_sections() {
done

if $load_async; then
(( ASYNC_INIT_DONE )) || {
builtin source "$SPACESHIP_ROOT/async.zsh"
spaceship::precompile "$SPACESHIP_ROOT/async.zsh"
}
spaceship::worker::load
fi
}

# Iterate over sections, start async jobs and store results in cache
# USAGE:
# spaceship::core::build_cache
spaceship::core::build_cache() {
# spaceship::core::start
spaceship::core::start() {
# Clear the cache before every render
spaceship::cache::clear

Expand All @@ -56,33 +50,51 @@ spaceship::core::build_cache() {
# spaceship::core::async_callback
spaceship::core::async_callback() {
local job="$1" ret="$2" output="$3" exec_time="$4" err="$5" has_next="$6"
local section

# ignore the async evals used to alter worker environment
if [[ "${job}" == "[async/eval]" ]] \
|| [[ "${job}" == ";" ]] \
|| [[ "${job}" == "[async]" ]]; then
# FIXME: Restart the async job if it failed
return
fi

section="${job#"spaceship_"}" # TODO: Move spaceship_ to a constant

SPACESHIP_JOBS=("${(@)SPACESHIP_JOBS:#${section}}")
local section="${job#"spaceship_"}" # TODO: Move spaceship_ to a constant

# Notify the worker that the job is done
spaceship::worker::callback "$@"

case $job in
"[async]")
# Handle all the errors that could indicate a crashed async worker.
# See zsh-async documentation for the definition of the exit codes.
if (( code == 2 )) || (( code == 3 )) || (( code == 130 )); then
# Our worker died unexpectedly, try to recover immediately.
spaceship::worker::init
spaceship::core::start
return
fi
;;
"[async/eval]")
if (( code )); then
# Looks like eval failed, rerun async tasks just in case.
spaceship::core::start
return
fi
;;
";")
# Ignore the async evals used to alter worker environment
return
;;
*)
# Hanlde regular successfully finished jobs
# Skip prompt re-rendering if section is empty
if [[ "$(spaceship::cache::get $section)" == "$output" ]]; then
return
fi

# Update section cache
spaceship::cache::set "$section" "$output"
;;
esac

# Refresh async section when the last async job has finished
if [[ "${#SPACESHIP_JOBS}" -eq 0 ]]; then
spaceship::core::refresh_section "async"
spaceship::core::render
fi

# Skip prompt re-rendering if section is empty
if [[ "$(spaceship::cache::get $section)" == "$output" ]]; then
return
fi

spaceship::cache::set "$section" "$output"

if [[ "$has_next" == 0 ]]; then
spaceship::core::render
fi
Expand Down Expand Up @@ -112,8 +124,7 @@ spaceship::core::refresh_section() {
fi

if spaceship::is_section_async "$section" && [[ -z $sync ]]; then
SPACESHIP_JOBS+=("$section")
async_job "spaceship" "spaceship_${section}"
spaceship::worker::run "spaceship_$section"
else
spaceship::cache::set "$section" "$(spaceship_$section)"
fi
Expand Down Expand Up @@ -150,5 +161,5 @@ spaceship::core::render() {
# .reset-prompt: bypass the zsh-syntax-highlighting wrapper
# https://github.com/sorin-ionescu/prezto/issues/1026
# https://github.com/zsh-users/zsh-autosuggestions/issues/107#issuecomment-183824034
zle .reset-prompt && zle -R
zle && zle .reset-prompt && zle -R
}
20 changes: 4 additions & 16 deletions lib/hooks.zsh
Original file line number Diff line number Diff line change
Expand Up @@ -39,18 +39,10 @@ prompt_spaceship_precmd() {
spaceship_exec_time_stop

# Restarts the async worker, in order to get an update-to-date shell environment
if spaceship::is_prompt_async; then
SPACESHIP_JOBS=()
# restart worker
async_stop_worker "spaceship"
async_start_worker "spaceship" -n -u
# setopt before call register to avoid callback by async_worker_eval
async_worker_eval "spaceship" 'setopt extendedglob'
async_register_callback "spaceship" "spaceship::core::async_callback"
fi
spaceship::worker::init

# Start building cache from sections
spaceship::core::build_cache
spaceship::core::start

# Initiate the first render
spaceship::populate
Expand All @@ -59,19 +51,15 @@ prompt_spaceship_precmd() {
# A hook right before the command is started executing
prompt_spaceship_preexec() {
# Stop running prompt async jobs
if spaceship::is_prompt_async; then
async_flush_jobs "spaceship"
fi
spaceship::worker::flush

# Start measuring exec_time right before executing the command
spaceship_exec_time_start
}

# A hook after changing the working directory
prompt_spaceship_chpwd() {
if spaceship::is_prompt_async; then
async_worker_eval "spaceship" 'cd' "$PWD"
fi
spaceship::worker::eval builtin cd -q $PWD

# Restart execution time recording once dir is changed
spaceship_exec_time_start
Expand Down
67 changes: 67 additions & 0 deletions lib/worker.zsh
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
# ------------------------------------------------------------------------------
# WORKER
# Spaceship wrapper around zsh-async
# ------------------------------------------------------------------------------

# Unique array of async jobs
typeset -ahU SPACESHIP_JOBS=()

# Load zsh-async if not loaded yet
spaceship::worker::load() {
if ! (( ASYNC_INIT_DONE )); then
builtin source "$SPACESHIP_ROOT/async.zsh"
spaceship::precompile "$SPACESHIP_ROOT/async.zsh"
fi
}

# Lower worker priority to avoid slowing down the prompt
spaceship::worker::renice() {
if command -v renice >/dev/null; then
command renice +15 -p $$
fi

if command -v ionice >/dev/null; then
command ionice -c 3 -p $$
fi
}

# This should be called to in callback to update the job counter
spaceship::worker::callback() {
SPACESHIP_JOBS=("${(@)SPACESHIP_JOBS:#${1}}")
}

# Start the worker and prepare the environment
spaceship::worker::init() {
if spaceship::is_prompt_async; then
SPACESHIP_JOBS=()
# restart worker
async_stop_worker "spaceship"
async_start_worker "spaceship" -n -u
# setopt before call register to avoid callback by async_worker_eval
async_worker_eval "spaceship" setopt extendedglob
async_worker_eval "spaceship" spaceship::worker::renice
async_register_callback "spaceship" spaceship::core::async_callback
fi
}

# Flush jobs for stopped worker
spaceship::worker::flush() {
if spaceship::is_prompt_async; then
async_flush_jobs "spaceship"
fi
}

# Eval command inside the worker
spaceship::worker::eval() {
if spaceship::is_prompt_async; then
async_worker_eval "spaceship" "$@"
fi
}

# Run a job in a worker
spaceship::worker::run() {
if spaceship::is_prompt_async; then
SPACESHIP_JOBS+=("$1")
async_job "spaceship" "$@"
fi
}
1 change: 1 addition & 0 deletions spaceship.zsh
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ SPACESHIP_PROMPT_DEFAULT_SUFFIX="${SPACESHIP_PROMPT_DEFAULT_SUFFIX=" "}"
SPACESHIP_LIBS=(
"lib/utils.zsh" # General porpuse utils
"lib/cache.zsh" # Cache utils
"lib/worker.zsh" # Async worker
"lib/hooks.zsh" # Zsh hooks
"lib/section.zsh" # Section utils
"lib/core.zsh" # Core functions for loading and rendering
Expand Down

0 comments on commit 772a3d8

Please sign in to comment.