Mercurial > public > mercurial-scm > hg-stable

diff mercurial/wireproto.py @ 37414:2d965bfeb8f6

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
wireproto: allow direct stream processing for unbundle Introduce a new option server.streamunbundle which starts a transaction immediately to apply a bundle instead of writing it to a temporary file first. This side steps the need for a large tmp directory at the cost of preventing concurrent pushes. This is a reasonable trade-off for many setups as concurrent pushes for the main branch at least are disallowed anyway. The option defaults to off to preserve existing behavior. Change the wireproto interface to provide a generator for reading the payload and make callers responsible for consuming all data. Differential Revision: https://phab.mercurial-scm.org/D2470
author Joerg Sonnenberger <joerg@bec.de>
date Tue, 27 Feb 2018 02:37:31 +0100
parents 33af46d639b4
children 0b7475ea38cf
line wrap: on
line diff
--- a/mercurial/wireproto.py	Fri Apr 06 22:22:19 2018 +0200
+++ b/mercurial/wireproto.py	Tue Feb 27 02:37:31 2018 +0100
@@ -1082,14 +1082,33 @@
     with proto.mayberedirectstdio() as output:
         try:
             exchange.check_heads(repo, their_heads, 'preparing changes')
+            cleanup = lambda: None
+            try:
+                payload = proto.getpayload()
+                if repo.ui.configbool('server', 'streamunbundle'):
+                    def cleanup():
+                        # Ensure that the full payload is consumed, so
+                        # that the connection doesn't contain trailing garbage.
+                        for p in payload:
+                            pass
+                    fp = util.chunkbuffer(payload)
+                else:
+                    # write bundle data to temporary file as it can be big
+                    fp, tempname = None, None
+                    def cleanup():
+                        if fp:
+                            fp.close()
+                        if tempname:
+                            os.unlink(tempname)
+                    fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
+                    repo.ui.debug('redirecting incoming bundle to %s\n' %
+                        tempname)
+                    fp = os.fdopen(fd, pycompat.sysstr('wb+'))
+                    r = 0
+                    for p in payload:
+                        fp.write(p)
+                    fp.seek(0)
 
-            # write bundle data to temporary file because it can be big
-            fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
-            fp = os.fdopen(fd, r'wb+')
-            r = 0
-            try:
-                proto.forwardpayload(fp)
-                fp.seek(0)
                 gen = exchange.readbundle(repo.ui, fp, None)
                 if (isinstance(gen, changegroupmod.cg1unpacker)
                     and not bundle1allowed(repo, 'push')):
@@ -1112,8 +1131,7 @@
                     r, output.getvalue() if output else '')
 
             finally:
-                fp.close()
-                os.unlink(tempname)
+                cleanup()
 
         except (error.BundleValueError, error.Abort, error.PushRaced) as exc:
             # handle non-bundle2 case first