bundle2: move processpart stream maintenance into part iterator
authorDurham Goode <durham@fb.com>
Thu, 14 Sep 2017 10:20:05 -0700
changeset 34258 e71890f27767
parent 34257 2844c4bd5a39
child 34259 07e4170f02f3
bundle2: move processpart stream maintenance into part iterator The processpart function also did some stream maintenance, so let's move it to the part iterator as well, as part of moving all part iteration logic into the class. There is one place processpart is called outside of the normal loop, so we manually handle the seek there. The now-empty try/finally will be removed in a later patch, for ease of review. Differential Revision: https://phab.mercurial-scm.org/D706
mercurial/bundle2.py
--- a/mercurial/bundle2.py	Mon Sep 18 14:12:20 2017 -0500
+++ b/mercurial/bundle2.py	Thu Sep 14 10:20:05 2017 -0700
@@ -354,13 +354,17 @@
         self.unbundler = unbundler
         self.iterator = None
         self.count = 0
+        self.current = None
 
     def __enter__(self):
         def func():
             itr = enumerate(self.unbundler.iterparts())
             for count, p in itr:
                 self.count = count
+                self.current = p
                 yield p
+                p.seek(0, 2)
+                self.current = None
         self.iterator = func()
         return self.iterator
 
@@ -369,6 +373,13 @@
             return
 
         if exc:
+            # If exiting or interrupted, do not attempt to seek the stream in
+            # the finally block below. This makes abort faster.
+            if (self.current and
+                not isinstance(exc, (SystemExit, KeyboardInterrupt))):
+                # consume the part content to not corrupt the stream.
+                self.current.seek(0, 2)
+
             # Any exceptions seeking to the end of the bundle at this point are
             # almost certainly related to the underlying stream being bad.
             # And, chances are that the exception we're handling is related to
@@ -455,7 +466,6 @@
     The part is guaranteed to have been fully consumed when the function exits
     (even if an exception is raised)."""
     status = 'unknown' # used by debug output
-    hardabort = False
     try:
         try:
             handler = parthandlermapping.get(part.type)
@@ -511,15 +521,8 @@
                                            mandatory=False)
                 outpart.addparam(
                     'in-reply-to', pycompat.bytestr(part.id), mandatory=False)
-    # If exiting or interrupted, do not attempt to seek the stream in the
-    # finally block below. This makes abort faster.
-    except (SystemExit, KeyboardInterrupt):
-        hardabort = True
-        raise
     finally:
-        # consume the part content to not corrupt the stream.
-        if not hardabort:
-            part.seek(0, 2)
+        pass
 
 
 def decodecaps(blob):
@@ -1147,7 +1150,15 @@
             return
         part = unbundlepart(self.ui, headerblock, self._fp)
         op = interruptoperation(self.ui)
-        _processpart(op, part)
+        hardabort = False
+        try:
+            _processpart(op, part)
+        except (SystemExit, KeyboardInterrupt):
+            hardabort = True
+            raise
+        finally:
+            if not hardabort:
+                part.seek(0, 2)
         self.ui.debug('bundle2-input-stream-interrupt:'
                       ' closing out of band context\n')