From e0a0e2ed421d26fe3477f9ffe0735a7477dd0238 Mon Sep 17 00:00:00 2001
From: Calvin Rose <calsrose@gmail.com>
Date: Sun, 2 Mar 2025 10:33:56 -0600
Subject: [PATCH] Prevent multi-scheduling with os/proc-wait. (address #1562)

---
 src/core/os.c       | 12 +++++++-----
 test/suite-ev.janet | 24 ++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/core/os.c b/src/core/os.c
index c814ef3a5..3808e3a5b 100644
--- a/src/core/os.c
+++ b/src/core/os.c
@@ -541,11 +541,12 @@ static void janet_proc_wait_cb(JanetEVGenericMessage args) {
         proc->flags &= ~JANET_PROC_WAITING;
         janet_gcunroot(janet_wrap_abstract(proc));
         janet_gcunroot(janet_wrap_fiber(args.fiber));
-        if ((status != 0) && (proc->flags & JANET_PROC_ERROR_NONZERO)) {
-            JanetString s = janet_formatc("command failed with non-zero exit code %d", status);
-            janet_cancel(args.fiber, janet_wrap_string(s));
-        } else {
-            if (janet_fiber_can_resume(args.fiber)) {
+        uint32_t sched_id = (uint32_t) args.argi;
+        if (janet_fiber_can_resume(args.fiber) && args.fiber->sched_id == sched_id) {
+            if ((status != 0) && (proc->flags & JANET_PROC_ERROR_NONZERO)) {
+                JanetString s = janet_formatc("command failed with non-zero exit code %d", status);
+                janet_cancel(args.fiber, janet_wrap_string(s));
+            } else {
                 janet_schedule(args.fiber, janet_wrap_integer(status));
             }
         }
@@ -603,6 +604,7 @@ os_proc_wait_impl(JanetProc *proc) {
     memset(&targs, 0, sizeof(targs));
     targs.argp = proc;
     targs.fiber = janet_root_fiber();
+    targs.argi = (uint32_t) targs.fiber->sched_id;
     janet_gcroot(janet_wrap_abstract(proc));
     janet_gcroot(janet_wrap_fiber(targs.fiber));
     janet_ev_threaded_call(janet_proc_wait_subr, targs, janet_proc_wait_cb);
diff --git a/test/suite-ev.janet b/test/suite-ev.janet
index c773e960a..ee14f4688 100644
--- a/test/suite-ev.janet
+++ b/test/suite-ev.janet
@@ -526,4 +526,28 @@
 (assert (= maxconn connect-count))
 (:close s)
 
+# Cancel os/proc-wait with ev/deadline
+(let [p (os/spawn [;run janet "-e" "(os/sleep 4)"] :p)]
+  (var terminated-normally false)
+  (assert-error "deadline expired"
+                (ev/with-deadline 0.01
+                  (os/proc-wait p)
+                  (print "uhoh")
+                  (set terminated-normally true)))
+  (assert (not terminated-normally) "early termination failure")
+  # Without this kill, janet will wait the full 4 seconds for the subprocess to complete before exiting.
+  (assert-no-error "kill proc after wait failed" (os/proc-kill p)))
+
+# Cancel os/proc-wait with ev/deadline 2
+(let [p (os/spawn [;run janet "-e" "(os/sleep 0.1)"] :p)]
+  (var terminated-normally false)
+  (assert-error "deadline expired"
+                (ev/with-deadline 0.05
+                  (os/proc-wait p)
+                  (print "uhoh")
+                  (set terminated-normally true)))
+  (assert (not terminated-normally) "early termination failure 2")
+  (ev/sleep 0.15)
+  (assert (not terminated-normally) "early termination failure 3"))
+
 (end-suite)