From: Zane Bitter Date: Tue, 7 May 2013 13:10:10 +0000 (+0200) Subject: Add a wrappertask decorator X-Git-Tag: 2014.1~640^2 X-Git-Url: https://review.fuel-infra.org/gitweb?a=commitdiff_plain;h=772f8b9e812ced3fe21425870bcde3b765cd73ab;p=openstack-build%2Fheat-build.git 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 a321d3e8..e0dbf4eb 100644 --- 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 8fd2630f..684e4eec 100644 --- 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()