From 0d644b4350f3e6cbb676360b0fc6a9e448137092 Mon Sep 17 00:00:00 2001 From: Adam Chlipala Date: Fri, 30 Aug 2019 16:32:14 -0400 Subject: Smooth rough edges of daemon command-line behavior --- src/main.mlton.sml | 247 +++++++++++++++++++++++++++-------------------------- 1 file changed, 128 insertions(+), 119 deletions(-) diff --git a/src/main.mlton.sml b/src/main.mlton.sml index bfa40265..e9317d46 100644 --- a/src/main.mlton.sml +++ b/src/main.mlton.sml @@ -255,7 +255,8 @@ fun oneRun args = ] val () = case args of - ["daemon", "stop"] => OS.Process.exit OS.Process.success + ["daemon", "stop"] => (OS.FileSys.remove socket handle OS.SysErr _ => (); + OS.Process.exit OS.Process.success) | _ => () val sources = parse_flags (flag_info ()) args @@ -318,127 +319,135 @@ fun send (sock, s) = send (sock, String.extract (s, n, NONE)) end -val () = (Globals.setResetTime (); - case CommandLine.arguments () of - ["daemon", "start"] => - (case Posix.Process.fork () of - SOME _ => () - | NONE => - let - val () = Elaborate.incremental := true - val listen = UnixSock.Strm.socket () - - fun loop () = - let - val (sock, _) = Socket.accept listen - - fun loop' (buf, args) = - let - val s = if CharVector.exists (fn ch => ch = #"\n") buf then - "" - else - MLton.CharVector.fromPoly (Vector.map (chr o Word8.toInt) (MLton.Word8Vector.toPoly (Socket.recvVec (sock, 1024)))) - val s = buf ^ s - val (befor, after) = Substring.splitl (fn ch => ch <> #"\n") (Substring.full s) - in - if Substring.isEmpty after then - loop' (s, args) - else - let - val cmd = Substring.string befor - val rest = Substring.string (Substring.slice (after, 1, NONE)) - in - case cmd of - "" => - (case args of - ["stop", "daemon"] => - (((Socket.close listen; - OS.FileSys.remove socket) handle OS.SysErr _ => ()); - OS.Process.exit OS.Process.success) - | _ => - let - val success = (oneRun (rev args)) - handle ex => (print "unhandled exception:\n"; - print (General.exnMessage ex ^ "\n"); - OS.Process.failure) - in - TextIO.flushOut TextIO.stdOut; - TextIO.flushOut TextIO.stdErr; - send (sock, if OS.Process.isSuccess success then - "\001" - else - "\002") - end) - | _ => loop' (rest, cmd :: args) - end - end handle OS.SysErr _ => () - - fun redirect old = - Posix.IO.dup2 {old = valOf (Posix.FileSys.iodToFD (Socket.ioDesc sock)), - new = old} - - val oldStdout = Posix.IO.dup Posix.FileSys.stdout - val oldStderr = Posix.IO.dup Posix.FileSys.stderr - in - (* Redirect the daemon's output to the socket. *) - redirect Posix.FileSys.stdout; - redirect Posix.FileSys.stderr; - - loop' ("", []); - Socket.close sock; - - Posix.IO.dup2 {old = oldStdout, new = Posix.FileSys.stdout}; - Posix.IO.dup2 {old = oldStderr, new = Posix.FileSys.stderr}; - Posix.IO.close oldStdout; - Posix.IO.close oldStderr; - - Settings.reset (); - MLton.GC.pack (); - loop () - end - in - OS.Process.atExit (fn () => OS.FileSys.remove socket); - Socket.bind (listen, UnixSock.toAddr socket); - Socket.listen (listen, 1); - loop () - end) - | args => +fun startDaemon () = + if OS.FileSys.access (socket, []) then + (print ("It looks like a daemon is already listening in this directory,\n" + ^ "though it's possible a daemon died without cleaning up its socket.\n"); + OS.Process.exit OS.Process.failure) + else case Posix.Process.fork () of + SOME _ => () + | NONE => let - val sock = UnixSock.Strm.socket () + val () = Elaborate.incremental := true + val listen = UnixSock.Strm.socket () - fun wait () = + fun loop () = let - val v = Socket.recvVec (sock, 1024) - in - if Word8Vector.length v = 0 then - OS.Process.failure - else + val (sock, _) = Socket.accept listen + + fun loop' (buf, args) = let - val s = MLton.CharVector.fromPoly (Vector.map (chr o Word8.toInt) (MLton.Word8Vector.toPoly v)) - val last = Word8Vector.sub (v, Word8Vector.length v - 1) - val (rc, s) = if last = Word8.fromInt 1 then - (SOME OS.Process.success, String.substring (s, 0, size s - 1)) - else if last = Word8.fromInt 2 then - (SOME OS.Process.failure, String.substring (s, 0, size s - 1)) - else - (NONE, s) + val s = if CharVector.exists (fn ch => ch = #"\n") buf then + "" + else + MLton.CharVector.fromPoly (Vector.map (chr o Word8.toInt) (MLton.Word8Vector.toPoly (Socket.recvVec (sock, 1024)))) + val s = buf ^ s + val (befor, after) = Substring.splitl (fn ch => ch <> #"\n") (Substring.full s) in - print s; - case rc of - NONE => wait () - | SOME rc => rc - end - end handle OS.SysErr _ => OS.Process.failure + if Substring.isEmpty after then + loop' (s, args) + else + let + val cmd = Substring.string befor + val rest = Substring.string (Substring.slice (after, 1, NONE)) + in + case cmd of + "" => + (case args of + ["stop", "daemon"] => + (((Socket.close listen; + OS.FileSys.remove socket) handle OS.SysErr _ => ()); + OS.Process.exit OS.Process.success) + | _ => + let + val success = (oneRun (rev args)) + handle ex => (print "unhandled exception:\n"; + print (General.exnMessage ex ^ "\n"); + OS.Process.failure) + in + TextIO.flushOut TextIO.stdOut; + TextIO.flushOut TextIO.stdErr; + send (sock, if OS.Process.isSuccess success then + "\001" + else + "\002") + end) + | _ => loop' (rest, cmd :: args) + end + end handle OS.SysErr _ => () + + fun redirect old = + Posix.IO.dup2 {old = valOf (Posix.FileSys.iodToFD (Socket.ioDesc sock)), + new = old} + + val oldStdout = Posix.IO.dup Posix.FileSys.stdout + val oldStderr = Posix.IO.dup Posix.FileSys.stderr + in + (* Redirect the daemon's output to the socket. *) + redirect Posix.FileSys.stdout; + redirect Posix.FileSys.stderr; + + loop' ("", []); + Socket.close sock; + + Posix.IO.dup2 {old = oldStdout, new = Posix.FileSys.stdout}; + Posix.IO.dup2 {old = oldStderr, new = Posix.FileSys.stderr}; + Posix.IO.close oldStdout; + Posix.IO.close oldStderr; + + Settings.reset (); + MLton.GC.pack (); + loop () + end in - if Socket.connectNB (sock, UnixSock.toAddr socket) - orelse not (List.null (#wrs (Socket.select {rds = [], - wrs = [Socket.sockDesc sock], - exs = [], - timeout = SOME (Time.fromSeconds 1)}))) then - (app (fn arg => send (sock, arg ^ "\n")) args; - send (sock, "\n"); - OS.Process.exit (wait ())) - else - (OS.FileSys.remove socket; - raise OS.SysErr ("", NONE)) - end handle OS.SysErr _ => OS.Process.exit (oneRun args)) + OS.Process.atExit (fn () => OS.FileSys.remove socket); + Socket.bind (listen, UnixSock.toAddr socket); + Socket.listen (listen, 1); + loop () + end + +fun oneCommandLine args = + let + val sock = UnixSock.Strm.socket () + + fun wait () = + let + val v = Socket.recvVec (sock, 1024) + in + if Word8Vector.length v = 0 then + OS.Process.failure + else + let + val s = MLton.CharVector.fromPoly (Vector.map (chr o Word8.toInt) (MLton.Word8Vector.toPoly v)) + val last = Word8Vector.sub (v, Word8Vector.length v - 1) + val (rc, s) = if last = Word8.fromInt 1 then + (SOME OS.Process.success, String.substring (s, 0, size s - 1)) + else if last = Word8.fromInt 2 then + (SOME OS.Process.failure, String.substring (s, 0, size s - 1)) + else + (NONE, s) + in + print s; + case rc of + NONE => wait () + | SOME rc => rc + end + end handle OS.SysErr _ => OS.Process.failure + in + if Socket.connectNB (sock, UnixSock.toAddr socket) + orelse not (List.null (#wrs (Socket.select {rds = [], + wrs = [Socket.sockDesc sock], + exs = [], + timeout = SOME (Time.fromSeconds 1)}))) then + (app (fn arg => send (sock, arg ^ "\n")) args; + send (sock, "\n"); + wait ()) + else + (OS.FileSys.remove socket; + raise OS.SysErr ("", NONE)) + end handle OS.SysErr _ => oneRun args + +val () = (Globals.setResetTime (); + case CommandLine.arguments () of + ["daemon", "start"] => startDaemon () + | args => OS.Process.exit (oneCommandLine args)) -- cgit v1.2.3