在Django中实现一个高性能未读消息计数器

挑战：实时更新你的计数器

为了让我们的计数器正常的工作，我们必须实时的更新它，这包括：

1. 当有新的未读消息过来的时候，为计数器 +1
2. 当消息被异常删除时，如果关联的消息为未读，为计数器 -1
3. 当阅读完一个新消息的时候，为计数器 -1

让我们一个一个来解决这些情况。

在抛出解决方案之前，我们需要先介绍Django中的一个功能： Signals ，Signals是django提供的一个事件通知机制，它可以让你在监听某些自定义或者 预设的事件，当这些事件发生的时候，调用实现定义好的方法。

比如 django.db.models.signals.pre_save & django.db.models.signals.post_save 表示的是 某个Model调用save方法之前和之后会触发的事件，它和Database提供的触发器在功能上有一点相似。

关于Signals更多的介绍可以参考官方文档，下面让我们来看看Signals能给我们的计数器带来什么好处。

from django.db.models.signals import post_save, post_delete

def incr_notifications_counter(sender, instance, created, **kwargs):
    # 只有当这个instance是新创建，而且has_readed是默认的false才更新
    if not (created and not instance.has_readed):
        return

    # 调用 update_unread_count 方法来更新计数器 +1
    NotificationController(instance.user_id).update_unread_count(1)

# 监听Notification Model的post_save信号
post_save.connect(incr_notifications_counter, sender=Notification)

def decr_notifications_counter(sender, instance, **kwargs):
    # 当删除的消息还没有被读过时，计数器 -1
    if not instance.has_readed:
        NotificationController(instance.user_id).update_unread_count(-1)

post_delete.connect(decr_notifications_counter, sender=Notification)

class NotificationController(object):

    ... ...

    def mark_as_readed(self, notification_id):
        notification = Notification.objects.get(pk=notification_id)
        # 没有必要重复标记一个已经读过的通知
        if notication.has_readed:
            return

        notification.has_readed = True
        notification.save()
        # 在这里更新我们的计数器，嗯，我感觉好极了
        self.update_unread_count(-1)

# 因为两个并发的请求，假设这两个方法几乎同时被调用
NotificationController(user_id).mark_as_readed(100)
NotificationController(user_id).mark_as_readed(100)

from django.db import transaction

class NotificationController(object):

    ... ...

    def mark_as_readed(self, notification_id):
        # 手动让select for update和update语句发生在一个完整的事务里面
        with transaction.commit_on_success():
            # 使用select_for_update来保证并发请求同时只有一个请求在处理，其他的请求
            # 等待锁释放
            notification = Notification.objects.select_for_update().get(pk=notification_id)
            # 没有必要重复标记一个已经读过的通知
            if notication.has_readed:
                return

            notification.has_readed = True
            notification.save()
            # 在这里更新我们的计数器，嗯，我感觉好极了
            self.update_unread_count(-1)

def mark_as_readed(self, notification_id):
        affected_rows = Notification.objects.filter(pk=notification_id, has_readed=False)\
                                            .update(has_readed=True)
        # affected_rows将会返回update语句修改的条目数
        self.update_unread_count(affected_rows)

使用sorted set来缓存计数器改动

redis是一个非常好用的内存数据库，其中的sorted set是它提供的一种数据类型：有序集合， 使用它，我们可以非常简单的缓存所有的计数器改动，然后批量回写到数据库。

RK_NOTIFICATIONS_COUNTER = 'ss_pending_counter_changes'

def update_unread_count(self, count):
    """修改过的update_unread_count方法"""
    redisdb.zincrby(RK_NOTIFICATIONS_COUNTER, str(self.user_id), count)

# 同时我们也需要修改获取用户未读消息数方法，使其获取redis中那些没有被回写
# 到数据库的缓冲区数据。在这里代码就省略了

通过以上的代码，我们把计数器的更新缓冲在了redis里面，我们还需要一个脚本来把这个缓冲区 里面的数据定时回写到数据库中。

通过自定义django的command，我们可以非常轻松的做到这一点：

# File: management/commands/notification_update_counter.py

# -*- coding: utf-8 -*-
from django.core.management.base import BaseCommand
from django.db.models import F

# Fix import prob
from notification.models import UserNotificationsCount
from notification.utils import RK_NOTIFICATIONS_COUNTER
from base_redis import redisdb

import logging
logger = logging.getLogger('stdout')


class Command(BaseCommand):
    help = 'Update UserNotificationsCounter objects, Write changes from redis to database'

    def handle(self, *args, **options):
        # 首先，通过 zrange 命令来获取缓冲区所有修改过的用户ID
        for user_id in redisdb.zrange(RK_NOTIFICATIONS_COUNTER, 0, -1):
            # 这里值得注意，为了保证操作的原子性，我们使用了redisdb的pipeline
            pipe = redisdb.pipeline()
            pipe.zscore(RK_NOTIFICATIONS_COUNTER, user_id)
            pipe.zrem(RK_NOTIFICATIONS_COUNTER, user_id)
            count, _ = pipe.execute()
            count = int(count)
            if not count:
                continue

            logger.info('Updating unread count user %s: count %s' % (user_id, count))
            UserNotificationsCount.objects.filter(pk=obj.pk)\
                                          .update(unread_count=F('unread_count') + count)

之后，通过 python manage.py notification_update_counter 这样的命令就可以把缓冲区 里面的改动批量回写到数据库了。我们还可以把这个命令配置到crontab中来定义执行。