SUBPROCESS: Simple subprocess runner/reaper

1 files changed, 275 insertions, 0 deletions

diff --git a/luxio/subprocess.lua b/luxio/subprocess.lua
new file mode 100644
index 0000000..f064593
--- /dev/null
+++ b/luxio/subprocess.lua
@@ -0,0 +1,275 @@
+-- Light Unix I/O for Lua
+-- Copyright 2012 Daniel Silverstone <dsilvers@digital-scurf.org>
+--
+-- Distributed under the same terms as Lua itself (MIT).
+--
+-- A wrapper for managing sub processes
+-------------------------------------------------------------------------------
+--
+-- sp = require "luxio.subprocess"
+--
+--
+-- proc = sp.spawn {
+--     "sleep", "10",
+--     exe = "/usr/bin/sleep"
+-- }
+--
+-- how, why = proc:wait(noblock)
+--     noblock, if true, won't block until proc exits.
+--
+-- You can redirect stdin/stdout/stderr for subprocesses
+-- by passing FDs or simple file handles as std{in,out,err} in
+-- the arguments to run or spawn.
+--
+-- If you provide sp.PIPE then a pipe will be created for you
+-- and the proc's std{in,out,err} will be the other end of it.
+--
+
+local l = require 'luxio'
+local sio = require 'luxio.simple'
+
+local fork = l.fork
+local waitpid = l.waitpid
+local exec = l.exec
+local execp = l.execp
+local pipe = l.pipe
+local close = l.close
+local dup = l.dup
+local dup2 = l.dup2
+local fcntl = l.fcntl
+local bclear = l.bit.bclear
+local wrap_fd = sio.wrap_fd
+
+local F_SETFD = l.F_SETFD
+local F_GETFD = l.F_GETFD
+
+local WIFEXITED = l.WIFEXITED
+local WEXITSTATUS = l.WEXITSTATUS
+local WIFSIGNALLED = l.WIFSIGNALLED
+local WTERMSIG = l.WTERMSIG
+local WNOHANG = l.WNOHANG
+
+local assert = assert
+local unpack = unpack
+local pairs = pairs
+local error = error
+
+l = nil
+sio = nil
+
+local _PIPE = {}
+
+local function _wait(proc, nohang)
+   local pid, status = waitpid(proc.pid, nohang and WNOHANG or 0)
+   if pid == -1 then
+      return "error", status
+   end
+   if pid == 0 then
+      return "running", 0
+   end
+
+   if WIFEXITED(status) ~= 0 then
+      return "exit", WEXITSTATUS(status)
+   end
+
+   if WIFSIGNALLED(status) ~= 0 then
+      return "signal", WTERMSIG(status)
+   end
+
+   return "unknown", status
+end
+
+local function _spawn(t)
+   -- Stage one, normalise our process
+   local proc = {args={}}
+   for k, v in pairs(t) do
+      proc.args[k] = v
+   end
+
+   if proc.args[0] == nil then
+      if proc.args.exe then 
+	 proc.args[0] = proc.args.exe
+      else
+	 proc.args[0] = proc.args[1]
+	 if proc.args[0]:match("^/") then
+	    proc.args.exe = proc.args[0]
+	 end
+      end
+   end
+
+   -- Choose which exec function to use, based on whether or not
+   -- we have an explicit exe.
+   local exec_fn = proc.args.exe and exec or execp
+   -- And prepare a local copy of the exe based on the args.
+   local exe = proc.args.exe or proc.args[0]
+
+   assert(exe, "Unable to continue, nothing to run")
+
+   -- Simple verification of the FDs.
+   local function check(v)
+      if proc.args[v] then
+	 -- Something to do for handle
+	 if proc.args[v] == _PIPE then
+	    -- For now, nothing, later we'll handle this
+	 elseif type(proc.args[v]) == "table" then
+	    -- Simple handle, extract the FD
+	    proc.args[v] = proc.args[v].fd
+	 elseif type(proc.args[v]) == "number" then
+	    -- FD, nothing to do for now
+	 else
+	    assert(false, v .. " isn't a PIPE, Simple Handle or FD")
+	 end
+      end
+   end
+
+   check "stdin"
+   check "stdout"
+   check "stderr"
+   
+   
+
+   -- Beyond here, if we get errors, we *MUST* clean up FDs
+   local parent_close = {}
+   local child_close = {}
+
+   local function _assert(m,v)
+      if not m then
+	 for fd in pairs(parent_close) do
+	    close(fd)
+	 end
+	 for fd in pairs(child_close) do
+	    close(fd)
+	 end
+	 error(v)
+      end
+   end
+
+   local function do_pipe(v)
+      if proc.args[v] == _PIPE then
+	 local fds = {}
+	 local res, errno = pipe(fds)
+	 _assert(res == 0, "Unable to pipe()")
+	 proc.args[v] = fds[2]
+	 proc[v] = fds[1]
+	 parent_close[proc.args[v]] = true
+	 child_close[proc[v]] = true
+      end
+   end
+   do_pipe "stdin"
+   do_pipe "stdout"
+   do_pipe "stderr"
+
+   if proc.stdin then
+      -- Swap stdin over.
+      proc.stdin, proc.args.stdin = proc.args.stdin, proc.stdin
+      parent_close[proc.stdin] = nil
+      parent_close[proc.args.stdin] = true
+      child_close[proc.stdin] = true
+      child_close[proc.args.stdin] = nil
+   end
+
+   proc.args.stdin = proc.args.stdin or 0
+   proc.args.stdout = proc.args.stdout or 1
+   proc.args.stderr = proc.args.stderr or 2
+
+   -- Finally, let's fork
+
+   local pid, errno = fork()
+
+   _assert(pid > -1, "Unable to fork()")
+
+   if pid == 0 then
+      local function check(m)
+	 if m == -1 then
+	    _exit(1)
+	 end
+	 return m
+      end
+
+      -- Set up the FDs.
+      if proc.args.stdin ~= 0 then
+	 check(dup2(proc.args.stdin, 0))
+      end
+      if proc.args.stdout ~= 1 then
+	 if proc.args.stderr == 1 then
+	    proc.args.old_stdout = check(dup(1))
+	 end
+	 check(dup2(proc.args.stdout, 1))
+      end
+      if proc.args.stderr ~= 2 then
+	 if proc.args.stderr == 1 and proc.args.old_stdout then
+	    check(dup2(proc.args.old_stdout, 2))
+	 else
+	    check(dup2(proc.args.stderr, 2))
+	 end
+      end
+      if proc.args.old_stdout then
+	 close(proc.args.old_stdout)
+      end
+      -- FDs 0, 1, 2 initialised, ensure they won't close on exec
+      fcntl(0, F_SETFD, bclear(fcntl(0, F_GETFD), 1))
+      fcntl(0, F_SETFD, bclear(fcntl(1, F_GETFD), 1))
+      fcntl(0, F_SETFD, bclear(fcntl(2, F_GETFD), 1))
+
+      -- Finally, close the parent halves of any pipes
+      for fd in pairs(child_close) do
+	 close(fd)
+      end
+
+      -- Run the child process
+      exec_fn(exe, unpack(proc.args))
+      
+      -- Exec failed, exit without doing lua cleanups
+      _exit(1)
+   end
+
+   -- The parent
+   for fd in pairs(parent_close) do
+      close(fd)
+   end
+
+   proc.pid = pid
+
+   proc.wait = _wait
+
+   return proc
+end
+
+local function _spawn_simple(t)
+   local t_copy = {}
+   for k, v in pairs(t) do
+      t_copy[k] = v
+   end
+
+   local piping = nil
+
+   if type(t_copy.stdin) == "string" then
+      piping = t_copy.stdin
+      t_copy.stdin = _PIPE
+   end
+
+   local proc = _spawn(t_copy)
+
+   local function doit(v)
+      if proc[v] then
+	 proc[v] = wrap_fd(proc[v])
+      end
+   end
+   doit "stdin"
+   doit "stdout"
+   doit "stderr"
+
+   if piping then
+      proc.stdin:write(piping)
+      proc.stdin:close()
+      proc.stdin = nil
+   end
+
+   return proc
+end
+
+return {
+   spawn = _spawn,
+   spawn_simple = _spawn_simple,
+   PIPE = _PIPE,
+}