Add AutoContext to automatically insert tasks in batches

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
+

furious/context/context.py

@@ -51,6 +51,8 @@
 
 from .. import errors
 
+DEFAULT_TASK_BATCH_SIZE = 100
+
 
 class Context(object):
     """Furious context object.
@@ -102,21 +104,27 @@ 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(
                 "This Context has already had its tasks inserted.")
 
         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))
 

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)