don't BLPOP each queue individually; do them all at once

Fixes issue binarymatt#85

Author: joeshaw

pyres/__init__.py

@@ -149,13 +149,16 @@ def push(self, queue, item):
         self.watch_queue(queue)
         self.redis.rpush("resque:queue:%s" % queue, ResQ.encode(item))
 
-    def pop(self, queue, timeout=10):
-        ret = self.redis.blpop("resque:queue:%s" % queue, timeout=timeout)
+    def pop(self, queues, timeout=10):
+        if isinstance(queues, basestring):
+            queues = [queues]
+        ret = self.redis.blpop(["resque:queue:%s" % q for q in queues],
+                               timeout=timeout)
         if ret:
-            if isinstance(ret, tuple):
-                q, ret = ret
-            return ResQ.decode(ret)
-        return ret
+            key, ret = ret
+            return key[13:], ResQ.decode(ret)  # trim "resque:queue:"
+        else:
+            return None, None
 
     def size(self, queue):
         return int(self.redis.llen("resque:queue:%s" % queue))
@@ -254,9 +257,9 @@ def keys(self):
         return [key.replace('resque:','')
                 for key in self.redis.keys('resque:*')]
 
-    def reserve(self, queue):
+    def reserve(self, queues):
         from pyres.job import Job
-        return Job.reserve(queue, self)
+        return Job.reserve(queues, self)
 
     def __str__(self):
         return "PyRes Client connected to %s" % self.redis.server

pyres/horde.py

@@ -73,12 +73,11 @@ def __str__(self):
         return '%s:%s:%s' % (self.hostname, self.pid, ','.join(self.queues))
 
     def reserve(self):
-        for q in self.queues:
-            self.logger.debug('checking queue: %s' % q)
-            job = Job.reserve(q, self.resq, self.__str__())
-            if job:
-                self.logger.info('Found job on %s' % q)
-                return job
+        self.logger.debug('checking queues: %s' % self.queues)
+        job = Job.reserve(self.queues, self.resq, self.__str__())
+        if job:
+            self.logger.info('Found job on %s' % job._queue)
+            return job
 
     def process(self, job):
         if not job:

pyres/job.py

@@ -83,11 +83,13 @@ def retry(self, payload_class, args):
         return False
 
     @classmethod
-    def reserve(cls, queue, res, worker=None, timeout=10):
-        """Reserve a job on the queue. This marks this job so that other workers
-        will not pick it up.
+    def reserve(cls, queues, res, worker=None, timeout=10):
+        """Reserve a job on one of the queues. This marks this job so
+        that other workers will not pick it up.
 
         """
-        payload = res.pop(queue, timeout=timeout)
+        if isinstance(queues, basestring):
+            queues = [queues]
+        queue, payload = res.pop(queues, timeout=timeout)
         if payload:
             return cls(queue, payload, res, worker)

pyres/worker.py

@@ -221,12 +221,11 @@ def process(self, job=None):
             self.done_working()
 
     def reserve(self, timeout=10):
-        for q in self.queues:
-            logger.debug('checking queue %s' % q)
-            job = self.job_class.reserve(q, self.resq, self.__str__(), timeout=timeout)
-            if job:
-                logger.info('Found job on %s' % q)
-                return job
+        logger.debug('checking queues %s' % self.queues)
+        job = self.job_class.reserve(self.queues, self.resq, self.__str__(), timeout=timeout)
+        if job:
+            logger.info('Found job on %s' % job._queue)
+            return job
 
     def working_on(self, job):
         logger.debug('marking as working on')

tests/test_resq.py

@@ -22,10 +22,27 @@ def test_pop(self):
         self.resq.push('pushq','content-newqueue')
         self.resq.push('pushq','content2-newqueue')
         assert self.redis.llen('resque:queue:pushq') == 2
-        assert self.resq.pop('pushq') == 'content-newqueue'
+        assert self.resq.pop('pushq') == ('pushq', 'content-newqueue')
         assert self.redis.llen('resque:queue:pushq') == 1
-        assert self.resq.pop('pushq') == 'content2-newqueue'
+        assert self.resq.pop(['pushq']) == ('pushq', 'content2-newqueue')
         assert self.redis.llen('resque:queue:pushq') == 0
+
+    def test_pop_two_queues(self):
+        self.resq.push('pushq1', 'content-q1-1')
+        self.resq.push('pushq1', 'content-q1-2')
+        self.resq.push('pushq2', 'content-q2-1')
+        assert self.redis.llen('resque:queue:pushq1') == 2
+        assert self.redis.llen('resque:queue:pushq2') == 1
+        assert self.resq.pop(['pushq1', 'pushq2']) == ('pushq1', 'content-q1-1')
+        assert self.redis.llen('resque:queue:pushq1') == 1
+        assert self.redis.llen('resque:queue:pushq2') == 1
+        assert self.resq.pop(['pushq2', 'pushq1']) == ('pushq2', 'content-q2-1')
+        assert self.redis.llen('resque:queue:pushq1') == 1
+        assert self.redis.llen('resque:queue:pushq2') == 0
+        assert self.resq.pop(['pushq2', 'pushq1']) == ('pushq1', 'content-q1-2')
+        assert self.redis.llen('resque:queue:pushq1') == 0
+        assert self.redis.llen('resque:queue:pushq2') == 0
+        assert self.resq.pop(['pushq1', 'pushq2'], timeout=1) == (None, None)
 
     def test_peek(self):
         self.resq.enqueue(Basic,"test1")