[PATCH 7 of 7] cmdserver: add service that listens on unix domain socket and forks process

[Yuya Nishihara]

# HG changeset patch
# User Yuya Nishihara <yuya at tcha.org>
# Date 1412408810 -32400
#      Sat Oct 04 16:46:50 2014 +0900
# Node ID 9b04af3305559a96ee310e9ba6766de1c4ffa63a
# Parent  0368d1677c97ad896468ab123b05760f9e34591a
cmdserver: add service that listens on unix domain socket and forks process

Typical use case of 'unix' mode is a background hg daemon.

    $ hg serve --cmdserver unix --cwd / -a /tmp/hg-`id -u`.sock

Unlike 'pipe' mode in which parent process keeps stdio channel, 'unix' server
can be detached.  So clients can freely connect and disconnect from server,
saving Python start-up time.

It might be better to write "--cmdserver socket -a unix:/sockpath" instead
of "--cmdserver unix -a /sockpath" in case hgweb gets the ability to listen
on unix domain socket.

diff --git a/mercurial/commandserver.py b/mercurial/commandserver.py
--- a/mercurial/commandserver.py
+++ b/mercurial/commandserver.py
@@ -7,7 +7,7 @@
 
 from i18n import _
 import struct
-import sys, os
+import sys, os, errno, traceback, SocketServer
 import dispatch, encoding, util
 
 logfile = None
@@ -256,8 +256,59 @@ class pipeservice(object):
     def run(self):
         return self.server.serve()
 
+class _requesthandler(SocketServer.StreamRequestHandler):
+    def handle(self):
+        ui = self.server.ui
+        repo = self.server.repo
+        sv = server(ui, repo, self.rfile, self.wfile)
+        try:
+            try:
+                sv.serve()
+            # handle exceptions that may be raised by command server. most of
+            # known exceptions are caught by dispatch.
+            except util.Abort, inst:
+                ui.warn(_('abort: %s\n') % inst)
+            except IOError, inst:
+                if inst.errno != errno.EPIPE:
+                    raise
+            except KeyboardInterrupt:
+                pass
+        except: # re-raises
+            # also write traceback to error channel. otherwise client cannot
+            # see it because it is written to server's stderr by default.
+            traceback.print_exc(file=sv.cerr)
+            raise
+
+class unixservice(object):
+    """
+    Listens on unix domain socket and forks server per connection
+    """
+    def __init__(self, ui, repo, opts):
+        self.ui = ui
+        self.repo = repo
+        self.address = opts['address']
+        if not util.safehasattr(SocketServer, 'UnixStreamServer'):
+            raise util.Abort(_('unsupported platform'))
+        if not self.address:
+            raise util.Abort(_('no socket path specified with --address'))
+
+    def init(self):
+        class cls(SocketServer.ForkingMixIn, SocketServer.UnixStreamServer):
+            ui = self.ui
+            repo = self.repo
+        self.server = cls(self.address, _requesthandler)
+        self.ui.status(_('listening at %s\n') % self.address)
+        self.ui.flush()  # avoid buffering of status message
+
+    def run(self):
+        try:
+            self.server.serve_forever()
+        finally:
+            os.unlink(self.address)
+
 _servicemap = {
     'pipe': pipeservice,
+    'unix': unixservice,
     }
 
 def createservice(ui, repo, opts):
diff --git a/tests/hghave.py b/tests/hghave.py
--- a/tests/hghave.py
+++ b/tests/hghave.py
@@ -1,5 +1,6 @@
 import os, stat
 import re
+import socket
 import sys
 import tempfile
 
@@ -258,6 +259,10 @@ def has_unix_permissions():
     finally:
         os.rmdir(d)
 
+ at check("unix-socket", "AF_UNIX socket family")
+def has_unix_socket():
+    return getattr(socket, 'AF_UNIX', None) is not None
+
 @check("root", "root permissions")
 def has_root():
     return getattr(os, 'geteuid', None) and os.geteuid() == 0
diff --git a/tests/test-commandserver.t b/tests/test-commandserver.t
--- a/tests/test-commandserver.t
+++ b/tests/test-commandserver.t
@@ -541,3 +541,63 @@ start without repository:
   *** runcommand init repo2
   *** runcommand id -R repo2
   000000000000 tip
+
+
+unix domain socket:
+
+  $ cd repo
+  $ hg update -q
+
+#if unix-socket
+
+  >>> import cStringIO
+  >>> from hgclient import unixserver, readchannel, runcommand, check
+  >>> server = unixserver('.hg/server.sock', '.hg/server.log')
+  >>> def hellomessage(conn):
+  ...     ch, data = readchannel(conn)
+  ...     print '%c, %r' % (ch, data)
+  ...     runcommand(conn, ['id'])
+  >>> check(hellomessage, server.connect)
+  o, 'capabilities: getencoding runcommand\nencoding: *' (glob)
+  *** runcommand id
+  eff892de26ec tip bm1/bm2/bm3
+  >>> def unknowncommand(conn):
+  ...     readchannel(conn)
+  ...     conn.stdin.write('unknowncommand\n')
+  >>> check(unknowncommand, server.connect)  # error sent to server.log
+  >>> def serverinput(conn):
+  ...     readchannel(conn)
+  ...     patch = """
+  ... # HG changeset patch
+  ... # User test
+  ... # Date 0 0
+  ... 2
+  ... 
+  ... diff -r eff892de26ec -r 1ed24be7e7a0 a
+  ... --- a/a
+  ... +++ b/a
+  ... @@ -1,1 +1,2 @@
+  ...  1
+  ... +2
+  ... """
+  ...     runcommand(conn, ['import', '-'], input=cStringIO.StringIO(patch))
+  ...     runcommand(conn, ['log', '-rtip', '-q'])
+  >>> check(serverinput, server.connect)
+  *** runcommand import -
+  applying patch from stdin
+  *** runcommand log -rtip -q
+  2:1ed24be7e7a0
+  >>> server.shutdown()
+
+  $ cat .hg/server.log
+  listening at .hg/server.sock
+  abort: unknown command unknowncommand
+  killed!
+
+#else
+
+  $ hg serve --cmdserver unix -a .hg/server.sock
+  abort: unsupported platform
+  [255]
+
+#endif