Queue groups with random distribution

example/__init__.py

@@ -35,6 +35,7 @@
 from .context_intro import ContextIntroHandler
 from .grep import GrepHandler
 from .simple_workflow import SimpleWorkflowHandler
+from .queue_group_intro import QueueGroupIntroHandler
 
 config = {
     'webapp2_extras.jinja2': {
@@ -55,5 +56,6 @@
     ('/batcher/run', BatcherHandler),
     ('/batcher/stats', BatcherStatsHandler),
     ('/grep', GrepHandler),
+    ('/queue_group', QueueGroupIntroHandler)
 ], config=config)
 

example/queue_group_intro.py

@@ -0,0 +1,48 @@
+#
+# Copyright 2013 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.
+#
+
+"""A very basic async example using a queue group.
+
+This will insert 500 async tasks to run the example_function in a queue
+group.
+"""
+
+import logging
+
+import webapp2
+
+
+class QueueGroupIntroHandler(webapp2.RequestHandler):
+    """Demonstrate the creation and insertion of 500 furious tasks to be run
+    in a queue group with random distribution."""
+    def get(self):
+        from furious import context
+
+        with context.new() as ctx:
+            for i in xrange(500):
+                ctx.add(target=example_function, args=[i],
+                        queue_group='workers')
+
+        logging.info('500 Async jobs inserted.')
+
+        self.response.out.write('Successfully inserted 500 Async jobs.')
+
+
+def example_function(*args, **kwargs):
+    """This function is called by furious tasks to demonstrate usage."""
+    logging.info('example_function executed with args: %r, kwargs: %r',
+                 args, kwargs)
+    return args

furious/async.py

@@ -77,6 +77,8 @@ def run_me(*args, **kwargs):
 from .job_utils import path_to_reference
 from .job_utils import reference_to_path
 
+from .queues import select_queue
+
 from . import errors
 
 
@@ -259,6 +261,10 @@ def get_headers(self):
 
     def get_queue(self):
         """Return the queue the task should run in."""
+        queue_group = self._options.get('queue_group')
+        if queue_group:
+            return select_queue(queue_group)
+
         return self._options.get('queue', ASYNC_DEFAULT_QUEUE)
 
     def get_task_args(self):

furious/batcher.py

@@ -21,6 +21,7 @@
 from google.appengine.api import memcache
 
 from .async import Async
+from .queues import select_queue
 
 MESSAGE_DEFAULT_QUEUE = 'default-pull'
 MESSAGE_PROCESSOR_NAME = 'processor'
@@ -45,6 +46,10 @@ def update_options(self, **options):
 
     def get_queue(self):
         """Return the queue the task should run in."""
+        queue_group = self._options.get('queue_group')
+        if queue_group:
+            return select_queue(queue_group)
+
         return self._options.get('queue', MESSAGE_DEFAULT_QUEUE)
 
     def get_task_args(self):

furious/config.py

@@ -131,7 +131,8 @@ def default_config():
     return {'secret_key':
             '931b8-i-f44330b4a5-am-3b9b733f-not-secure-043e96882',
             'persistence': 'ndb',
-            'task_system': 'appengine_taskqueue'}
+            'task_system': 'appengine_taskqueue',
+            'queue_counts': {}}
 
 
 def _load_yaml_config(path=None):

furious/queues.py

@@ -0,0 +1,73 @@
+#
+# 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.
+#
+
+from random import shuffle
+from threading import local
+
+from furious.config import get_config
+
+QUEUE_COUNTS = 'queue_counts'
+
+
+def select_queue(queue_group):
+    """Select an optimal queue to run a task in from the given queue group. By
+    default, this simply randomly selects a queue from the group.
+
+    TODO: leverage the taskqueue API to try and determine the best queue to
+    use via a kwarg.
+    """
+
+    if not queue_group:
+        return None
+
+    queue_counts = get_config().get(QUEUE_COUNTS)
+    queue_count = queue_counts.get(queue_group, 0)
+
+    if queue_count <= 0:
+        raise Exception('Queue group must have at least 1 queue.')
+
+    queue_lists = _get_from_cache('queues', default={})
+    group_queues = queue_lists.setdefault(queue_group, [])
+
+    if not group_queues:
+        group_queues.extend('%s-%d' % (queue_group, i)
+                            for i in xrange(queue_count))
+
+        shuffle(group_queues)
+
+    return group_queues.pop()
+
+
+# Use a thread local cache to optimize performance and queue distribution.
+_cache = local()
+
+
+def _get_from_cache(key, default=None):
+    """Fetch the value for the given key from the thread local cache. If the
+    key does not exist, set it for the given default value and return the
+    default.
+    """
+
+    if not key:
+        return default
+
+    if hasattr(_cache, key):
+        return getattr(_cache, key)
+
+    setattr(_cache, key, default)
+
+    return default
+

furious/tests/test_async.py

@@ -327,6 +327,25 @@ def test_get_default_queue(self):
 
         self.assertEqual('default', job.get_queue())
 
+    @patch('furious.async.select_queue')
+    def test_get_queue_group_queue(self, mock_select_queue):
+        """Ensure get_queue calls through to select_queue when a queue group
+        is given.
+        """
+        from furious.async import Async
+        from furious.async import ASYNC_DEFAULT_QUEUE
+
+        queue_group = 'foo-queue'
+        queue_count = 5
+
+        job = Async('nonexistant', queue_group=queue_group,
+                    queue_count=queue_count, default=ASYNC_DEFAULT_QUEUE)
+
+        expected = '%s-2' % queue_group
+        mock_select_queue.return_value = expected
+
+        self.assertEqual(expected, job.get_queue())
+
     def test_get_task_args(self):
         """Ensure get_task_args returns the job task_args."""
         from furious.async import Async

furious/tests/test_batcher.py

@@ -71,6 +71,25 @@ def test_get_default_queue(self):
 
         self.assertEqual('default-pull', message.get_queue())
 
+    @patch('furious.batcher.select_queue')
+    def test_get_queue_group_queue(self, mock_select_queue):
+        """Ensure get_queue calls through to select_queue when a queue group
+        is given.
+        """
+        from furious.batcher import Message
+        from furious.batcher import MESSAGE_DEFAULT_QUEUE
+
+        queue_group = 'foo-queue'
+        queue_count = 5
+
+        message = Message(queue_group=queue_group, queue_count=queue_count,
+                          default=MESSAGE_DEFAULT_QUEUE)
+
+        expected = '%s-2' % queue_group
+        mock_select_queue.return_value = expected
+
+        self.assertEqual(expected, message.get_queue())
+
     def test_get_task_args(self):
         """Ensure get_task_args returns the message task_args."""
         from furious.batcher import Message

furious/tests/test_config.py

@@ -62,13 +62,16 @@ def test_get_config(self):
 
         example_yaml = str('secret_key: "blah"\n'
                            'persistence: bubble\n'
-                           'task_system: flah\n')
+                           'task_system: flah\n'
+                           'queue_counts:\n'
+                           '    queue-group: 5')
 
         my_config = _parse_yaml_config(example_yaml)
 
         self.assertEqual(my_config, {'secret_key': 'blah',
                                      'persistence': 'bubble',
-                                     'task_system': 'flah'})
+                                     'task_system': 'flah',
+                                     'queue_counts': {'queue-group': 5}})
 
     def test_get_configured_persistence_exists(self):
         """Ensure a chosen persistence module is selected."""

furious/tests/test_queues.py

@@ -0,0 +1,140 @@
+#
+# 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 mock import patch
+
+
+class TestSelectQueue(unittest.TestCase):
+    """Ensure select_queue() works correctly."""
+
+    def setUp(self):
+        from threading import local
+        from furious import async
+
+        async._cache = local()
+
+    def test_none(self):
+        """Ensure that None is returned when the queue group is None."""
+        from furious.async import select_queue
+
+        actual = select_queue(None)
+
+        self.assertEqual(actual, None)
+
+    @patch('furious.queues.get_config')
+    def test_invalid_queue_count(self, mock_get_config):
+        """Ensure that an exception is raised when a bad queue count is given.
+        """
+        from furious.async import select_queue
+
+        queue_group = 'foo-queue'
+
+        mock_get_config.return_value = {'queue_counts':
+                                        {queue_group: 0}}
+
+        with self.assertRaises(Exception) as context:
+            select_queue(queue_group)
+
+        self.assertEqual(context.exception.message,
+                         'Queue group must have at least 1 queue.')
+
+    @patch('furious.queues.get_config')
+    @patch('furious.queues._get_from_cache')
+    def test_random_from_cache(self, mock_cache, mock_get_config):
+        """Ensure that a random queue is selected from the group when there are
+        cached queues.
+        """
+        from furious.async import select_queue
+
+        queue_group = 'foo-queue'
+        queue_count = 5
+        expected = '%s-4' % queue_group
+
+        mock_get_config.return_value = {'queue_counts':
+                                        {queue_group: queue_count}}
+
+        mock_cache.return_value = {queue_group: [expected]}
+
+        actual = select_queue(queue_group)
+
+        self.assertEqual(actual, expected)
+        mock_cache.assert_called_once_with('queues', default={})
+
+    @patch('furious.queues.get_config')
+    @patch('furious.queues.shuffle')
+    @patch('furious.queues._get_from_cache')
+    def test_random_no_cache(self, mock_cache, mock_shuffle, mock_get_config):
+        """Ensure that a random queue is selected from the group when there are
+        no cached queues.
+        """
+        from furious.async import select_queue
+
+        queue_group = 'foo-queue'
+        queue_count = 5
+        expected = '%s-4' % queue_group
+
+        mock_get_config.return_value = {'queue_counts':
+                                        {queue_group: queue_count}}
+
+        mock_cache.return_value = {}
+
+        actual = select_queue(queue_group)
+
+        self.assertEqual(actual, expected)
+        mock_cache.assert_called_once_with('queues', default={})
+        self.assertTrue(mock_shuffle.called)
+
+    def test_get_from_cache_no_key(self):
+        """Ensure the default value is returned when the key is None."""
+        from furious.queues import _get_from_cache
+
+        key = None
+        expected = 'default'
+
+        actual = _get_from_cache(key, default=expected)
+
+        self.assertEqual(actual, expected)
+
+    def test_get_from_cache_present(self):
+        """Ensure the correct value is returned when the key is present."""
+        from furious.queues import _cache
+        from furious.queues import _get_from_cache
+
+        key = 'key'
+        expected = 'value'
+        setattr(_cache, key, expected)
+
+        actual = _get_from_cache(key)
+
+        self.assertEqual(actual, expected)
+
+    def test_get_from_cache_not_present(self):
+        """Ensure the default value is returned and set when the key is not
+        present.
+        """
+        from furious.queues import _cache
+        from furious.queues import _get_from_cache
+
+        key = 'key'
+        expected = 'default'
+
+        actual = _get_from_cache(key, default=expected)
+
+        self.assertEqual(actual, expected)
+        self.assertTrue(hasattr(_cache, key))
+