[PATCH RFC] wireproto: add support for batching of some commands

[Peter Arrenbrecht]

# HG changeset patch
# User Peter Arrenbrecht <peter.arrenbrecht at gmail.com>
# Date 1300093067 -3600
wireproto: add support for batching of some commands

Allows clients to send batches of server commands in one roundtrip. An
example is provided in the contained test.

This is desirable for the new discovery protocol. For example, in a later
patch we will batch the request for the server's heads with the one where
the server checks if it knows the client's heads, eliminating one roundtrip
out of 4 for a pull where local is a subset of remote.

diff --git a/mercurial/wireproto.py b/mercurial/wireproto.py
--- a/mercurial/wireproto.py
+++ b/mercurial/wireproto.py
@@ -20,9 +20,80 @@
 def encodelist(l, sep=' '):
     return sep.join(map(hex, l))
 
+# batch encoding
+
+def escapearg(plain):
+    # FIXME might be faster using regexp
+    return (plain
+            .replace(':', '::')
+            .replace(',', ':,')
+            .replace(';', ':;')
+            .replace('=', ':='))
+
+def unescapearg(escaped):
+    # FIXME might be faster using regexp
+    return (escaped
+            .replace(':=', '=')
+            .replace(':;', ';')
+            .replace(':,', ',')
+            .replace('::', ':'))
+
 # client side
 
+# FIXME maybe merge into commands below
+batchcoders = {
+          'branches': (
+            {'nodes': encodelist},
+            lambda rsp: [tuple(decodelist(b)) for b in rsp.splitlines()],
+            ),
+          'heads': (
+            {},
+            lambda rsp: decodelist(rsp[:-1]),
+            ),
+          }
+
 class wirerepository(repo.repository):
+
+    def batch(self, cmds):
+        """sends a batch of commands to the server in one request
+
+        cmds is [(op, {name: val},)]
+
+        The command arguments are encoded normally first. Then commands and their
+        arguments are packed into a string of the form:
+
+          <cmd> <name>=<arg>,<name>=<arg>,...;...
+
+        where all ";" chars in <args> are escaped as ":;", all "," as ":,",
+        all "=" as ":=", and all ":" as "::".
+        """
+        reqs = []
+        for op, args in cmds:
+            enc, _dec = batchcoders[op]
+            args = ','.join(n + '=' + escapearg(enc[n](v))
+                            for n, v in args.iteritems())
+            reqs.append(op + ' ' + args)
+        rsp = self._call("batch", cmds=';'.join(reqs))
+        try:
+            res = []
+            for pair in rsp.split(';'):
+                op, result = pair.split(' ',1)
+                _enc, dec = batchcoders[op]
+                result = dec(unescapearg(result))
+                res.append((op, result,))
+            return res
+        except:
+            self._abort(error.ResponseError(_("unexpected response:"), rsp))
+
+    def _single(self, op, **args):
+        enc, dec = batchcoders[op]
+        req = dict((n, enc[n](v),) for n, v in args.iteritems())
+        rsp = self._call(op, **req)
+        try:
+            return dec(rsp)
+        except:
+            self._abort(error.ResponseError(_("unexpected response:"), rsp))
+
     def lookup(self, key):
         self.requirecap('lookup', _('look up remote revision'))
         d = self._call("lookup", key=encoding.fromlocal(key))
@@ -32,11 +103,7 @@
         self._abort(error.RepoError(data))
 
     def heads(self):
-        d = self._call("heads")
-        try:
-            return decodelist(d[:-1])
-        except:
-            self._abort(error.ResponseError(_("unexpected response:"), d))
+        return self._single('heads')
 
     def branchmap(self):
         d = self._call("branchmap")
@@ -52,13 +119,7 @@
             self._abort(error.ResponseError(_("unexpected response:"), d))
 
     def branches(self, nodes):
-        n = encodelist(nodes)
-        d = self._call("branches", nodes=n)
-        try:
-            br = [tuple(decodelist(b)) for b in d.splitlines()]
-            return br
-        except:
-            self._abort(error.ResponseError(_("unexpected response:"), d))
+        return self._single('branches', nodes=nodes)
 
     def between(self, pairs):
         batch = 8 # avoid giant requests
@@ -152,6 +213,24 @@
     args = proto.getargs(spec)
     return func(repo, proto, *args)
 
+def batch(repo, proto, cmds):
+    res = []
+    for pair in cmds.split(';'):
+        op, args = pair.split(' ', 1)
+        vals = {}
+        for a in args.split(','):
+            if a:
+                n, v = a.split('=')
+                vals[n] = unescapearg(v)
+        func, spec = commands[op]
+        if spec:
+            spec = spec.split(' ')
+            result = func(repo, proto, *[vals[n] for n in spec])
+        else:
+            result = func(repo, proto)
+        res.append(op + ' ' + escapearg(result))
+    return ';'.join(res)
+
 def between(repo, proto, pairs):
     pairs = [decodelist(p, '-') for p in pairs.split(" ")]
     r = []
@@ -176,7 +255,7 @@
     return "".join(r)
 
 def capabilities(repo, proto):
-    caps = 'lookup changegroupsubset branchmap pushkey'.split()
+    caps = 'lookup changegroupsubset batch branchmap pushkey'.split()
     if _allowstream(repo.ui):
         requiredformats = repo.requirements & repo.supportedformats
         # if our local revlogs are just revlogv1, add 'stream' cap
@@ -337,6 +416,7 @@
         os.unlink(tempname)
 
 commands = {
+    'batch': (batch, 'cmds'),
     'between': (between, 'pairs'),
     'branchmap': (branchmap, ''),
     'branches': (branches, 'nodes'),
diff --git a/tests/test-hgweb-commands.t b/tests/test-hgweb-commands.t
--- a/tests/test-hgweb-commands.t
+++ b/tests/test-hgweb-commands.t
@@ -905,7 +905,7 @@
   $ "$TESTDIR/get-with-headers.py" 127.0.0.1:$HGPORT '?cmd=capabilities'; echo
   200 Script output follows
   
-  lookup changegroupsubset branchmap pushkey unbundle=HG10GZ,HG10BZ,HG10UN
+  lookup changegroupsubset batch branchmap pushkey unbundle=HG10GZ,HG10BZ,HG10UN
 
 heads
 
diff --git a/tests/test-wireprotocol.py b/tests/test-wireprotocol.py
new file mode 100644
--- /dev/null
+++ b/tests/test-wireprotocol.py
@@ -0,0 +1,37 @@
+from mercurial import wireproto
+
+class proto():
+    def __init__(self, args):
+        self.args = args
+    def getargs(self, spec):
+        names = spec.split(' ')
+        return [self.args[n] for n in names]
+
+class clientrepo(wireproto.wirerepository):
+    def __init__(self, serverrepo):
+        self.serverrepo = serverrepo
+    def _call(self, cmd, **args):
+        return wireproto.dispatch(self.serverrepo, proto(args), cmd)
+    def greet(self, name):
+        return self._single('greet', name=name)
+
+class serverrepo():
+    def greet(self, name):
+        return "Hello, " + name
+
+def mangle(s):
+    return ''.join(chr(ord(c) + 1) for c in s)
+def unmangle(s):
+    return ''.join(chr(ord(c) - 1) for c in s)
+
+def greet(repo, proto, name):
+    return mangle(repo.greet(unmangle(name)))
+
+wireproto.batchcoders['greet'] = ({'name': mangle}, unmangle,)
+wireproto.commands['greet'] = (greet, 'name',)
+
+srv = serverrepo()
+clt = clientrepo(srv)
+
+print clt.greet("Foobar")
+print clt.batch([('greet', {'name':"Fo, =;o"},), ('greet', {'name':"Bar"},)])
diff --git a/tests/test-wireprotocol.py.out b/tests/test-wireprotocol.py.out
new file mode 100644
--- /dev/null
+++ b/tests/test-wireprotocol.py.out
@@ -0,0 +1,2 @@
+Hello, Foobar
+[('greet', 'Hello, Fo, =;o'), ('greet', 'Hello, Bar')]