wdye: enable waiting for processes

3 files changed, 69 insertions, 1 deletions

diff --git a/tools/wdye/test.lua b/tools/wdye/test.lua
index 5e2fe36..c255c72 100644
--- a/tools/wdye/test.lua
+++ b/tools/wdye/test.lua
@@ -18,7 +18,14 @@ expect(cat:regex {"A(.*)3", nocase=true, function (p)
 	assert(p[1] == "bc12", "wrong regex group #1")
 end})
 
+assert(not cat:wait (true), "process reports exiting early")
+
 -- Send EOF (^D), test method chaining.
 cat:send("Closing...\r"):send("\004")
 local v = expect(cat:eof {true},
 	cat:default {.5, function (p) error "expected EOF, got a timeout" end})
+
+local s1, exit, signal = cat:wait ()
+assert(s1 == 0 and exit == 0 and not signal, "unexpected exit status")
+local s2 = cat:wait (true)
+assert(s1 == s2, "exit status not remembered")
diff --git a/tools/wdye/wdye.adoc b/tools/wdye/wdye.adoc
index 6282af2..d901400 100644
--- a/tools/wdye/wdye.adoc
+++ b/tools/wdye/wdye.adoc
@@ -94,6 +94,16 @@ PROCESS:eof {[notransfer=true, ] [value1, ...]}
 Returns a new end-of-file _pattern_, which matches the entire read buffer
 contents once the child process closes the terminal.
 
+PROCESS:wait ([nowait])
+~~~~~~~~~~~~~~~~~~~~~~~
+Waits for the program to terminate, and returns three values:
+a combined status as used by `$?` in shells,
+an exit status, and a termination signal number.
+One of the latter two values will be _nil_, as appropriate.
+
+When the *nowait* option is _true_, the function returns immediately.
+If the process hasn't terminated yet, the function then returns no values.
+
 PROCESS:default {[timeout, ] [notransfer=true, ] [value1, ...]}
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Returns a new _pattern_ combining *wdye.timeout* with *eof*.
diff --git a/tools/wdye/wdye.c b/tools/wdye/wdye.c
index b801c89..ece3038 100644
--- a/tools/wdye/wdye.c
+++ b/tools/wdye/wdye.c
@@ -431,9 +431,10 @@ static luaL_Reg xlua_pattern_table[] =
 struct process
 {
 	int terminal_fd;                    ///< Process stdin/stdout/stderr
-	pid_t pid;                          ///< Process ID
+	pid_t pid;                          ///< Process ID or -1 if collected
 	int ref_term;                       ///< Terminal information
 	struct str buffer;                  ///< Terminal input buffer
+	int status;                         ///< Process status iff pid is -1
 };
 
 static struct process *
@@ -598,6 +599,54 @@ xlua_process_default (lua_State *L)
 	return 1;
 }
 
+static int
+xlua_process_wait (lua_State *L)
+{
+	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
+	bool nowait = luaL_opt(L, lua_toboolean, 2, false);
+	if (lua_gettop (L) > 2)
+		return luaL_error (L, "too many arguments");
+
+	int status = self->status;
+restart:
+	if (self->pid != -1)
+	{
+		int options = 0;
+		if (nowait)
+			options |= WNOHANG;
+
+		pid_t pid = waitpid (self->pid, &status, options);
+		if (!pid)
+			return 0;
+
+		if (pid < 0)
+		{
+			if (errno == EINTR)
+				goto restart;
+			return luaL_error (L, "waitpid: %s", strerror (errno));
+		}
+
+		// We lose the ability to reliably kill the whole process group.
+		self->status = status;
+		self->pid = -1;
+	}
+	if (WIFEXITED (status))
+	{
+		lua_pushinteger (L, WEXITSTATUS (status));
+		lua_pushinteger (L, WEXITSTATUS (status));
+		lua_pushnil (L);
+		return 3;
+	}
+	if (WIFSIGNALED (status))
+	{
+		lua_pushinteger (L, 128 + WTERMSIG (status));
+		lua_pushnil (L);
+		lua_pushinteger (L, WTERMSIG (status));
+		return 3;
+	}
+	return 0;
+}
+
 static bool
 process_feed (struct process *self)
 {
@@ -632,6 +681,7 @@ static luaL_Reg xlua_process_table[] =
 	{ "exact",      xlua_process_exact    },
 	{ "eof",        xlua_process_eof      },
 	{ "default",    xlua_process_default  },
+	{ "wait",       xlua_process_wait     },
 	{ NULL,         NULL                  }
 };
 
@@ -851,6 +901,7 @@ spawn_protected (lua_State *L)
 		execvpe (ctx->argv.vector[0], ctx->argv.vector, ctx->envv.vector);
 		print_error ("failed to spawn %s: %s",
 			ctx->argv.vector[0], strerror (errno));
+		// Or we could figure out when exactly to use statuses 126 and 127.
 		_exit (EXIT_FAILURE);
 	}