unbundle: extract checkheads in its own function

We are going to refactor the unbundle function to have it working on
a local repository too. Having this function extracted will ease the process.

In the case of non-matching heads, the function now directly raises an
exception. The top level of the function is catching it.

mercurial/exchange.py

@@ -611,3 +611,20 b' def getbundle(repo, source, heads=None, '
         temp.write(c)
     temp.seek(0)
     return bundle2.unbundle20(repo.ui, temp)
+
+class PushRaced(RuntimeError):
+    """An exception raised during unbunding that indicate a push race"""
+
+def check_heads(repo, their_heads, context):
+    """check if the heads of a repo have been modified
+
+    Used by peer for unbundling.
+    """
+    heads = repo.heads()
+    heads_hash = util.sha1(''.join(sorted(heads))).digest()
+    if not (their_heads == ['force'] or their_heads == heads or
+            their_heads == ['hashed', heads_hash]):
+        # someone else committed/pushed/unbundled while we
+        # were transferring data
+        raise PushRaced('repository changed while %s - '
+                        'please try again' % context)

mercurial/wireproto.py

@@ -9,7 +9,7 b' import urllib, tempfile, os, sys'
 from i18n import _
 from node import bin, hex
 import changegroup as changegroupmod
-import peer, error, encoding, util, store
+import peer, error, encoding, util, store, exchange
 
 
 class abstractserverproto(object):
@@ -754,46 +754,36 b' def stream(repo, proto):'
 def unbundle(repo, proto, heads):
     their_heads = decodelist(heads)
 
-    def check_heads():
-        heads = repo.heads()
-        heads_hash = util.sha1(''.join(sorted(heads))).digest()
-        return (their_heads == ['force'] or their_heads == heads or
-                their_heads == ['hashed', heads_hash])
+    try:
+        proto.redirect()
 
-    proto.redirect()
+        exchange.check_heads(repo, their_heads, 'preparing changes')
 
-    # fail early if possible
-    if not check_heads():
-        return pusherr('repository changed while preparing changes - '
-                       'please try again')
-
-    # write bundle data to temporary file because it can be big
-    fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
-    fp = os.fdopen(fd, 'wb+')
-    r = 0
-    try:
-        proto.getfile(fp)
-        lock = repo.lock()
+        # write bundle data to temporary file because it can be big
+        fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
+        fp = os.fdopen(fd, 'wb+')
+        r = 0
         try:
-            if not check_heads():
-                # someone else committed/pushed/unbundled while we
-                # were transferring data
-                return pusherr('repository changed while uploading changes - '
-                               'please try again')
-
-            # push can proceed
-            fp.seek(0)
-            gen = changegroupmod.readbundle(fp, None)
-
+            proto.getfile(fp)
+            lock = repo.lock()
             try:
-                r = changegroupmod.addchangegroup(repo, gen, 'serve',
-                                                  proto._client())
-            except util.Abort, inst:
-                sys.stderr.write("abort: %s\n" % inst)
+                exchange.check_heads(repo, their_heads, 'uploading changes')
+
+                # push can proceed
+                fp.seek(0)
+                gen = changegroupmod.readbundle(fp, None)
+
+                try:
+                    r = changegroupmod.addchangegroup(repo, gen, 'serve',
+                                                      proto._client())
+                except util.Abort, inst:
+                    sys.stderr.write("abort: %s\n" % inst)
+            finally:
+                lock.release()
+            return pushres(r)
+
         finally:
-            lock.release()
-        return pushres(r)
-
-    finally:
-        fp.close()
-        os.unlink(tempname)
+            fp.close()
+            os.unlink(tempname)
+    except exchange.PushRaced, exc:
+        return pusherr(str(exc))