[PATCH 6 of 6] bundle2: read the whole bundle from stream on abort

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

# HG changeset patch
# User Pierre-Yves David <pierre-yves.david at fb.com>
# Date 1395706815 25200
#      Mon Mar 24 17:20:15 2014 -0700
# Node ID cb8e491405095e50481a5a3380801c6edf79229d
# Parent  8f0da5aa2129a78b4e7891e29ec583d9296ac10c
bundle2: read the whole bundle from stream on abort

When the bundle processing abort on unknown mandatory parts, we now makes sure
all the bundle content is read. This avoid leaving the communication channel in
an unrecoverable state.

diff --git a/mercurial/bundle2.py b/mercurial/bundle2.py
--- a/mercurial/bundle2.py
+++ b/mercurial/bundle2.py
@@ -128,12 +128,15 @@ for a certain part type.
 
 The matching of a part to its handler is case insensitive. The case of the
 part type is used to know if a part is mandatory or advisory. If the Part type
 contains any uppercase char it is considered mandatory. When no handler is
 known for a Mandatory part, the process is aborted and an exception is raised.
-If the part is advisory and no handler is known, the part is ignored.
-
+If the part is advisory and no handler is known, the part is ignored. When the
+process is aborted, the full bundle is still read from the stream to keep the
+channel usable. But none of the part read from an abort are processed. In the
+future, dropping the stream may become an option for channel we do not care to
+preserve.
 """
 
 import util
 import struct
 import urllib
@@ -203,27 +206,33 @@ def processbundle(repo, stream):
     unbundler = unbundle20(ui, stream)
     # todo:
     # - replace this is a init function soon.
     # - exception catching
     unbundler.params
-    for part in unbundler:
-        parttype = part.type
-        # part key are matched lower case
-        key = parttype.lower()
-        try:
-            handler = parthandlermapping[key]
-            ui.debug('found an handler for part %r\n' % parttype)
-        except KeyError:
-            if key != parttype: # mandatory parts
+    iterparts = iter(unbundler)
+    try:
+        for part in iterparts:
+            parttype = part.type
+            # part key are matched lower case
+            key = parttype.lower()
+            try:
+                handler = parthandlermapping[key]
+                ui.debug('found an handler for part %r\n' % parttype)
+            except KeyError:
+                if key != parttype: # mandatory parts
+                    # todo:
+                    # - use a more precise exception
+                    raise
+                ui.debug('ignoring unknown advisory part %r\n' % key)
                 # todo:
-                # - use a more precise exception
-                # - consume the bundle2 stream anyway.
-                raise
-            ui.debug('ignoring unknown advisory part %r\n' % key)
-            # todo: consume the part (once we use streamed parts)
-            continue
-        handler(repo, part)
+                # - consume the part once we use streaming
+                continue
+            handler(repo, part)
+    except Exception:
+        for part in iterparts:
+            pass # consume the bundle content
+        raise
 
 class bundle20(object):
     """represent an outgoing bundle2 container
 
     Use the `addparam` method to add stream level parameter. and `addpart` to
diff --git a/tests/test-bundle2.t b/tests/test-bundle2.t
--- a/tests/test-bundle2.t
+++ b/tests/test-bundle2.t
@@ -70,13 +70,17 @@ Create an extension to test bundle2 API
   > 
   > @command('unbundle2', [], '')
   > def cmdunbundle2(ui, repo):
   >     """process a bundle2 stream from stdin on the current repo"""
   >     try:
-  >         bundle2.processbundle(repo, sys.stdin)
-  >     except KeyError, exc:
-  >         raise util.Abort('missing support for %s' % exc)
+  >         try:
+  >             bundle2.processbundle(repo, sys.stdin)
+  >         except KeyError, exc:
+  >             raise util.Abort('missing support for %s' % exc)
+  >     finally:
+  >         remains = sys.stdin.read()
+  >         ui.write('%i unread bytes\n' % len(remains))
   > 
   > @command('statbundle2', [], '')
   > def cmdstatbundle2(ui, repo):
   >     """print statistic on the bundle2 container read from stdin"""
   >     unbundler = bundle2.unbundle20(ui, sys.stdin)
@@ -383,16 +387,18 @@ Process the bundle
   payload chunk size: 2
   payload chunk size: 0
   ignoring unknown advisory part 'test:math'
   part header size: 0
   end of bundle2 stream
+  0 unread bytes
 
 
   $ hg bundle2 --parts --unknown ../unknown.hg2
 
   $ hg unbundle2 < ../unknown.hg2
   The choir start singing:
       Patali Dirapata, Cromda Cromda Ripalo, Pata Pata, Ko Ko Ko
       Bokoro Dipoulito, Rondi Rondi Pepino, Pata Pata, Ko Ko Ko
       Emana Karassoli, Loucra Loucra Ponponto, Pata Pata, Ko Ko Ko.
+  0 unread bytes
   abort: missing support for 'test:unknown'
   [255]