switch wait > Semaphore

Author: samuelcolvin

HISTORY.rst

@@ -3,10 +3,12 @@
 History
 -------
 
-v0.8.0 (TBC)
-............
+v0.8.0 (2017-06-05)
+...................
 * add ``async-timeout`` dependency and use async timeout around ``shadow_factory``
 * change logger name for control process log messages
+* use ``Semaphore`` rather than ``asyncio.wait(...return_when=asyncio.FIRST_COMPLETED)`` for improved performance
+* improve log display
 
 v0.7.0 (2017-06-01)
 ...................

arq/drain.py

@@ -47,6 +47,7 @@ def __init__(self, *,
         self.redis_pool = redis_pool
         self.loop = redis_pool._loop
         self.max_concurrent_tasks = max_concurrent_tasks
+        self.task_semaphore = asyncio.Semaphore(value=max_concurrent_tasks, loop=self.loop)
         self.shutdown_delay = max(shutdown_delay, 0.1)
         self.timeout_seconds = timeout_seconds
         self.burst_mode = burst_mode
@@ -71,6 +72,10 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
             e = self.task_exception
             raise TaskError(f'A processed task failed: {e.__class__.__name__}, {e}') from e
 
+    @property
+    def jobs_in_progress(self):
+        return self.max_concurrent_tasks - self.task_semaphore._value
+
     async def iter(self, *raw_queues: bytes, pop_timeout=1):
         """
         blpop jobs from redis queues and yield them. Waits for the number of tasks to drop below max_concurrent_tasks
@@ -88,17 +93,21 @@ async def iter(self, *raw_queues: bytes, pop_timeout=1):
             work_logger.debug('populating quit queue to prompt exit: %s', quit_queue.decode())
             await self.redis.rpush(quit_queue, b'1')
             raw_queues = tuple(raw_queues) + (quit_queue,)
-        while self.running:
+        while True:
+            await self.task_semaphore.acquire()
+            if not self.running:
+                break
             msg = await self.redis.blpop(*raw_queues, timeout=pop_timeout)
             if msg is None:
                 yield None, None
+                self.task_semaphore.release()
                 continue
             raw_queue, raw_data = msg
             if self.burst_mode and raw_queue == quit_queue:
                 work_logger.debug('got job from the quit queue, stopping')
                 break
+            work_logger.debug('jobs in progress %d', self.jobs_in_progress)
             yield raw_queue, raw_data
-            await self.wait()
 
     def add(self, coro, job, re_enqueue=False):
         """
@@ -115,18 +124,6 @@ def add(self, coro, job, re_enqueue=False):
         self.loop.call_later(self.timeout_seconds, self._cancel_job, task, job)
         self.pending_tasks.add(task)
 
-    async def wait(self):
-        """
-        Wait for a the number of pending tasks to drop bellow ``max_concurrent_tasks``
-        """
-        while True:
-            pt_cnt = len(self.pending_tasks)
-            if pt_cnt < self.max_concurrent_tasks:
-                return
-            work_logger.info('%d pending tasks, waiting for one to finish', pt_cnt)
-            _, self.pending_tasks = await asyncio.wait(self.pending_tasks, loop=self.loop,
-                                                       return_when=asyncio.FIRST_COMPLETED)
-
     async def finish(self, timeout=None):
         """
         Cancel all pending tasks and optionally re-enqueue jobs which haven't finished after the timeout.
@@ -150,6 +147,7 @@ async def finish(self, timeout=None):
                 self.pending_tasks = set()
 
     def _job_callback(self, task):
+        self.task_semaphore.release()
         self.jobs_complete += 1
         task_exception = task.exception()
         if task_exception:

arq/logs.py

@@ -9,13 +9,17 @@
 
 __all__ = ['ColourHandler', 'default_log_config']
 
-LOG_COLOURS = {
-    logging.DEBUG: 'white',
-    logging.INFO: 'green',
-    logging.WARN: 'yellow',
+LOG_FORMATS = {
+    logging.DEBUG: {'fg': 'white', 'dim': True},
+    logging.INFO: {'fg': 'green'},
+    logging.WARN: {'fg': 'yellow'},
 }
 
 
+def get_log_format(record):
+    return LOG_FORMATS.get(record.levelno, {'fg': 'red'})
+
+
 class ColourHandler(logging.StreamHandler):
     """
     Coloured log handler. Levels: debug: white, info: green, warning: yellow, else: red.
@@ -24,14 +28,13 @@ class ColourHandler(logging.StreamHandler):
     """
     def emit(self, record):
         log_entry = self.format(record)
-        colour = LOG_COLOURS.get(record.levelno, 'red')
         m = re.match('^(.*?: )', log_entry)
         if m:
             prefix = click.style(m.groups()[0], fg='magenta')
-            msg = click.style(log_entry[m.end():], fg=colour)
+            msg = click.style(log_entry[m.end():], **get_log_format(record))
             click.echo(prefix + msg)
         else:
-            click.secho(log_entry, fg=colour)
+            click.secho(log_entry, **get_log_format(record))
 
 
 def default_log_config(verbose: bool) -> dict:

arq/utils.py

@@ -106,6 +106,8 @@ def create_tz(utcoffset=0) -> timezone:
 
 def timestamp() -> float:
     """
+    This should be exactly the same as time.time(), we use this approach for consistency with
+    other methods and possibly greater accuracy.
     :return: now in unix time, eg. seconds since 1970
     """
     return (datetime.utcnow() - EPOCH).total_seconds()

tests/test_main.py

@@ -75,7 +75,9 @@ async def test_dispatch_work(tmpworkdir, loop, caplog, redis_conn):
             'recording health: <date time2> j_complete=0 j_failed=0 j_timedout=0 j_ongoing=0 q_high=1 q_dft=1 q_low=0\n'
             'starting main blpop loop\n'
             'populating quit queue to prompt exit: arq:quit-<random>\n'
+            'jobs in progress 1\n'
             'scheduling job <Job MockRedisDemoActor.high_add_numbers(3, 4, c=5) on high>, re-enqueue: False\n'
+            'jobs in progress 2\n'
             'scheduling job <Job MockRedisDemoActor.add_numbers(1, 2) on dft>, re-enqueue: False\n'
             'got job from the quit queue, stopping\n'
             'drain waiting 5.0s for 2 tasks to finish\n'

tests/test_worker.py

@@ -157,7 +157,7 @@ async def test_run_quit(tmpworkdir, redis_conn, actor, caplog):
     await worker.run()
     # the third job should be remove from the queue and readded
     assert tmpworkdir.join('save_slow').read() == '2'
-    assert '1 pending tasks, waiting for one to finish' in caplog
+    # assert '1 pending tasks, waiting for one to finish' in caplog
     assert 2 == await redis_conn.llen(b'arq:q:dft')