pdk-rate-limit-lib 1.7.0

PDK Rate Limit Library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use pdk_core::log::debug;
use pdk_core::policy_context::api::Tier;
use serde::{Deserialize, Serialize};
use std::cmp::max;

use super::super::distribution_formula::DistributionFormula;
use super::{LimitStats, QuotaInfo, RequestAllowed};

/// Cluster bucket for the rate limiting.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ClusterBucket {
    limits: Vec<ClusterLimit>,
}

impl ClusterBucket {
    /// Create a new cluster bucket with the specified start time and tiers.
    pub fn new(start_time: u128, tiers: &[Tier]) -> Self {
        Self {
            limits: tiers
                .iter()
                .map(|tier| ClusterLimit::from(start_time, tier))
                .collect(),
        }
    }

    /// To use in case of the bucket being used in cluster mode
    pub fn request_allowed(&mut self, now: u128, amount: usize) -> RequestAllowed {
        self.set_time(now);

        if self.limits.iter().all(|l| l.remaining(amount)) {
            self.limits.iter_mut().for_each(|l| l.allow_request(amount));
            RequestAllowed::Allowed
        } else if self.limits.iter().all(|l| l.cluster_remaining(amount)) {
            RequestAllowed::OutOfLocalQuota
        } else {
            RequestAllowed::OutOfQuota
        }
    }

    /// Get the status of the cluster bucket.
    pub fn status(&self) -> Option<LimitStats> {
        self.limits
            .iter()
            .map(ClusterLimit::stats)
            .reduce(LimitStats::most_restrictive)
    }

    /// To use in the primary bucket to get quota.
    pub fn get_quota(
        &mut self,
        now: u128,
        bucket: &ClusterBucket,
        formula: &DistributionFormula,
        amount: usize,
    ) -> Vec<QuotaInfo> {
        self.set_time(now);

        let mut quota_info = Vec::new();
        for (cluster, local) in self.limits.iter_mut().zip(bucket.limits.iter()) {
            if local.available >= amount as u64 && local.reset == cluster.reset {
                quota_info.push(QuotaInfo::new(0, cluster.cluster_remaining, cluster.reset))
            } else {
                quota_info.push(cluster.get_quota_portion(max(
                    formula.portion(cluster),
                    amount as u64 - local.available,
                )));
            }
        }
        quota_info
    }

    /// To use in the secondary bucket to get quota.
    pub fn update_quota(&mut self, quota: &[QuotaInfo]) {
        //TODO: this will misbehave if the sizes of the vecs are different
        for (local, quota) in self.limits.iter_mut().zip(quota.iter()) {
            local.update_quota(quota)
        }
    }

    fn set_time(&mut self, now: u128) {
        self.limits.iter_mut().for_each(|limit| limit.set_time(now));
    }
}

/// Cluster limit for the rate limiting.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ClusterLimit {
    /// Amount of requests available in the current period.
    pub available: u64,
    /// Amount of requests remaining in the cluster.
    pub cluster_remaining: u64,
    /// Maximum requests allowed in a single period.
    pub max_requests: u64,
    /// Length of the time window in milliseconds.
    pub period: u128,
    /// Timestamp when the rate limit will reset.
    pub reset: u128,
}

impl ClusterLimit {
    fn from(start_time: u128, tier: &Tier) -> Self {
        Self {
            available: 0,
            cluster_remaining: tier.requests,
            max_requests: tier.requests,
            period: tier.period_in_millis as u128,
            reset: start_time + tier.period_in_millis as u128,
        }
    }

    fn stats(&self) -> LimitStats {
        LimitStats::new(
            self.available + self.cluster_remaining,
            self.max_requests,
            self.reset,
        )
    }

    fn set_time(&mut self, now: u128) {
        while self.reset <= now {
            self.reset += self.period;
            self.available = 0;
            self.cluster_remaining = self.max_requests;
        }
    }

    fn update_quota(&mut self, quota: &QuotaInfo) {
        if quota.reset != self.reset {
            self.available = quota.given;
            self.reset = quota.reset;
            self.cluster_remaining = quota.remaining;
        } else {
            self.available += quota.given;
            self.cluster_remaining = quota.remaining;
        }
    }

    fn remaining(&self, amount: usize) -> bool {
        self.available >= amount as u64
    }

    fn cluster_remaining(&self, amount: usize) -> bool {
        self.available + self.cluster_remaining >= amount as u64
    }

    fn allow_request(&mut self, amount: usize) {
        self.available -= amount as u64;
    }

    fn get_quota_portion(&mut self, portion: u64) -> QuotaInfo {
        let mut given: u64 = portion;

        if given > self.cluster_remaining {
            given = self.cluster_remaining;
        }

        if given < 1 && self.cluster_remaining >= 1 {
            given = 1;
        }

        self.cluster_remaining -= given;
        debug!(
            "Portion requested: {}, given: {}, remaining: {}",
            portion, given, self.cluster_remaining
        );
        QuotaInfo::new(given, self.cluster_remaining, self.reset)
    }
}