bug fix for closing consumer - added tests

Author: lsbardel

.travis.yml

@@ -3,7 +3,7 @@ language: python
 sudo: false
 
 python:
- - 3.5.1
+ - 3.5.2
 
 services:
  - redis-server
@@ -14,12 +14,7 @@ install:
 script:
  - python setup.py test --coverage -q
  - flake8
- - if [[ $TRAVIS_PYTHON_VERSION == '3.5.1' ]]; then python setup.py test --coveralls; fi
+ - python setup.py test --coveralls
 
 notifications:
  email: false
-
-branches:
- only:
- - master
- - dev

README.rst

@@ -331,6 +331,12 @@ value.
  Data store used for publishing and subscribing to messages (redis is the
  only backend available at the moment)
 
+* **max_requests** (``--max-requests 0``)
+
+ The maximum number of tasks a worker will process before restarting.
+ A 0 value (the default) means no maximum number, workers will process
+ all tasks forever.
+
 * **message_broker** (``--message-broker ...``)
 
  Data store used as distributed task queue. If not provided (default) the

pq/server/apps.py

@@ -38,6 +38,8 @@ async def monitor_start(self, monitor, exc=None):
  if not exc:
  if self.cfg.schedule_periodic:
  self.backend_factory = TaskScheduler
+ elif self.cfg.workers:
+ self.backend_factory = TaskProducer
  self._backend = await self._start(monitor)
 
  def monitor_task(self, monitor):

pq/server/consumer.py

@@ -132,7 +132,6 @@ def __new__(cls, *args, **kwargs):
  o = super().__new__(cls)
  o._processed = 0
  o._next_time = 1
- o._closing_msg = None
  o._closing_waiter = None
  o._concurrent_tasks = set()
  return o
@@ -161,7 +160,7 @@ async def start(self, worker):
  '''Starts consuming tasks
  '''
  await self.pubsub.start()
- self._pool_tasks(worker)
+ self._poll_tasks(worker)
  self.logger.warning('%s started polling tasks', self)
  return self
 
@@ -171,24 +170,22 @@ def closing(self):
  def close(self, msg=None):
  if not self.closing():
  self._closing_waiter = Future(loop=self._loop)
- self._closing_msg = msg
+ self.logger.warning(msg)
  return self._closing_waiter
 
  # #######################################################################
  # # PRIVATE METHODS
  # #######################################################################
- def _pool_tasks(self, worker, next_time=None):
- if self.closing():
- if not self._concurrent_tasks:
- self._do_close()
- else:
- if worker.is_running() and not next_time:
- ensure_future(self._may_pool_task(worker), loop=worker._loop)
- elif not worker.after_run():
- next_time = next_time or 0
- worker._loop.call_later(next_time, self._pool_tasks, worker)
-
- async def _may_pool_task(self, worker):
+ def _poll_tasks(self, worker, next_time=None):
+ if self.closing() and not self._concurrent_tasks:
+ self._do_close()
+ elif worker.is_running() and not next_time:
+ ensure_future(self._may_poll_task(worker), loop=worker._loop)
+ elif not worker.after_run():
+ next_time = next_time or 0
+ worker._loop.call_later(next_time, self._poll_tasks, worker)
+
+ async def _may_poll_task(self, worker):
  # Called in the ``worker`` event loop.
  #
  # It pools a new task if possible, and add it to the queue of
@@ -200,7 +197,8 @@ async def _may_pool_task(self, worker):
  max_tasks = self.cfg.max_requests
  if max_tasks and self._processed >= max_tasks:
  self.close(
- 'Processed %s tasks. Restarting.' % self._processed
+ 'Processed %s tasks. Stop polling tasks.'
+ % self._processed
  )
 
  if not self.closing():
@@ -238,7 +236,7 @@ async def _may_pool_task(self, worker):
  next_time = self._next_time
  await self._broadcast(worker)
 
- self._pool_tasks(worker, next_time)
+ self._poll_tasks(worker, next_time)
 
  def _broadcast(self, worker):
  info = self.info()
@@ -247,7 +245,7 @@ def _broadcast(self, worker):
  return self.pubsub.publish(consumer_event, info)
 
  def _do_close(self):
- if self._closing_msg:
- self.logger.warning(self._closing_msg)
+ self.logger.warning('Closing %s', self)
  self.manager.close()
  self._closing_waiter.set_result(True)
+ self._loop.call_later(1, self._loop.stop)

pq/server/producer.py

@@ -29,8 +29,7 @@ class TaskProducer(models.RegistryMixin, ExecutorMixin):
  def __init__(self, cfg, *, logger=None, **kw):
  self.cfg = cfg
  self.logger = logger or logging.getLogger('pulsar.queue')
- self._closing_waiter = None
- self._closing = None
+ self._closing = False
  loop = cfg.params.pop('loop', None)
  store = create_store(cfg.data_store, loop=loop)
  if not cfg.message_broker:
@@ -88,12 +87,17 @@ def flush_queues(self, *queues):
  def on_events(self, callback):
  self.pubsub.on_events(callback)
 
+ def closing(self):
+ return self._closing
+
  def close(self):
  '''Close this :class:`.TaskBackend`.
 
  Invoked by the :class:`.Actor` when stopping.
  '''
- self.manager.close()
+ if not self._closing:
+ self._closing = True
+ self.manager.close()
 
  def queue_task(self, jobname, callback=True, **kwargs):
  '''Try to queue a new :task

tests/app.py

@@ -29,6 +29,8 @@ class TaskQueueBase:
  # used for both keep-alive and timeout in JsonProxy
  # long enough to allow to wait for tasks
  rpc_timeout = 500
+ max_requests = 0
+ concurrent_tasks = 5
  tq_app = None
  rpc = None
  schedule_periodic = False
@@ -46,15 +48,17 @@ def rpc_name(cls):
  async def setUpClass(cls):
  # The name of the task queue application
  params = cls.params()
- params.update(dict(
+ params.update(
  wsgi=True,
  schedule_periodic=cls.schedule_periodic,
  rpc_bind='127.0.0.1:0',
  concurrency=cls.concurrency,
+ concurrent_tasks=cls.concurrent_tasks,
+ max_requests=cls.max_requests,
  message_serializer=cls.message_serializer,
  rpc_concurrency=cls.concurrency,
  rpc_keep_alive=cls.rpc_timeout
- ))
+ )
  pq = api.PulsarQueue(**params)
  cfgs = await pq.start()
  cls.tq_app = cfgs[0].app()
@@ -285,6 +289,10 @@ async def test_thread_io(self):
  self.assertNotEqual(task.result['thread'], threading.get_ident())
  self.assertEqual(task.result['text'], 306)
 
+ async def test_bad_task(self):
+ task = await self.tq.queue_task('asynchronous', sleep=2)
+ self.assertEqual(task.status_string, 'FAILURE')
+
  def _test_sync(self):
  loop = asyncio.new_event_loop()
  tasks = api.TaskApp(loop=loop, **self.params()).api()

tests/example/manage.py

@@ -1,11 +1,6 @@
-task_paths = ['sampletasks.*']
-
-
-def app():
- from pq.api import TaskApp
- return TaskApp(config=__file__)
+task_paths = ['jobs.*', 'pq.jobs']
 
 
 if __name__ == '__main__': # pragma nocover
  from pq.api import PulsarQueue
- PulsarQueue('taskqueue', config='manage.py').start()
+ PulsarQueue(name='taskqueue', config=__file__).start()

tests/test_api.py

@@ -38,3 +38,26 @@ def test_format_time(self):
  st2 = format_time(timestamp)
  self.assertEqual(st, st2)
  self.assertEqual(format_time(None), '?')
+
+ def test_close(self):
+ t = api.TaskApp().api()
+ self.assertEqual(t.closing(), False)
+ t.close()
+ self.assertEqual(t.closing(), True)
+ warn = mock.MagicMock()
+ t.logger.warning = warn
+ self.assertFalse(t.queue_task('foo'))
+ self.assertEqual(warn.call_count, 1)
+ self.assertEqual(
+ warn.call_args[0][0],
+ 'Cannot queue task, task backend closing'
+ )
+
+ def test_task_not_available(self):
+ t = api.TaskApp().api()
+ self.assertRaises(api.TaskNotAvailable,
+ t.queue_task, 'jsdbcjsdhbc')
+
+ def test_no_queues(self):
+ t = api.TaskApp().api()
+ self.assertEqual(t.queues(), ())

tests/test_close.py

@@ -0,0 +1,36 @@
+"""Tests closing the worker after 10 requests"""
+import unittest
+import asyncio
+from random import random
+
+from tests import app
+
+
+class TestMsgPackQueue(app.TaskQueueBase, unittest.TestCase):
+ max_requests = 10
+ concurrent_tasks = 20
+
+ async def test_max_requests(self):
+ tasks = [self.tq.queue_task('asynchronous', lag=random())
+ for _ in range(18)]
+ tasks = await asyncio.gather(*tasks)
+ self.assertEqual(len(tasks), 18)
+ workers = set()
+ for task in tasks:
+ self.assertEqual(task.status_string, 'SUCCESS')
+ workers.add(task.worker)
+
+ self.assertEqual(len(workers), 2)
+
+ # FAILURES
+ tasks = [self.tq.queue_task('asynchronous', sleep=1)
+ for _ in range(6)]
+
+ tasks = await asyncio.gather(*tasks)
+ self.assertEqual(len(tasks), 6)
+ workers = set()
+ for task in tasks:
+ self.assertEqual(task.status_string, 'FAILURE')
+ workers.add(task.worker)
+
+ self.assertEqual(len(workers), 2)