cmdserver: make server streams switchable

In 'unix' mode, server instance will be created per connection, and fin/fout
are set to socket files.

mercurial/commandserver.py

@@ -1,268 +1,268 b''
 # commandserver.py - communicate with Mercurial's API over a pipe
 #
 #  Copyright Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 from i18n import _
 import struct
 import sys, os
 import dispatch, encoding, util
 
 logfile = None
 
 def log(*args):
     if not logfile:
         return
 
     for a in args:
         logfile.write(str(a))
 
     logfile.flush()
 
 class channeledoutput(object):
     """
     Write data to out in the following format:
 
     data length (unsigned int),
     data
     """
     def __init__(self, out, channel):
         self.out = out
         self.channel = channel
 
     def write(self, data):
         if not data:
             return
         self.out.write(struct.pack('>cI', self.channel, len(data)))
         self.out.write(data)
         self.out.flush()
 
     def __getattr__(self, attr):
         if attr in ('isatty', 'fileno'):
             raise AttributeError(attr)
         return getattr(self.out, attr)
 
 class channeledinput(object):
     """
     Read data from in_.
 
     Requests for input are written to out in the following format:
     channel identifier - 'I' for plain input, 'L' line based (1 byte)
     how many bytes to send at most (unsigned int),
 
     The client replies with:
     data length (unsigned int), 0 meaning EOF
     data
     """
 
     maxchunksize = 4 * 1024
 
     def __init__(self, in_, out, channel):
         self.in_ = in_
         self.out = out
         self.channel = channel
 
     def read(self, size=-1):
         if size < 0:
             # if we need to consume all the clients input, ask for 4k chunks
             # so the pipe doesn't fill up risking a deadlock
             size = self.maxchunksize
             s = self._read(size, self.channel)
             buf = s
             while s:
                 s = self._read(size, self.channel)
                 buf += s
 
             return buf
         else:
             return self._read(size, self.channel)
 
     def _read(self, size, channel):
         if not size:
             return ''
         assert size > 0
 
         # tell the client we need at most size bytes
         self.out.write(struct.pack('>cI', channel, size))
         self.out.flush()
 
         length = self.in_.read(4)
         length = struct.unpack('>I', length)[0]
         if not length:
             return ''
         else:
             return self.in_.read(length)
 
     def readline(self, size=-1):
         if size < 0:
             size = self.maxchunksize
             s = self._read(size, 'L')
             buf = s
             # keep asking for more until there's either no more or
             # we got a full line
             while s and s[-1] != '\n':
                 s = self._read(size, 'L')
                 buf += s
 
             return buf
         else:
             return self._read(size, 'L')
 
     def __iter__(self):
         return self
 
     def next(self):
         l = self.readline()
         if not l:
             raise StopIteration
         return l
 
     def __getattr__(self, attr):
         if attr in ('isatty', 'fileno'):
             raise AttributeError(attr)
         return getattr(self.in_, attr)
 
 class server(object):
     """
-    Listens for commands on stdin, runs them and writes the output on a channel
-    based stream to stdout.
+    Listens for commands on fin, runs them and writes the output on a channel
+    based stream to fout.
     """
-    def __init__(self, ui, repo):
+    def __init__(self, ui, repo, fin, fout):
         self.cwd = os.getcwd()
 
         logpath = ui.config("cmdserver", "log", None)
         if logpath:
             global logfile
             if logpath == '-':
                 # write log on a special 'd' (debug) channel
-                logfile = channeledoutput(sys.stdout, 'd')
+                logfile = channeledoutput(fout, 'd')
             else:
                 logfile = open(logpath, 'a')
 
         if repo:
             # the ui here is really the repo ui so take its baseui so we don't
             # end up with its local configuration
             self.ui = repo.baseui
             self.repo = repo
             self.repoui = repo.ui
         else:
             self.ui = ui
             self.repo = self.repoui = None
 
-        self.cerr = channeledoutput(sys.stdout, 'e')
-        self.cout = channeledoutput(sys.stdout, 'o')
-        self.cin = channeledinput(sys.stdin, sys.stdout, 'I')
-        self.cresult = channeledoutput(sys.stdout, 'r')
+        self.cerr = channeledoutput(fout, 'e')
+        self.cout = channeledoutput(fout, 'o')
+        self.cin = channeledinput(fin, fout, 'I')
+        self.cresult = channeledoutput(fout, 'r')
 
-        self.client = sys.stdin
+        self.client = fin
 
     def _read(self, size):
         if not size:
             return ''
 
         data = self.client.read(size)
 
         # is the other end closed?
         if not data:
             raise EOFError
 
         return data
 
     def runcommand(self):
         """ reads a list of \0 terminated arguments, executes
         and writes the return code to the result channel """
 
         length = struct.unpack('>I', self._read(4))[0]
         if not length:
             args = []
         else:
             args = self._read(length).split('\0')
 
         # copy the uis so changes (e.g. --config or --verbose) don't
         # persist between requests
         copiedui = self.ui.copy()
         uis = [copiedui]
         if self.repo:
             self.repo.baseui = copiedui
             # clone ui without using ui.copy because this is protected
             repoui = self.repoui.__class__(self.repoui)
             repoui.copy = copiedui.copy # redo copy protection
             uis.append(repoui)
             self.repo.ui = self.repo.dirstate._ui = repoui
             self.repo.invalidateall()
 
         for ui in uis:
             # any kind of interaction must use server channels
             ui.setconfig('ui', 'nontty', 'true', 'commandserver')
 
         req = dispatch.request(args[:], copiedui, self.repo, self.cin,
                                self.cout, self.cerr)
 
         ret = (dispatch.dispatch(req) or 0) & 255 # might return None
 
         # restore old cwd
         if '--cwd' in args:
             os.chdir(self.cwd)
 
         self.cresult.write(struct.pack('>i', int(ret)))
 
     def getencoding(self):
         """ writes the current encoding to the result channel """
         self.cresult.write(encoding.encoding)
 
     def serveone(self):
         cmd = self.client.readline()[:-1]
         if cmd:
             handler = self.capabilities.get(cmd)
             if handler:
                 handler(self)
             else:
                 # clients are expected to check what commands are supported by
                 # looking at the servers capabilities
                 raise util.Abort(_('unknown command %s') % cmd)
 
         return cmd != ''
 
     capabilities = {'runcommand'  : runcommand,
                     'getencoding' : getencoding}
 
     def serve(self):
         hellomsg = 'capabilities: ' + ' '.join(sorted(self.capabilities))
         hellomsg += '\n'
         hellomsg += 'encoding: ' + encoding.encoding
 
         # write the hello msg in -one- chunk
         self.cout.write(hellomsg)
 
         try:
             while self.serveone():
                 pass
         except EOFError:
             # we'll get here if the client disconnected while we were reading
             # its request
             return 1
 
         return 0
 
 class pipeservice(object):
     def __init__(self, ui, repo, opts):
-        self.server = server(ui, repo)
+        self.server = server(ui, repo, sys.stdin, sys.stdout)
 
     def init(self):
         pass
 
     def run(self):
         return self.server.serve()
 
 _servicemap = {
     'pipe': pipeservice,
     }
 
 def createservice(ui, repo, opts):
     mode = opts['cmdserver']
     try:
         return _servicemap[mode](ui, repo, opts)
     except KeyError:
         raise util.Abort(_('unknown mode %s') % mode)