基于 Redis 的优先级队列

问题

如何使用 Redis 实现一个优先级队列呢？

sorted set 简介

Sorted set 是 Redis 提供的一种数据结构类型，兼具 set 和 hash 的特点。首先，sorted set 中的每个元素是唯一的；其次，sorted set 中的每个元素都关联一个浮点值（叫做 score）。除此之外，sorted set 中的元素是有序的，这些元素依照如下规则排序：

• 如果A和B是两个 score 不一样的元素，则当且仅当 A.score > B.score 的时候，A > B
• 如果A和B的 score 一样，则当且仅当A的字典序大于B的时候，A > B

sorted set 操作

• ZADD：向 sorted set 中添加元素
• ZCOUNT： sorted set 中 score 等于指定值的元素有多少个
• ZSCORE：sorted set 中指定元素的 score 是多少
• ZCARD： sorted set 中总共有多少个元素
• ZREM：删除 sorted set 中的指定元素
• ZREVRANGE：按照从大到小的顺序返回指定索引区间内的元素
• ZRANGE: 按照从小到大的顺序返回指定索引区间内的元素

sorted set 支持的操作还有很多，上面只是列举了一些对实现优先级队列有用的操作。

优先级队列实现

基于 sorted set，可以实现一个简单的优先级队列。代码如下：

class PriorityQueue(object):

    def __init__(self, redis_client, queue, namespace='priority-queue'):
        self.client = redis_client
        self.key = '%s:%s' % (namespace, queue)

    def enqueue(self, score, member):
        return self.client.zadd(self.key, score, member)

    def dequeue(self):
        """注意：该方法不是线程安全的"""
        score = None
        member = None
        result = self.client.zrevrange(self.key, 0, 0, withscores=True)
        # [( member, score), (member, score), ...]
        if result:
            member, score = result[0]
            ret = self.client.zrem(self.key, member)
            assert ret == 1
        return score, member

    def card(self):
        return self.client.zcard(self.key)

    def __repr__(self):
        return u"PriorityQueue<%s>" % self.key

优先级队列主要操作是入队和出队，sorted set 根据元素的 score 维护了优先级顺序。需要注意的是，上述代码中的出队操作不是线程安全的，因为取优先级最高的元素以及删除这个元素是两次操作，不是原子性的。