[PATCH 2 of 4] unbundle: extract checkheads in its own function

[pierre-yves.david at ens-lyon.org]

# HG changeset patch
# User Pierre-Yves David <pierre-yves.david at fb.com>
# Date 1396656925 25200
#      Fri Apr 04 17:15:25 2014 -0700
# Node ID f3a8d0e8c734e5aaf76d57d5a080c13f37d64979
# Parent  a049f28c37ddc81665f40e4ce84fcffa205f3c45
unbundle: extract checkheads in its own function

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

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

diff --git a/mercurial/exchange.py b/mercurial/exchange.py
--- a/mercurial/exchange.py
+++ b/mercurial/exchange.py
@@ -609,5 +609,23 @@ def getbundle(repo, source, heads=None, 
     temp = cStringIO.StringIO()
     for c in bundler.getchunks():
         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)
diff --git a/mercurial/wireproto.py b/mercurial/wireproto.py
--- a/mercurial/wireproto.py
+++ b/mercurial/wireproto.py
@@ -7,11 +7,11 @@
 
 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):
     """abstract class that summarizes the protocol API
 
@@ -752,48 +752,38 @@ def stream(repo, proto):
 
 @wireprotocommand('unbundle', 'heads')
 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()
+            try:
+                exchange.check_heads(repo, their_heads, 'uploading changes')
 
-    # 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()
-        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)
 
-            # 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)
 
-            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:
-        fp.close()
-        os.unlink(tempname)
+            fp.close()
+            os.unlink(tempname)
+    except exchange.PushRaced, exc:
+        return pusherr(exc.message)