基于 Redis 的优先级队列

妖狐艹你老母 2021-10-30 08:20 455阅读 0赞

## 问题 ##

如何使用 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 维护了优先级顺序。需要注意的是，上述代码中的出队操作不是线程安全的，因为取优先级最高的元素以及删除这个元素是两次操作，不是原子性的。