diff --git a/furious/context/auto_context.py b/furious/context/auto_context.py new file mode 100644 index 0000000..12dbd25 --- /dev/null +++ b/furious/context/auto_context.py @@ -0,0 +1,87 @@ +# +# Copyright 2012 WebFilings, LLC +# +# 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. +# + +""" +Furious AutoContext is used batch inserts of tasks into groups. + +It is similar to Context, but inserts automatically before the context +is exited. +""" + + +from furious.context.context import Context + + +class AutoContext(Context): + """Similar to context, but automatically inserts tasks asynchronously as + they are added to the context. Inserted in batches if specified. + """ + + def __init__(self, batch_size=None, **options): + """Setup this context in addition to accepting a batch_size.""" + + Context.__init__(self, **options) + + self.batch_size = batch_size + + def add(self, target, args=None, kwargs=None, **options): + """Add an Async job to this context. + + Like Context.add(): creates an Async and adds it to our list of tasks. + but also calls _auto_insert_check() to add tasks to queues + automatically. + """ + + # In superclass, add new task to our list of tasks + target = super( + AutoContext, self).add(target, args, kwargs, **options) + + self._auto_insert_check() + + return target + + def _auto_insert_check(self): + """Automatically insert tasks asynchronously. + Depending on batch_size, insert or wait until next call. + """ + + if not self.batch_size: + return + + if len(self._tasks) >= self.batch_size: + self._handle_tasks() + + def _handle_tasks(self): + """Convert Async's into tasks, then insert them into queues. + Similar to the default _handle_tasks, but don't mark all + tasks inserted. + """ + + self._handle_tasks_insert(batch_size=self.batch_size) + self._tasks = [] + + def __exit__(self, exc_type, exc_val, exc_tb): + """In addition to the default __exit__(), also mark all tasks + inserted. + """ + + super(AutoContext, self).__exit__(exc_type, exc_val, exc_tb) + + # Mark all tasks inserted. + self._tasks_inserted = True + + return False + diff --git a/furious/context/context.py b/furious/context/context.py index 2584f56..1328f62 100644 --- a/furious/context/context.py +++ b/furious/context/context.py @@ -51,6 +51,8 @@ from .. import errors +DEFAULT_TASK_BATCH_SIZE = 100 + class Context(object): """Furious context object. @@ -102,7 +104,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): return False - def _handle_tasks(self): + def _handle_tasks_insert(self, batch_size=None): """Convert all Async's into tasks, then insert them into queues.""" if self._tasks_inserted: raise errors.ContextAlreadyStartedError( @@ -110,13 +112,19 @@ def _handle_tasks(self): task_map = self._get_tasks_by_queue() for queue, tasks in task_map.iteritems(): - for batch in _task_batcher(tasks): + for batch in _task_batcher(tasks, batch_size=batch_size): inserted = self._insert_tasks(batch, queue=queue) if isinstance(inserted, (int, long)): # Don't blow up on insert_tasks that don't return counts. self._insert_success_count += inserted self._insert_failed_count += len(batch) - inserted + def _handle_tasks(self): + """Convert all Async's into tasks, then insert them into queues. + Also mark all tasks inserted to ensure they are not reinserted later. + """ + self._handle_tasks_insert() + self._tasks_inserted = True def _get_tasks_by_queue(self): @@ -133,8 +141,8 @@ def _get_tasks_by_queue(self): def add(self, target, args=None, kwargs=None, **options): """Add an Async job to this context. - Takes an Async object or the argumnets to construct an Async - object as argumnets. Returns the newly added Async object. + Takes an Async object or the arguments to construct an Async + object as arguments. Returns the newly added Async object. """ from ..async import Async from ..batcher import Message @@ -256,12 +264,15 @@ def _insert_tasks(tasks, queue, transactional=False): return inserted -def _task_batcher(tasks): +def _task_batcher(tasks, batch_size=None): """Batches large task lists into groups of 100 so that they can all be inserted. """ from itertools import izip_longest - args = [iter(tasks)] * 100 + if not batch_size: + batch_size = DEFAULT_TASK_BATCH_SIZE + + args = [iter(tasks)] * batch_size return ([task for task in group if task] for group in izip_longest(*args)) diff --git a/furious/tests/context/test_auto_context.py b/furious/tests/context/test_auto_context.py new file mode 100644 index 0000000..09d3716 --- /dev/null +++ b/furious/tests/context/test_auto_context.py @@ -0,0 +1,128 @@ +# +# Copyright 2012 WebFilings, LLC +# +# 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 unittest + +from google.appengine.ext import testbed + +from mock import Mock +from mock import patch + + +class TestAutoContextTestCase(unittest.TestCase): + """Test the AutoContext class.""" + + def setUp(self): + """Setup the test harness and a request hash.""" + import os + import uuid + + self.harness = testbed.Testbed() + self.harness.activate() + self.harness.init_taskqueue_stub() + + # Backup environment + self._orig_environ = os.environ.copy() + # Ensure each test looks like it is in a new request. + os.environ['REQUEST_ID_HASH'] = uuid.uuid4().hex + + def tearDown(self): + """Deactive the test harness, and restore the environment.""" + import os + + self.harness.deactivate() + os.environ.clear() + os.environ.update(self._orig_environ) + + @patch('google.appengine.api.taskqueue.Queue.add', auto_spec=True) + def test_add_job_to_context_multiple_batches(self, queue_add_mock): + """Ensure adding more tasks than the batch_size causes multiple batches + to get inserted. + + Adding 3 asyncs with a batch_size of 2 should result in two queue.adds, + containing 2 and 1 task respectively. + + Also ensure that the remaining task is inserted upon exiting the + context. + """ + from furious.async import Async + from furious.context.auto_context import AutoContext + + batch_size = 2 + + with AutoContext(batch_size) as ctx: + # First batch + job1 = ctx.add('test', args=[1, 2]) + job2 = ctx.add('test2', args=[1, 2]) + + # Second batch (not added until "with" section ends)) + job3 = ctx.add('test3', args=[1, 2]) + + # Ensure the first two jobs were inserted in the first batch. + self.assertIsInstance(job1, Async) + self.assertIsInstance(job2, Async) + queue_add_mock.assert_called_once() + #Ensure only two tasks were inserted + tasks_added = queue_add_mock.call_args[0][0] + self.assertEqual(2, len(tasks_added)) + + # Ensure the third job was inserted when the context exited. + self.assertIsInstance(job3, Async) + # Ensure add has now been called twice. + self.assertEqual(2, queue_add_mock.call_count) + # Ensure only one task was inserted + tasks_added = queue_add_mock.call_args[0][0] + self.assertEqual(1, len(tasks_added)) + + @patch('google.appengine.api.taskqueue.Queue.add', auto_spec=True) + def test_add_job_to_context_batch_size_unspecified(self, queue_add_mock): + """When batch_size is None or 0, the default behavior of Context is + used. All the tasks are added to the queue when the context is exited. + """ + from furious.async import Async + from furious.context.auto_context import AutoContext + + with AutoContext() as ctx: + + job1 = ctx.add('test', args=[1, 2]) + job2 = ctx.add('test2', args=[1, 2]) + job3 = ctx.add('test3', args=[1, 2]) + + self.assertIsInstance(job1, Async) + self.assertIsInstance(job2, Async) + self.assertIsInstance(job3, Async) + + # Ensure no batches of tasks are inserted yet. + self.assertFalse(queue_add_mock.called) + + # Ensure the list of tasks added when the context exited. + self.assertEqual(1, queue_add_mock.call_count) + # Ensure the three tasks were inserted + tasks_added = queue_add_mock.call_args[0][0] + self.assertEqual(3, len(tasks_added)) + + @patch('google.appengine.api.taskqueue.Queue.add', auto_spec=True) + def test_no_jobs(self, queue_add_mock): + """When no Asyncs are added to the context, ensure that there are no + errors and nothing is added to the task queue. + """ + from furious.context.auto_context import AutoContext + + with AutoContext(3): + pass + + # Ensure queue.add() was never called. + self.assertEqual(0, queue_add_mock.call_count)