Add a wrappertask decorator

It is common that we might want to create a task that is a light wrapper
around another task. This allows us to, for example, call some clean-up
function after a PollingTaskGroup has completed. We really want to make
sure that any exceptions caused by timeout or cancellation get passed on to
the subtask as well.

Python 3 sports the "yield from" keyword (PEP 380) that effectively allows
the sort of nesting of generators that we want to do here. This change
provides a decorator function that allows the decorated generator function
to do the same thing using the "yield" keyword, even in Python 2.

Change-Id: I110141f2891ed35b8ffe98ec6ae8b42738f7db64

diff --git a/heat/engine/scheduler.py b/heat/engine/scheduler.py
index a321d3e82f6d5e1f577ebfae1af4051f64a8e0d1..e0dbf4ebe77518c64f6719e37267a3b01c7bc06e 100644 (file)
--- a/heat/engine/scheduler.py
+++ b/heat/engine/scheduler.py
@@ -16,6 +16,7 @@
 import eventlet
 import functools
 import itertools
+import sys
 import types
 
 from heat.openstack.common import excutils
@@ -144,6 +145,55 @@ class TaskRunner(object):
         return not self.done()
 
 
+def wrappertask(task):
+    """
+    Decorator for a task that needs to drive a subtask.
+
+    This is essentially a replacement for the Python 3-only "yield from"
+    keyword (PEP 380), using the "yield" keyword that is supported in
+    Python 2. For example:
+
+        @wrappertask
+        def parent_task(self):
+            self.setup()
+
+            yield self.child_task()
+
+            self.cleanup()
+    """
+
+    @functools.wraps(task)
+    def wrapper(*args, **kwargs):
+        parent = task(*args, **kwargs)
+
+        for subtask in parent:
+            try:
+                if subtask is not None:
+                    for step in subtask:
+                        try:
+                            yield step
+                        except GeneratorExit as exit:
+                            subtask.close()
+                            raise exit
+                        except:
+                            try:
+                                subtask.throw(*sys.exc_info())
+                            except StopIteration:
+                                break
+                else:
+                    yield
+            except GeneratorExit as exit:
+                parent.close()
+                raise exit
+            except:
+                try:
+                    parent.throw(*sys.exc_info())
+                except StopIteration:
+                    break
+
+    return wrapper
+
+
 class PollingTaskGroup(object):
     """
     A task which manages a group of subtasks.

diff --git a/heat/tests/test_scheduler.py b/heat/tests/test_scheduler.py
index 8fd2630febe58b7de25d323315a26cfd40ac50f6..684e4eec495034793d6aad4669325da19ac3ec21 100644 (file)
--- a/heat/tests/test_scheduler.py
+++ b/heat/tests/test_scheduler.py
@@ -351,3 +351,250 @@ class TaskTest(mox.MoxTestBase):
         runner.start()
         self.assertTrue(runner.step())
         self.assertTrue(runner.step())
+
+
+class WrapperTaskTest(mox.MoxTestBase):
+
+    def test_wrap(self):
+        child_tasks = [DummyTask() for i in range(3)]
+
+        @scheduler.wrappertask
+        def task():
+            for child_task in child_tasks:
+                yield child_task()
+
+            yield
+
+        for child_task in child_tasks:
+            self.mox.StubOutWithMock(child_task, 'do_step')
+        self.mox.StubOutWithMock(scheduler.TaskRunner, '_sleep')
+
+        for child_task in child_tasks:
+            child_task.do_step(1).AndReturn(None)
+            scheduler.TaskRunner._sleep(mox.IsA(int)).AndReturn(None)
+            child_task.do_step(2).AndReturn(None)
+            scheduler.TaskRunner._sleep(mox.IsA(int)).AndReturn(None)
+            child_task.do_step(3).AndReturn(None)
+            scheduler.TaskRunner._sleep(mox.IsA(int)).AndReturn(None)
+
+        self.mox.ReplayAll()
+
+        scheduler.TaskRunner(task)()
+
+    def test_child_exception(self):
+        class MyException(Exception):
+            pass
+
+        def child_task():
+            yield
+
+            raise MyException()
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield child_task()
+            except MyException:
+                raise
+            else:
+                self.fail('No exception raised in parent_task')
+
+        task = parent_task()
+        task.next()
+        self.assertRaises(MyException, task.next)
+
+    def test_child_exception_exit(self):
+        class MyException(Exception):
+            pass
+
+        def child_task():
+            yield
+
+            raise MyException()
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield child_task()
+            except MyException:
+                return
+            else:
+                self.fail('No exception raised in parent_task')
+
+        task = parent_task()
+        task.next()
+        self.assertRaises(StopIteration, task.next)
+
+    def test_child_exception_swallow(self):
+        class MyException(Exception):
+            pass
+
+        def child_task():
+            yield
+
+            raise MyException()
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield child_task()
+            except MyException:
+                yield
+            else:
+                self.fail('No exception raised in parent_task')
+
+            yield
+
+        task = parent_task()
+        task.next()
+        task.next()
+
+    def test_parent_exception(self):
+        class MyException(Exception):
+            pass
+
+        def child_task():
+            yield
+
+        @scheduler.wrappertask
+        def parent_task():
+            yield child_task()
+            raise MyException()
+
+        task = parent_task()
+        task.next()
+        self.assertRaises(MyException, task.next)
+
+    def test_parent_throw(self):
+        class MyException(Exception):
+            pass
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield DummyTask()()
+            except MyException:
+                raise
+            else:
+                self.fail('No exception raised in parent_task')
+
+        task = parent_task()
+        task.next()
+        self.assertRaises(MyException, task.throw, MyException())
+
+    def test_parent_throw_exit(self):
+        class MyException(Exception):
+            pass
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield DummyTask()()
+            except MyException:
+                return
+            else:
+                self.fail('No exception raised in parent_task')
+
+        task = parent_task()
+        task.next()
+        self.assertRaises(StopIteration, task.throw, MyException())
+
+    def test_parent_cancel(self):
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield
+            except GeneratorExit:
+                raise
+            else:
+                self.fail('parent_task not closed')
+
+        task = parent_task()
+        task.next()
+        task.close()
+
+    def test_parent_cancel_exit(self):
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield
+            except GeneratorExit:
+                return
+            else:
+                self.fail('parent_task not closed')
+
+        task = parent_task()
+        task.next()
+        task.close()
+
+    def test_cancel(self):
+        def child_task():
+            yield
+
+            try:
+                yield
+            except GeneratorExit:
+                raise
+            else:
+                self.fail('child_task not closed')
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield DummyTask()()
+            except GeneratorExit:
+                raise
+            else:
+                self.fail('parent_task not closed')
+
+        task = parent_task()
+        task.next()
+        task.close()
+
+    def test_cancel_exit(self):
+        def child_task():
+            yield
+
+            try:
+                yield
+            except GeneratorExit:
+                return
+            else:
+                self.fail('child_task not closed')
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield DummyTask()()
+            except GeneratorExit:
+                raise
+            else:
+                self.fail('parent_task not closed')
+
+        task = parent_task()
+        task.next()
+        task.close()
+
+    def test_cancel_parent_exit(self):
+        def child_task():
+            yield
+
+            try:
+                yield
+            except GeneratorExit:
+                return
+            else:
+                self.fail('child_task not closed')
+
+        @scheduler.wrappertask
+        def parent_task():
+            try:
+                yield DummyTask()()
+            except GeneratorExit:
+                return
+            else:
+                self.fail('parent_task not closed')
+
+        task = parent_task()
+        task.next()
+        task.close()