# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@fb.com>
# Date 2014-04-10 17:53:43
# Node ID 62d35f251c602d3c470f02358fc559d1db5cef1f
# Parent  0bea9db7543bb15f08275b700d7e5e478cc73e54

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.

diff --git a/mercurial/exchange.py b/mercurial/exchange.py
--- a/mercurial/exchange.py
+++ b/mercurial/exchange.py
@@ -106,7 +106,9 @@ def push(repo, remote, force=False, revs
                 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 @@ def _pushcheckoutgoing(pushop):
                              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 @@ def unbundle(repo, cg, heads, source, ur
 
     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
diff --git a/tests/test-bundle2.t b/tests/test-bundle2.t
--- a/tests/test-bundle2.t
+++ b/tests/test-bundle2.t
@@ -668,11 +668,21 @@ clone --pull
 
 pull
 
-  $ hg -R other pull
+  $ hg -R other pull -r 24b6387c8c8c
   pulling from $TESTTMP/main (glob)
   searching for changes
   adding changesets
   adding manifests
   adding file changes
-  added 7 changesets with 6 changes to 6 files (+3 heads)
+  added 1 changesets with 1 changes to 1 files (+1 heads)
   (run 'hg heads' to see heads, 'hg merge' to merge)
+
+push
+
+  $ hg -R main push other --rev eea13746799a
+  pushing to other
+  searching for changes
+  adding changesets
+  adding manifests
+  adding file changes
+  added 1 changesets with 0 changes to 0 files (-1 heads)