diff mercurial/exchange.py @ 21061:62d35f251c60

bundle2: allow using bundle2 for push

We now support bundle2 for local push. The unbundle function has to detect
which version of the bundle to use since the return type is different.

Note that push error handling is currently nonexistent. This is one of the
reasons why bundle2 is still disabled by default.

--- a/mercurial/exchange.py	Fri Apr 11 06:43:01 2014 -0700
+++ b/mercurial/exchange.py	Thu Apr 10 10:53:43 2014 -0700
@@ -106,7 +106,9 @@
                 pushop.repo.prepushoutgoinghooks(pushop.repo,
                                                  pushop.remote,
                                                  pushop.outgoing)
-                _pushchangeset(pushop)
+                if pushop.remote.capable('bundle2'):
+                    _pushbundle2(pushop)
+                else:
             _pushcomputecommonheads(pushop)
             _pushsyncphase(pushop)
             _pushobsolete(pushop)
@@ -172,6 +174,35 @@
                              newbm)
     return True
 
+def _pushbundle2(pushop):
+    """push data to the remote using bundle2
+
+    The only currently supported type of data is changegroup but this will
+    evolve in the future."""
+    # Send known head to the server for race detection.
+    bundler = bundle2.bundle20(pushop.ui)
+    if not pushop.force:
+        part = bundle2.bundlepart('CHECK:HEADS', data=iter(pushop.remoteheads))
+        bundler.addpart(part)
+    # add the changegroup bundle
+    cg = changegroup.getlocalbundle(pushop.repo, 'push', pushop.outgoing)
+    def cgchunks(cg=cg):
+        yield 'HG10UN'
+        for c in cg.getchunks():
+            yield c
+    cgpart = bundle2.bundlepart('CHANGEGROUP', data=cgchunks())
+    bundler.addpart(cgpart)
+    stream = util.chunkbuffer(bundler.getchunks())
+    sent = bundle2.unbundle20(pushop.repo.ui, stream)
+    reply = pushop.remote.unbundle(sent, ['force'], 'push')
+    try:
+        op = bundle2.processbundle(pushop.repo, reply)
+    except KeyError, exc:
+        raise util.Abort('missing support for %s' % exc)
+    cgreplies = op.records.getreplies(cgpart.id)
+    assert len(cgreplies['changegroup']) == 1
+    pushop.ret = cgreplies['changegroup'][0]['return']
+
 def _pushchangeset(pushop):
     """Make the actual push of changeset bundle to remote repo"""
     outgoing = pushop.outgoing
@@ -637,11 +668,22 @@
 
     If the push was raced as PushRaced exception is raised."""
     r = 0
+    # need a transaction when processing a bundle2 stream
+    tr = None
     lock = repo.lock()
     try:
         check_heads(repo, heads, 'uploading changes')
         # push can proceed
-        r = changegroup.addchangegroup(repo, cg, source, url)
+        if util.safehasattr(cg, 'params'):
+            tr = repo.transaction('unbundle')
+            ret = bundle2.processbundle(repo, cg, lambda: tr)
+            tr.close()
+            stream = util.chunkbuffer(ret.reply.getchunks())
+            r = bundle2.unbundle20(repo.ui, stream)
+        else:
+            r = changegroup.addchangegroup(repo, cg, source, url)
     finally:
+        if tr is not None:
+            tr.release()
         lock.release()
     return r