streamclone: move payload header generation into own function
authorGregory Szorc <gregory.szorc@gmail.com>
Sun, 04 Oct 2015 19:06:06 -0700
changeset 26469 fb743268510e
parent 26468 19bbd53af46d
child 26470 4b5647d9ee13
streamclone: move payload header generation into own function The stream clone data over the wire protocol contains a header line indicating total file count and data size. In bundle2, this metadata can be captured by a part parameter and doesn't need to be in the body. In preparation for bundle2, have generatev1() return the raw metadata and move the header generation to its own function.
mercurial/streamclone.py
mercurial/wireproto.py
--- a/mercurial/streamclone.py	Sun Oct 04 18:44:46 2015 -0700
+++ b/mercurial/streamclone.py	Sun Oct 04 19:06:06 2015 -0700
@@ -164,13 +164,12 @@
 def generatev1(repo):
     """Emit content for version 1 of a streaming clone.
 
-    This is a generator of raw chunks that constitute a streaming clone.
+    This returns a 3-tuple of (file count, byte size, data iterator).
 
-    The stream begins with a line of 2 space-delimited integers containing the
-    number of entries and total bytes size.
+    The data iterator consists of N entries for each file being transferred.
+    Each file entry starts as a line with the file name and integer size
+    delimited by a null byte.
 
-    Next, are N entries for each file being transferred. Each file entry starts
-    as a line with the file name and integer size delimited by a null byte.
     The raw file data follows. Following the raw file data is the next file
     entry, or EOF.
 
@@ -196,31 +195,44 @@
 
     repo.ui.debug('%d files, %d bytes to transfer\n' %
                   (len(entries), total_bytes))
-    yield '%d %d\n' % (len(entries), total_bytes)
 
     svfs = repo.svfs
     oldaudit = svfs.mustaudit
     debugflag = repo.ui.debugflag
     svfs.mustaudit = False
 
-    try:
-        for name, size in entries:
-            if debugflag:
-                repo.ui.debug('sending %s (%d bytes)\n' % (name, size))
-            # partially encode name over the wire for backwards compat
-            yield '%s\0%d\n' % (store.encodedir(name), size)
-            if size <= 65536:
-                fp = svfs(name)
-                try:
-                    data = fp.read(size)
-                finally:
-                    fp.close()
-                yield data
-            else:
-                for chunk in util.filechunkiter(svfs(name), limit=size):
-                    yield chunk
-    finally:
-        svfs.mustaudit = oldaudit
+    def emitrevlogdata():
+        try:
+            for name, size in entries:
+                if debugflag:
+                    repo.ui.debug('sending %s (%d bytes)\n' % (name, size))
+                # partially encode name over the wire for backwards compat
+                yield '%s\0%d\n' % (store.encodedir(name), size)
+                if size <= 65536:
+                    fp = svfs(name)
+                    try:
+                        data = fp.read(size)
+                    finally:
+                        fp.close()
+                    yield data
+                else:
+                    for chunk in util.filechunkiter(svfs(name), limit=size):
+                        yield chunk
+        finally:
+            svfs.mustaudit = oldaudit
+
+    return len(entries), total_bytes, emitrevlogdata()
+
+def generatev1wireproto(repo):
+    """Emit content for version 1 of streaming clone suitable for the wire.
+
+    This is the data output from ``generatev1()`` with a header line
+    indicating file count and byte size.
+    """
+    filecount, bytecount, it = generatev1(repo)
+    yield '%d %d\n' % (filecount, bytecount)
+    for chunk in it:
+        yield chunk
 
 def consumev1(repo, fp, filecount, bytecount):
     """Apply the contents from version 1 of a streaming clone file handle.
--- a/mercurial/wireproto.py	Sun Oct 04 18:44:46 2015 -0700
+++ b/mercurial/wireproto.py	Sun Oct 04 19:06:06 2015 -0700
@@ -718,7 +718,7 @@
     try:
         # LockError may be raised before the first result is yielded. Don't
         # emit output until we're sure we got the lock successfully.
-        it = streamclone.generatev1(repo)
+        it = streamclone.generatev1wireproto(repo)
         return streamres(getstream(it))
     except error.LockError:
         return '2\n'