# -*- coding: utf-8 -*- # Copyright (C) 2014 Yahoo! Inc. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); you may # not use this file except in compliance with the License. You may obtain # a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the # License for the specific language governing permissions and limitations # under the License. import collections import contextlib import threading import futurist import testscenarios from zake import fake_client from taskflow.conductors import backends from taskflow import engines from taskflow.jobs.backends import impl_zookeeper from taskflow.jobs import base from taskflow.patterns import linear_flow as lf from taskflow.persistence.backends import impl_memory from taskflow import states as st from taskflow import test from taskflow.tests import utils as test_utils from taskflow.utils import persistence_utils as pu from taskflow.utils import threading_utils @contextlib.contextmanager def close_many(*closeables): try: yield finally: for c in closeables: c.close() def test_factory(blowup): f = lf.Flow("test") if not blowup: f.add(test_utils.ProgressingTask('test1')) else: f.add(test_utils.FailingTask("test1")) return f def sleep_factory(): f = lf.Flow("test") f.add(test_utils.SleepTask('test1')) f.add(test_utils.ProgressingTask('test2')) return f def test_store_factory(): f = lf.Flow("test") f.add(test_utils.TaskMultiArg('task1')) return f def single_factory(): return futurist.ThreadPoolExecutor(max_workers=1) ComponentBundle = collections.namedtuple('ComponentBundle', ['board', 'client', 'persistence', 'conductor']) class ManyConductorTest(testscenarios.TestWithScenarios, test_utils.EngineTestBase, test.TestCase): scenarios = [ ('blocking', {'kind': 'blocking', 'conductor_kwargs': {'wait_timeout': 0.1}}), ('nonblocking_many_thread', {'kind': 'nonblocking', 'conductor_kwargs': {'wait_timeout': 0.1}}), ('nonblocking_one_thread', {'kind': 'nonblocking', 'conductor_kwargs': { 'executor_factory': single_factory, 'wait_timeout': 0.1, }}) ] def make_components(self): client = fake_client.FakeClient() persistence = impl_memory.MemoryBackend() board = impl_zookeeper.ZookeeperJobBoard('testing', {}, client=client, persistence=persistence) conductor_kwargs = self.conductor_kwargs.copy() conductor_kwargs['persistence'] = persistence conductor = backends.fetch(self.kind, 'testing', board, **conductor_kwargs) return ComponentBundle(board, client, persistence, conductor) def test_connection(self): components = self.make_components() components.conductor.connect() with close_many(components.conductor, components.client): self.assertTrue(components.board.connected) self.assertTrue(components.client.connected) self.assertFalse(components.board.connected) self.assertFalse(components.client.connected) def test_run_empty(self): components = self.make_components() components.conductor.connect() with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() components.conductor.stop() self.assertTrue( components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) t.join() def test_run(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() job_consumed_event = threading.Event() job_abandoned_event = threading.Event() def on_consume(state, details): consumed_event.set() def on_job_consumed(event, details): if event == 'job_consumed': job_consumed_event.set() def on_job_abandoned(event, details): if event == 'job_abandoned': job_abandoned_event.set() components.board.notifier.register(base.REMOVAL, on_consume) components.conductor.notifier.register("job_consumed", on_job_consumed) components.conductor.notifier.register("job_abandoned", on_job_abandoned) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence) engines.save_factory_details(fd, test_factory, [False], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid}) self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT)) self.assertTrue(job_consumed_event.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(job_abandoned_event.wait(1)) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) persistence = components.persistence with contextlib.closing(persistence.get_connection()) as conn: lb = conn.get_logbook(lb.uuid) fd = lb.find(fd.uuid) self.assertIsNotNone(fd) self.assertEqual(st.SUCCESS, fd.state) def test_run_max_dispatches(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() def on_consume(state, details): consumed_event.set() components.board.notifier.register(base.REMOVAL, on_consume) with close_many(components.client, components.conductor): t = threading_utils.daemon_thread( lambda: components.conductor.run(max_dispatches=5)) t.start() lb, fd = pu.temporary_flow_detail(components.persistence) engines.save_factory_details(fd, test_factory, [False], {}, backend=components.persistence) for _ in range(5): components.board.post('poke', lb, details={'flow_uuid': fd.uuid}) self.assertTrue(consumed_event.wait( test_utils.WAIT_TIMEOUT)) components.board.post('poke', lb, details={'flow_uuid': fd.uuid}) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) def test_fail_run(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() job_consumed_event = threading.Event() job_abandoned_event = threading.Event() def on_consume(state, details): consumed_event.set() def on_job_consumed(event, details): if event == 'job_consumed': job_consumed_event.set() def on_job_abandoned(event, details): if event == 'job_abandoned': job_abandoned_event.set() components.board.notifier.register(base.REMOVAL, on_consume) components.conductor.notifier.register("job_consumed", on_job_consumed) components.conductor.notifier.register("job_abandoned", on_job_abandoned) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence) engines.save_factory_details(fd, test_factory, [True], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid}) self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT)) self.assertTrue(job_consumed_event.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(job_abandoned_event.wait(1)) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) persistence = components.persistence with contextlib.closing(persistence.get_connection()) as conn: lb = conn.get_logbook(lb.uuid) fd = lb.find(fd.uuid) self.assertIsNotNone(fd) self.assertEqual(st.REVERTED, fd.state) def test_missing_store(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() def on_consume(state, details): consumed_event.set() components.board.notifier.register(base.REMOVAL, on_consume) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence) engines.save_factory_details(fd, test_store_factory, [], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid}) self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT)) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) persistence = components.persistence with contextlib.closing(persistence.get_connection()) as conn: lb = conn.get_logbook(lb.uuid) fd = lb.find(fd.uuid) self.assertIsNotNone(fd) self.assertIsNone(fd.state) def test_job_store(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() def on_consume(state, details): consumed_event.set() store = {'x': True, 'y': False, 'z': None} components.board.notifier.register(base.REMOVAL, on_consume) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence) engines.save_factory_details(fd, test_store_factory, [], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid, 'store': store}) self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT)) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) persistence = components.persistence with contextlib.closing(persistence.get_connection()) as conn: lb = conn.get_logbook(lb.uuid) fd = lb.find(fd.uuid) self.assertIsNotNone(fd) self.assertEqual(st.SUCCESS, fd.state) def test_flowdetails_store(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() def on_consume(state, details): consumed_event.set() store = {'x': True, 'y': False, 'z': None} components.board.notifier.register(base.REMOVAL, on_consume) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence, meta={'store': store}) engines.save_factory_details(fd, test_store_factory, [], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid}) self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT)) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) persistence = components.persistence with contextlib.closing(persistence.get_connection()) as conn: lb = conn.get_logbook(lb.uuid) fd = lb.find(fd.uuid) self.assertIsNotNone(fd) self.assertEqual(st.SUCCESS, fd.state) def test_combined_store(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() def on_consume(state, details): consumed_event.set() flow_store = {'x': True, 'y': False} job_store = {'z': None} components.board.notifier.register(base.REMOVAL, on_consume) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence, meta={'store': flow_store}) engines.save_factory_details(fd, test_store_factory, [], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid, 'store': job_store}) self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT)) components.conductor.stop() self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT)) self.assertFalse(components.conductor.dispatching) persistence = components.persistence with contextlib.closing(persistence.get_connection()) as conn: lb = conn.get_logbook(lb.uuid) fd = lb.find(fd.uuid) self.assertIsNotNone(fd) self.assertEqual(st.SUCCESS, fd.state) def test_stop_aborts_engine(self): components = self.make_components() components.conductor.connect() consumed_event = threading.Event() job_consumed_event = threading.Event() job_abandoned_event = threading.Event() running_start_event = threading.Event() def on_running_start(event, details): running_start_event.set() def on_consume(state, details): consumed_event.set() def on_job_consumed(event, details): if event == 'job_consumed': job_consumed_event.set() def on_job_abandoned(event, details): if event == 'job_abandoned': job_abandoned_event.set() components.board.notifier.register(base.REMOVAL, on_consume) components.conductor.notifier.register("job_consumed", on_job_consumed) components.conductor.notifier.register("job_abandoned", on_job_abandoned) components.conductor.notifier.register("running_start", on_running_start) with close_many(components.conductor, components.client): t = threading_utils.daemon_thread(components.conductor.run) t.start() lb, fd = pu.temporary_flow_detail(components.persistence) engines.save_factory_details(fd, sleep_factory, [], {}, backend=components.persistence) components.board.post('poke', lb, details={'flow_uuid': fd.uuid, 'store': {'duration': 2}}) running_start_event.wait(test_utils.WAIT_TIMEOUT) components.conductor.stop() job_abandoned_event.wait(test_utils.WAIT_TIMEOUT) self.assertTrue(job_abandoned_event.is_set()) self.assertFalse(job_consumed_event.is_set()) self.assertFalse(consumed_event.is_set()) class NonBlockingExecutorTest(test.TestCase): def test_bad_wait_timeout(self): persistence = impl_memory.MemoryBackend() client = fake_client.FakeClient() board = impl_zookeeper.ZookeeperJobBoard('testing', {}, client=client, persistence=persistence) self.assertRaises(ValueError, backends.fetch, 'nonblocking', 'testing', board, persistence=persistence, wait_timeout='testing') def test_bad_factory(self): persistence = impl_memory.MemoryBackend() client = fake_client.FakeClient() board = impl_zookeeper.ZookeeperJobBoard('testing', {}, client=client, persistence=persistence) self.assertRaises(ValueError, backends.fetch, 'nonblocking', 'testing', board, persistence=persistence, executor_factory='testing')