eval: Implement `jobwait()` vimscript function

3 files changed, 127 insertions, 4 deletions

diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index 05a452c39d..e4acaa0c53 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -448,6 +448,7 @@ typedef struct {
   int refcount;
   ufunc_T *on_stdout, *on_stderr, *on_exit;
   dict_T *self;
+  int *status_ptr;
 } TerminalJobData;
 
 
@@ -470,6 +471,7 @@ typedef struct {
 #define JobEventFreer(x)
 KMEMPOOL_INIT(JobEventPool, JobEvent, JobEventFreer)
 static kmempool_t(JobEventPool) *job_event_pool = NULL;
+static bool defer_job_callbacks = true;
 
 /*
  * Initialize the global and v: variables.
@@ -6537,6 +6539,7 @@ static struct fst {
   {"jobsend",         2, 2, f_jobsend},
   {"jobstart",        1, 2, f_jobstart},
   {"jobstop",         1, 1, f_jobstop},
+  {"jobwait",         1, 2, f_jobwait},
   {"join",            1, 2, f_join},
   {"keys",            1, 1, f_keys},
   {"last_buffer_nr",  0, 0, f_last_buffer_nr},  /* obsolete */
@@ -10841,6 +10844,105 @@ static void f_jobstop(typval_T *argvars, typval_T *rettv)
   rettv->vval.v_number = 1;
 }
 
+// "jobwait(ids[, timeout])" function
+static void f_jobwait(typval_T *argvars, typval_T *rettv)
+{
+  rettv->v_type = VAR_NUMBER;
+  rettv->vval.v_number = 0;
+
+  if (check_restricted() || check_secure()) {
+    return;
+  }
+
+  if (argvars[0].v_type != VAR_LIST || (argvars[1].v_type != VAR_NUMBER
+        && argvars[1].v_type != VAR_UNKNOWN)) {
+    EMSG(_(e_invarg));
+    return;
+  }
+
+  list_T *args = argvars[0].vval.v_list;
+  list_T *rv = list_alloc();
+
+  // must temporarily disable job event deferring so the callbacks are
+  // processed while waiting.
+  defer_job_callbacks = false;
+  // For each item in the input list append an integer to the output list. -3
+  // is used to represent an invalid job id, -2 is for a interrupted job and
+  // -1 for jobs that were skipped or timed out.
+  for (listitem_T *arg = args->lv_first; arg != NULL; arg = arg->li_next) {
+    Job *job = NULL;
+    if (arg->li_tv.v_type != VAR_NUMBER
+        || !(job = job_find(arg->li_tv.vval.v_number))
+        || !is_user_job(job)) {
+      list_append_number(rv, -3);
+    } else {
+      TerminalJobData *data = job_data(job);
+      // append the list item and set the status pointer so we'll collect the
+      // status code when the job exits
+      list_append_number(rv, -1);
+      data->status_ptr = &rv->lv_last->li_tv.vval.v_number;
+    }
+  }
+
+  int remaining = -1;
+  uint64_t before = 0;
+  if (argvars[1].v_type == VAR_NUMBER && argvars[1].vval.v_number >= 0) {
+    remaining = argvars[1].vval.v_number;
+    before = os_hrtime();
+  }
+
+  for (listitem_T *arg = args->lv_first; arg != NULL; arg = arg->li_next) {
+    Job *job = NULL;
+    if (remaining == 0) {
+      // timed out
+      break;
+    }
+    if (arg->li_tv.v_type != VAR_NUMBER
+        || !(job = job_find(arg->li_tv.vval.v_number))
+        || !is_user_job(job)) {
+      continue;
+    }
+    TerminalJobData *data = job_data(job);
+    int status = job_wait(job, remaining);
+    if (status < 0) {
+      // interrupted or timed out, skip remaining jobs.
+      if (status == -2) {
+        // set the status so the user can distinguish between interrupted and
+        // skipped/timeout jobs.
+        *data->status_ptr = -2;
+      }
+      break;
+    }
+    if (remaining > 0) {
+      uint64_t now = os_hrtime();
+      remaining -= (int) ((now - before) / 1000000);
+      before = now;
+      if (remaining <= 0) {
+        break;
+      }
+    }
+  }
+
+  for (listitem_T *arg = args->lv_first; arg != NULL; arg = arg->li_next) {
+    Job *job = NULL;
+    if (arg->li_tv.v_type != VAR_NUMBER
+        || !(job = job_find(arg->li_tv.vval.v_number))
+        || !is_user_job(job)) {
+      continue;
+    }
+    TerminalJobData *data = job_data(job);
+    // remove the status pointer because the list may be freed before the
+    // job exits
+    data->status_ptr = NULL;
+  }
+  // restore defer flag
+  defer_job_callbacks = true;
+
+  rv->lv_refcount++;
+  rettv->v_type = VAR_LIST;
+  rettv->vval.v_list = rv;
+}
+
 /*
  * "join()" function
  */
@@ -19951,6 +20053,16 @@ static inline void free_term_job_data(TerminalJobData *data) {
   free(data);
 }
 
+static inline bool is_user_job(Job *job)
+{
+  if (!job) {
+    return false;
+  }
+
+  JobOptions *opts = job_opts(job);
+  return opts->exit_cb == on_job_exit;
+}
+
 // vimscript job callbacks must be executed on Nvim main loop
 static inline void push_job_event(Job *job, ufunc_T *callback,
     const char *type, char *data, size_t count, int status)
@@ -19990,7 +20102,7 @@ static inline void push_job_event(Job *job, ufunc_T *callback,
   event_push((Event) {
     .handler = on_job_event,
     .data = event_data
-  }, true);
+  }, defer_job_callbacks);
 }
 
 static void on_job_stdout(RStream *rstream, void *job, bool eof)
@@ -20039,6 +20151,10 @@ static void on_job_exit(Job *job, int status, void *d)
         _("\r\n[Program exited, press any key to close]"));
   }
 
+  if (data->status_ptr) {
+    *data->status_ptr = status;
+  }
+
   push_job_event(job, data->on_exit, "exit", NULL, 0, status);
 }
 
diff --git a/src/nvim/os/job.c b/src/nvim/os/job.c
index ccd7891601..7d9a91a1b2 100644
--- a/src/nvim/os/job.c
+++ b/src/nvim/os/job.c
@@ -24,7 +24,6 @@
 // before we send SIGNAL to it
 #define TERM_TIMEOUT 1000000000
 #define KILL_TIMEOUT (TERM_TIMEOUT * 2)
-#define MAX_RUNNING_JOBS 100
 #define JOB_BUFFER_SIZE 0xFFFF
 
 #define close_job_stream(job, stream, type)                                \
@@ -234,11 +233,12 @@ void job_stop(Job *job)
 /// @return returns the status code of the exited job. -1 if the job is
 ///         still running and the `timeout` has expired. Note that this is
 ///         indistinguishable from the process returning -1 by itself. Which
-///         is possible on some OS.
+///         is possible on some OS. Returns -2 if the job was interrupted.
 int job_wait(Job *job, int ms) FUNC_ATTR_NONNULL_ALL
 {
   // The default status is -1, which represents a timeout
   int status = -1;
+  bool interrupted = false;
 
   // Increase refcount to stop the job from being freed before we have a
   // chance to get the status.
@@ -251,6 +251,7 @@ int job_wait(Job *job, int ms) FUNC_ATTR_NONNULL_ALL
   // we'll assume that a user frantically hitting interrupt doesn't like
   // the current job. Signal that it has to be killed.
   if (got_int) {
+    interrupted = true;
     got_int = false;
     job_stop(job);
     if (ms == -1) {
@@ -265,7 +266,7 @@ int job_wait(Job *job, int ms) FUNC_ATTR_NONNULL_ALL
   if (job->refcount == 1) {
     // Job exited, collect status and manually invoke close_cb to free the job
     // resources
-    status = job->status;
+    status = interrupted ? -2 : job->status;
     job_close_streams(job);
     job_decref(job);
   } else {
@@ -357,6 +358,11 @@ void job_close_streams(Job *job)
   close_job_err(job);
 }
 
+JobOptions *job_opts(Job *job)
+{
+  return &job->opts;
+}
+
 /// Iterates the table, sending SIGTERM to stopped jobs and SIGKILL to those
 /// that didn't die from SIGTERM after a while(exit_timeout is 0).
 static void job_stop_timer_cb(uv_timer_t *handle)
diff --git a/src/nvim/os/job_defs.h b/src/nvim/os/job_defs.h
index 200cf75e59..7fee900ac0 100644
--- a/src/nvim/os/job_defs.h
+++ b/src/nvim/os/job_defs.h
@@ -5,6 +5,7 @@
 #include "nvim/os/rstream_defs.h"
 #include "nvim/os/wstream_defs.h"
 
+#define MAX_RUNNING_JOBS 100
 typedef struct job Job;
 
 /// Function called when the job reads data