FIFO Set Queue for Redis (Python)

fifo_set_queue.py

__author__ = "Todd Lewellen"

"""
The MIT License (MIT)

Copyright (c) 2014 Todd Lewellen

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
import redis
import time


class SetQueue(object):
    """
    A Redis-backed Set Queue with FIFO ordering.  Useful for queueing up recurring tasks and ensuring that
    a task is only queued once at any given time.  If all tasks are equally important, yet certain tasks are
    produced (enqueued) more than others, this FIFO Set Queue may be a good solution.
    
    Redis transactions are used to help ensure atomicity.  
    
    This system relies on the assumption that the system times of the tasks producers are synced.
    
    Initial inspiration: http://www.rediscookbook.org/implement_a_fifo_queue.html
    """
    def __init__(self, host="localhost", port=6379, queue_name=None):
        self.r = redis.Redis(host=host, port=port)
        _id = queue_name or self.r.incr("queue_space")
        self.queue = "queue:{}".format(_id)

    def push(self, element):
        """Push an element to the tail of the queue"""
        self._zaddnx(element)

    def pop(self):
        """Pop an element from the head of the queue"""
        return self._zpop()

    def _zaddnx(self, element):
        pipe = self.r.pipeline()
        while 1:
            try:
                pipe.watch(self.queue)
                score = pipe.zscore(self.queue, element)
                if score is None:
                    pipe.multi()
                    pipe.zadd(self.queue, element, time.time())
                    pipe.execute()
                break
            except redis.WatchError:
                continue

    def _zpop(self):
        pipe = self.r.pipeline()
        while 1:
            try:
                pipe.watch(self.queue)
                rs = pipe.zrange(self.queue, 0, 0)
                if len(rs) > 0:
                    element = rs[0]
                    pipe.multi()
                    pipe.zrem(self.queue, element)
                    pipe.execute()
                    return element
            except (redis.WatchError, IndexError):
                continue