Smooth rough edges of daemon command-line behavior

1 files 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))