Django Rest Framework 源碼解析--節流

Django Rest Framework 源碼解析--節流

接上一篇博客，restframework重寫的dispatch()方法中，執行了inital()函數。inital()中check_throttles((request) 方法實現了請求的訪問頻率控制功能。

1、check_throttles(request)函數中，循環了限流類的對象列表，依次執行限流對象的 allow_request() 方法

def check_throttles(self, request):
    """
    Check if request should be throttled.
    Raises an appropriate exception if the request is throttled.
    
    檢查請求是否應該被節流。
    如果被節流，則拋出相應異常。
    """
    # 遍歷 限流對象列表，如果返回Fasle,則被限流，拋出異常（可傳參數throttle.wait()的返回值）
    for throttle in self.get_throttles():
        # 如果被節流，返回False，則拋出相應異常
        if not throttle.allow_request(request, self):
            self.throttled(request, throttle.wait())

2、self.get_throttles()限流對象列表生成式

def get_throttles(self):
    """
    Instantiates and returns the list of throttles that this view uses.
    """
    # 可以在view中重寫throttle_classes，指定限流對象列表
    # 也可以在setting.py中定義
    return [throttle() for throttle in self.throttle_classes]

3、如果被節流，則通過throttled()方法，拋出相應異常，可傳入wait作爲等待時間的參數

def throttled(self, request, wait):
    """
    If request is throttled, determine what kind of exception to raise.
    """
    # wait參數，出入的值是 節流類的wait()方法的返回值（單位：秒）
    raise exceptions.Throttled(wait)

使用示例：

1、自定義限流類，繼承BaseThrottle類，重寫 allow_request() 和 wait() 這兩個方法

from rest_framework.throttling import BaseThrottle

import time

# 保存訪問記錄
VISIT_RECORD = {}

class VisitThrottle(BaseThrottle):
    """
    自定義限流類：60秒內只能訪問3次，超過就限流
    
    返回True，允許請求訪問
    返回False，禁止請求訪問
    """
    
    
    # 通過 self.history這個對象的成員變量，
    # 在allow_request()和 wait()這兩個成員方法之間傳遞history的值
    def __init__(self):
        self.history = None  # 初始化訪問記錄

    def allow_request(self, request, view):
        # 獲取用戶ip  
        remote_addr = self.get_ident(request)
        timer = time.time()
        if remote_addr not in REQ_RECORD:
            VISIT_RECORD[remote_addr]=[timer]
            return True
        # 獲取當前ip的歷史訪問記錄
        history = VISIT_RECORD[remote_addr]
        self.history = history
        # 如果有歷史訪問記錄，並且最早一次的訪問記錄離當前時間超過60s，就刪除最早的那個訪問記錄，
        # 只要爲True，就一直循環刪除最早的一次訪問記錄
        while history and history[-1] < timer - 60:
            history.pop()

        # 60秒內的訪問記錄，是否超過3次
        # 如果沒有超過，則記錄這次訪問，並返回True,允許訪問
        # 如果超過，則返回False,禁止訪問
        if len(history) < 3:
            history.insert(0, timer)
            return True
        return False

    def wait(self):
        '''還需要等多久才能訪問'''
        timer = time.time()
        return 60 - (timer - self.history[-1])

2、在View中調用限流類

class TestView(APIView):
    # 在View中重寫throttle_classes限流類列表，一般只寫一個限流，
    # 或者不限流，使列表爲空，throttle_classes = []
    throttle_classes = [VisitThrottle, ]

    def get(self,request,*args,**kwargs):
        pass

3、或者在setting.py中指定全站默認使用的限流類的路徑

REST_FRAMEWORK = {
    "DEFAULT_THROTTLE_CLASSES": ('plugin.restframework.throttling.VisitThrottle')
}

內置的限流類

與認證類、鑑權類通常繼承BaseXXX類，高度自定義不同，限流類我個人覺得繼承restframework提供的其他內置的類更方便

例如繼承 SimpleRateThrottle 類

class TestTrottle(SimpleRateThrottle):
    # 設定規定時間內能訪問的次數，例如 3/m, 1/s, 1000/h, 9999/day
    # 通常設定在setting.py中
    THROTTLE_RATES = {
        "Test": '5/m'
    }
    # 指定scope值爲 查找THROTTLE_RATES的key
    scope = "Test"
    
    def get_cache_key(self, request, view):
        # 通過ip限制節流
        return self.get_ident(request)
        
        # return request.user.pk   # 通常也使用requser.user作爲標識一個用戶的ID