Merge pull request #3604 from gmbnomis/asyncio_stages_cancel

DeclarativeVersion: Cancel all remaining stages if one stage fails with an exception

plugin/pulpcore/plugin/stages/api.py

@@ -52,9 +52,21 @@ async def create_pipeline(stages, maxsize=100):
     futures = []
     for stage in stages:
         out_q = asyncio.Queue(maxsize=maxsize)
-        futures.append(stage(in_q, out_q))
+        futures.append(asyncio.ensure_future(stage(in_q, out_q)))
         in_q = out_q
-    await asyncio.gather(*futures)
+    try:
+        await asyncio.gather(*futures)
+    except Exception:
+        # One of the stages raised an exception, cancel all stages...
+        pending = []
+        for task in futures:
+            if not task.done():
+                task.cancel()
+                pending.append(task)
+        # ...and run until all Exceptions show up
+        if pending:
+            await asyncio.wait(pending, timeout=60)
+        raise
 
 
 class EndStage(Stage):

pulpcore/pulpcore/app/models/progress.py

@@ -2,6 +2,8 @@
 Django models related to progress reporting
 """
 from gettext import gettext as _
+
+from asyncio import CancelledError
 import logging
 import datetime
 
@@ -91,11 +93,14 @@ def __enter__(self):
 
     def __exit__(self, type, value, traceback):
         """
-        Update the progress report state to COMPLETED or FAILED.
+        Update the progress report state to COMPLETED, CANCELED, or FAILED.
+
+        If an exception occurs the progress report state is saved as:
+        - CANCELED if the exception is `asyncio.CancelledError`
+        - FAILED otherwise.
 
-        If an exception occurs the progress report state is saved as FAILED and the exception is
-        not suppressed. If the context manager exited without exception the progress report state
-        is saved as COMPLETED.
+        The exception is not suppressed. If the context manager exited without
+        exception the progress report state is saved as COMPLETED.
 
         See the context manager documentation for more info on __exit__ parameters
         """
@@ -104,10 +109,11 @@ def __exit__(self, type, value, traceback):
             self.total = self.done
         if type is None:
             self.state = TASK_STATES.COMPLETED
-            self.save()
+        elif type is CancelledError:
+            self.state = TASK_STATES.CANCELED
         else:
             self.state = TASK_STATES.FAILED
-            self.save()
+        self.save()
 
 
 class ProgressSpinner(ProgressReport):
@@ -128,8 +134,8 @@ class ProgressSpinner(ProgressReport):
         >>> metadata_progress.save()
 
     The ProgressSpinner() is a context manager that provides automatic state transitions and saving
-    for the RUNNING COMPLETED and FAILED states. When ProgressSpinner() is used as a context
-    manager progress reporting is rate limited to every 500 milliseconds.
+    for the RUNNING CANCELED COMPLETED and FAILED states. When ProgressSpinner() is used as a
+    context manager progress reporting is rate limited to every 500 milliseconds.
     Use it as follows:
 
         >>> spinner = ProgressSpinner(message='Publishing Metadata')
@@ -171,9 +177,9 @@ class ProgressBar(ProgressReport):
         >>> progress_bar.save()
 
     The ProgressBar() is a context manager that provides automatic state transitions and saving for
-    the RUNNING COMPLETED and FAILED states. The increment() method can be called in the loop as
-    work is completed. When ProgressBar() is used as a context manager progress reporting is rate
-    limited to every 500 milliseconds.
+    the RUNNING CANCELED COMPLETED and FAILED states. The increment() method can be called in the
+    loop as work is completed. When ProgressBar() is used as a context manager progress reporting
+    is rate limited to every 500 milliseconds.
     Use it as follows:
 
         >>> progress_bar = ProgressBar(message='Publishing files', total=len(files_iterator))