D706: bundle2: move processpart stream maintenance into part iterator

[durham (Durham Goode)]

durham updated this revision to Diff 1818.

REPOSITORY
  rHG Mercurial

CHANGES SINCE LAST UPDATE
  https://phab.mercurial-scm.org/D706?vs=1802&id=1818

REVISION DETAIL
  https://phab.mercurial-scm.org/D706

AFFECTED FILES
  mercurial/bundle2.py

CHANGE DETAILS

diff --git a/mercurial/bundle2.py b/mercurial/bundle2.py
--- a/mercurial/bundle2.py
+++ b/mercurial/bundle2.py
@@ -354,21 +354,32 @@
         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
 
     def __exit__(self, type, exc, tb):
         if not self.iterator:
             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):
@@ -1143,7 +1146,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')
 



To: durham, #hg-reviewers, indygreg
Cc: indygreg, mercurial-devel