throttle/
limit.rs

1use crate::{Error, LimitSpec, REDIS};
2use anyhow::{anyhow, Context};
3use mod_redis::{RedisConnection, Script};
4use parking_lot::Mutex;
5use std::collections::HashMap;
6use std::sync::{Arc, LazyLock};
7use std::time::{Duration, Instant, SystemTime};
8use tokio::sync::Notify;
9use uuid::Uuid;
10
11static MEMORY: LazyLock<Mutex<MemoryStore>> = LazyLock::new(|| Mutex::new(MemoryStore::new()));
12
13static ACQUIRE_SCRIPT: LazyLock<Script> = LazyLock::new(|| Script::new(include_str!("limit.lua")));
14
15pub struct LimitSpecWithDuration {
16    pub spec: LimitSpec,
17    /// Maximum lease duration for a single count
18    pub duration: Duration,
19}
20
21#[derive(Debug)]
22pub struct LimitLease {
23    /// Name of the element to release on Drop
24    name: String,
25    uuid: Uuid,
26    armed: bool,
27    backend: Backend,
28}
29
30#[derive(Debug, PartialEq, Clone, Copy)]
31enum Backend {
32    Memory,
33    Redis,
34}
35
36impl LimitSpecWithDuration {
37    pub async fn acquire_lease<S: AsRef<str>>(
38        &self,
39        key: S,
40        deadline: Instant,
41    ) -> Result<LimitLease, Error> {
42        match (self.spec.force_local, REDIS.get()) {
43            (false, Some(redis)) => {
44                self.acquire_lease_redis(redis, key.as_ref(), deadline)
45                    .await
46            }
47            (true, _) | (false, None) => self.acquire_lease_memory(key.as_ref(), deadline).await,
48        }
49    }
50
51    pub async fn acquire_lease_redis(
52        &self,
53        conn: &RedisConnection,
54        key: &str,
55        deadline: Instant,
56    ) -> Result<LimitLease, Error> {
57        loop {
58            let now_ts = SystemTime::now()
59                .duration_since(SystemTime::UNIX_EPOCH)
60                .map(|d| d.as_secs_f64())
61                .unwrap_or(0.0);
62
63            let expires_ts = now_ts + self.duration.as_secs_f64();
64            let uuid = Uuid::new_v4();
65            let uuid_str = uuid.to_string();
66
67            let mut script = ACQUIRE_SCRIPT.prepare_invoke();
68            script
69                .key(key)
70                .arg(now_ts)
71                .arg(expires_ts)
72                .arg(self.spec.limit)
73                .arg(uuid_str);
74
75            match conn.invoke_script(script).await.with_context(|| {
76                format!(
77                    "error invoking redis lease acquisition script \
78                     key={key} now={now_ts} expires={expires_ts} \
79                     limit={} uuid={uuid}",
80                    self.spec.limit
81                )
82            })? {
83                mod_redis::RedisValue::Okay => {
84                    return Ok(LimitLease {
85                        name: key.to_string(),
86                        uuid,
87                        armed: true,
88                        backend: Backend::Redis,
89                    });
90                }
91                mod_redis::RedisValue::Int(next_expiration_interval) => {
92                    if Instant::now() >= deadline {
93                        return Err(Error::TooManyLeases(Duration::from_secs(
94                            next_expiration_interval as u64,
95                        )));
96                    }
97
98                    tokio::time::sleep(Duration::from_secs(3)).await;
99                }
100                mod_redis::RedisValue::Double(next_expiration_interval) => {
101                    if Instant::now() >= deadline {
102                        return Err(Error::TooManyLeases(Duration::from_secs(
103                            next_expiration_interval as u64,
104                        )));
105                    }
106
107                    tokio::time::sleep(Duration::from_secs(3)).await;
108                }
109                value => {
110                    return Err(anyhow!("acquire script succeeded but returned {value:?}").into());
111                }
112            }
113        }
114    }
115
116    pub async fn acquire_lease_memory(
117        &self,
118        key: &str,
119        deadline: Instant,
120    ) -> Result<LimitLease, Error> {
121        let uuid = Uuid::new_v4();
122
123        fn resolve_set(key: &str) -> Arc<LeaseSet> {
124            MEMORY.lock().get_or_create(key)
125        }
126
127        let set = resolve_set(key);
128
129        set.acquire(uuid, self.spec.limit, self.duration, deadline)
130            .await?;
131
132        Ok(LimitLease {
133            name: key.to_string(),
134            uuid,
135            armed: true,
136            backend: Backend::Memory,
137        })
138    }
139}
140
141impl LimitLease {
142    pub fn name(&self) -> &str {
143        &self.name
144    }
145
146    pub async fn release(&mut self) {
147        self.armed = false;
148        match self.backend {
149            Backend::Memory => self.release_memory().await,
150            Backend::Redis => {
151                if let Some(redis) = REDIS.get() {
152                    self.release_redis(redis).await;
153                } else {
154                    eprintln!("LimitLease::release: backend is Redis but REDIS is not set");
155                }
156            }
157        }
158    }
159
160    pub async fn extend(&self, duration: Duration) -> Result<(), Error> {
161        match self.backend {
162            Backend::Memory => self.extend_memory(duration).await,
163            Backend::Redis => {
164                if let Some(redis) = REDIS.get() {
165                    self.extend_redis(redis, duration).await
166                } else {
167                    Err(anyhow::anyhow!(
168                        "LimitLease::extend: backend is Redis but REDIS is not set"
169                    )
170                    .into())
171                }
172            }
173        }
174    }
175
176    pub fn take(&mut self) -> Self {
177        let armed = self.armed;
178        self.armed = false;
179        Self {
180            name: self.name.clone(),
181            uuid: self.uuid,
182            armed,
183            backend: self.backend,
184        }
185    }
186
187    async fn extend_memory(&self, duration: Duration) -> Result<(), Error> {
188        let store = MEMORY.lock();
189        if let Some(set) = store.get(&self.name) {
190            set.extend(self.uuid, duration)
191        } else {
192            Err(Error::NonExistentLease)
193        }
194    }
195
196    async fn extend_redis(&self, conn: &RedisConnection, duration: Duration) -> Result<(), Error> {
197        let now_ts = SystemTime::now()
198            .duration_since(SystemTime::UNIX_EPOCH)
199            .map(|d| d.as_secs_f64())
200            .unwrap_or(0.0);
201
202        let expires = now_ts + duration.as_secs_f64();
203
204        let mut cmd = mod_redis::cmd("ZADD");
205        cmd.arg(&self.name)
206            .arg("XX") // only allow updating existing
207            .arg("CH") // return number of changed entries
208            .arg(expires)
209            .arg(self.uuid.to_string());
210        let value = conn.query(cmd).await?;
211
212        if value != mod_redis::RedisValue::Int(1) {
213            return Err(anyhow!("Failed to extend lease").into());
214        }
215
216        Ok(())
217    }
218
219    async fn release_memory(&self) {
220        let store = MEMORY.lock();
221        if let Some(set) = store.get(&self.name) {
222            set.release(self.uuid);
223        }
224    }
225
226    async fn release_redis(&mut self, conn: &RedisConnection) {
227        let mut cmd = mod_redis::cmd("ZREM");
228        cmd.arg(&self.name).arg(self.uuid.to_string());
229        conn.query(cmd).await.ok();
230    }
231}
232
233impl Drop for LimitLease {
234    fn drop(&mut self) {
235        if self.armed {
236            self.armed = false;
237            let mut deferred = Self {
238                armed: false,
239                name: self.name.clone(),
240                uuid: self.uuid,
241                backend: self.backend,
242            };
243            tokio::task::Builder::new()
244                .name("LimitLeaseDropper")
245                .spawn(async move {
246                    deferred.release().await;
247                })
248                .ok();
249        }
250    }
251}
252
253struct LeaseSet {
254    members: Mutex<HashMap<Uuid, Instant>>,
255    notify: Notify,
256}
257
258impl LeaseSet {
259    fn new() -> Self {
260        Self {
261            members: Mutex::new(HashMap::new()),
262            notify: Notify::new(),
263        }
264    }
265
266    fn acquire_immediate(&self, uuid: Uuid, limit: u64, duration: Duration) -> bool {
267        let mut members = self.members.lock();
268        let now = Instant::now();
269        members.retain(|_k, expiry| *expiry > now);
270
271        if members.len() as u64 + 1 <= limit {
272            members.insert(uuid, now + duration);
273            return true;
274        }
275
276        false
277    }
278
279    async fn acquire(
280        &self,
281        uuid: Uuid,
282        limit: u64,
283        duration: Duration,
284        deadline: Instant,
285    ) -> Result<(), Error> {
286        loop {
287            if self.acquire_immediate(uuid, limit, duration) {
288                return Ok(());
289            }
290
291            match tokio::time::timeout_at(deadline.into(), self.notify.notified()).await {
292                Err(_) => {
293                    if self.acquire_immediate(uuid, limit, duration) {
294                        return Ok(());
295                    }
296                    let min_expiration = self
297                        .members
298                        .lock()
299                        .values()
300                        .cloned()
301                        .min()
302                        .expect("some elements");
303                    return Err(Error::TooManyLeases(min_expiration - Instant::now()));
304                }
305                Ok(_) => {
306                    // Try to acquire again
307                    continue;
308                }
309            }
310        }
311    }
312
313    fn extend(&self, uuid: Uuid, duration: Duration) -> Result<(), Error> {
314        match self.members.lock().get_mut(&uuid) {
315            Some(entry) => {
316                *entry = Instant::now() + duration;
317                Ok(())
318            }
319            None => Err(Error::NonExistentLease),
320        }
321    }
322
323    fn release(&self, uuid: Uuid) {
324        let mut members = self.members.lock();
325        members.remove(&uuid);
326        self.notify.notify_one();
327    }
328}
329
330struct MemoryStore {
331    sets: HashMap<String, Arc<LeaseSet>>,
332}
333
334impl MemoryStore {
335    fn new() -> Self {
336        Self {
337            sets: HashMap::new(),
338        }
339    }
340
341    fn get(&self, name: &str) -> Option<Arc<LeaseSet>> {
342        self.sets.get(name).map(Arc::clone)
343    }
344
345    fn get_or_create(&mut self, name: &str) -> Arc<LeaseSet> {
346        self.sets
347            .entry(name.to_string())
348            .or_insert_with(|| Arc::new(LeaseSet::new()))
349            .clone()
350    }
351}
352
353#[cfg(test)]
354mod test {
355    use super::*;
356    use mod_redis::test::{RedisCluster, RedisServer};
357
358    #[tokio::test]
359    async fn test_memory() {
360        let limit = LimitSpecWithDuration {
361            spec: LimitSpec::new(2),
362            duration: Duration::from_secs(2),
363        };
364
365        let key = format!("test_memory-{}", Uuid::new_v4());
366        let lease1 = limit
367            .acquire_lease_memory(&key, Instant::now())
368            .await
369            .unwrap();
370        eprintln!("lease1: {lease1:?}");
371        let mut lease2 = limit
372            .acquire_lease_memory(&key, Instant::now())
373            .await
374            .unwrap();
375        eprintln!("lease2: {lease2:?}");
376        // Cannot acquire a 3rd lease while the other two are alive
377        assert!(limit
378            .acquire_lease_memory(&key, Instant::now())
379            .await
380            .is_err());
381
382        // Release and try to get a third
383        lease2.release().await;
384        let _lease3 = limit
385            .acquire_lease_memory(&key, Instant::now())
386            .await
387            .unwrap();
388
389        // Cannot acquire while the other two are alive
390        assert!(limit
391            .acquire_lease_memory(&key, Instant::now())
392            .await
393            .is_err());
394
395        let start = Instant::now();
396
397        // We can acquire another after waiting for some number of leases to expire
398        let _lease4 = limit
399            .acquire_lease_memory(&key, start + limit.duration + limit.duration)
400            .await
401            .unwrap();
402
403        assert!(
404            start.elapsed() > limit.duration,
405            "elapsed is {:?}",
406            start.elapsed()
407        );
408    }
409
410    #[tokio::test]
411    async fn test_redis() {
412        if !RedisServer::is_available() {
413            return;
414        }
415        let redis = RedisServer::spawn("").await.unwrap();
416        let conn = redis.connection().await.unwrap();
417
418        let limit = LimitSpecWithDuration {
419            spec: LimitSpec::new(2),
420            duration: Duration::from_secs(2),
421        };
422
423        let key = format!("test_redis-{}", Uuid::new_v4());
424        let mut lease1 = limit
425            .acquire_lease_redis(&conn, &key, Instant::now())
426            .await
427            .unwrap();
428        eprintln!("lease1: {lease1:?}");
429        let mut lease2 = limit
430            .acquire_lease_redis(&conn, &key, Instant::now())
431            .await
432            .unwrap();
433        eprintln!("lease2: {lease2:?}");
434        // Cannot acquire a 3rd lease while the other two are alive
435        assert!(limit
436            .acquire_lease_redis(&conn, &key, Instant::now())
437            .await
438            .is_err());
439
440        // Release and try to get a third
441        lease2.release_redis(&conn).await;
442        let mut lease3 = limit
443            .acquire_lease_redis(&conn, &key, Instant::now())
444            .await
445            .unwrap();
446
447        // Cannot acquire while the other two are alive
448        assert!(limit
449            .acquire_lease_redis(&conn, &key, Instant::now())
450            .await
451            .is_err());
452
453        let start = Instant::now();
454
455        // We can acquire another after waiting for some number of leases to expire
456        let mut lease4 = limit
457            .acquire_lease_redis(&conn, &key, start + limit.duration + limit.duration)
458            .await
459            .unwrap();
460
461        assert!(
462            start.elapsed() > limit.duration,
463            "elapsed is {:?}",
464            start.elapsed()
465        );
466
467        lease1.release_redis(&conn).await;
468        lease3.release_redis(&conn).await;
469        lease4.release_redis(&conn).await;
470    }
471
472    #[tokio::test]
473    async fn test_redis_cluster() {
474        if !RedisCluster::is_available().await {
475            return;
476        }
477        let redis = RedisCluster::spawn().await.unwrap();
478        let conn = redis.connection().await.unwrap();
479
480        let limit = LimitSpecWithDuration {
481            spec: LimitSpec::new(2),
482            duration: Duration::from_secs(2),
483        };
484
485        let key = format!("test_redis-{}", Uuid::new_v4());
486        let mut lease1 = limit
487            .acquire_lease_redis(&conn, &key, Instant::now())
488            .await
489            .unwrap();
490        eprintln!("lease1: {lease1:?}");
491        let mut lease2 = limit
492            .acquire_lease_redis(&conn, &key, Instant::now())
493            .await
494            .unwrap();
495        eprintln!("lease2: {lease2:?}");
496        // Cannot acquire a 3rd lease while the other two are alive
497        assert!(limit
498            .acquire_lease_redis(&conn, &key, Instant::now())
499            .await
500            .is_err());
501
502        // Release and try to get a third
503        lease2.release_redis(&conn).await;
504        let mut lease3 = limit
505            .acquire_lease_redis(&conn, &key, Instant::now())
506            .await
507            .unwrap();
508
509        // Cannot acquire while the other two are alive
510        assert!(limit
511            .acquire_lease_redis(&conn, &key, Instant::now())
512            .await
513            .is_err());
514
515        // Wait for some number of leases to expire
516        tokio::time::sleep(limit.duration + limit.duration).await;
517
518        let mut lease4 = limit
519            .acquire_lease_redis(&conn, &key, Instant::now())
520            .await
521            .unwrap();
522
523        lease1.release_redis(&conn).await;
524        lease3.release_redis(&conn).await;
525        lease4.release_redis(&conn).await;
526    }
527
528    #[tokio::test]
529    async fn test_memory_extension() {
530        let limit = LimitSpecWithDuration {
531            spec: LimitSpec::new(1),
532            duration: Duration::from_secs(2),
533        };
534
535        let key = format!("test_redis-{}", Uuid::new_v4());
536        let lease1 = limit
537            .acquire_lease_memory(&key, Instant::now())
538            .await
539            .unwrap();
540        eprintln!("lease1: {lease1:?}");
541        // Cannot acquire a 2nd lease while the first is are alive
542        assert!(limit
543            .acquire_lease_memory(&key, Instant::now())
544            .await
545            .is_err());
546
547        tokio::time::sleep(Duration::from_secs(1)).await;
548
549        lease1.extend_memory(Duration::from_secs(6)).await.unwrap();
550
551        // Wait for original lease duration to expire
552        tokio::time::sleep(limit.duration + limit.duration).await;
553
554        // Cannot acquire because we have an extended lease
555        assert!(limit
556            .acquire_lease_memory(&key, Instant::now())
557            .await
558            .is_err());
559
560        // Wait for extension to pass
561        tokio::time::sleep(limit.duration + limit.duration).await;
562
563        let _lease2 = limit
564            .acquire_lease_memory(&key, Instant::now())
565            .await
566            .unwrap();
567    }
568
569    #[tokio::test]
570    async fn test_redis_extension() {
571        if !RedisServer::is_available() {
572            return;
573        }
574        let redis = RedisServer::spawn("").await.unwrap();
575        let conn = redis.connection().await.unwrap();
576
577        let limit = LimitSpecWithDuration {
578            spec: LimitSpec::new(1),
579            duration: Duration::from_secs(2),
580        };
581
582        let key = format!("test_redis-{}", Uuid::new_v4());
583        let mut lease1 = limit
584            .acquire_lease_redis(&conn, &key, Instant::now())
585            .await
586            .unwrap();
587        eprintln!("lease1: {lease1:?}");
588        // Cannot acquire a 2nd lease while the first is are alive
589        assert!(limit
590            .acquire_lease_redis(&conn, &key, Instant::now())
591            .await
592            .is_err());
593
594        tokio::time::sleep(Duration::from_secs(1)).await;
595
596        lease1
597            .extend_redis(&conn, Duration::from_secs(6))
598            .await
599            .unwrap();
600
601        // Wait for original lease duration to expire
602        tokio::time::sleep(limit.duration + limit.duration).await;
603
604        // Cannot acquire because we have an extended lease
605        assert!(limit
606            .acquire_lease_redis(&conn, &key, Instant::now())
607            .await
608            .is_err());
609
610        // Wait for extension to pass
611        tokio::time::sleep(limit.duration + limit.duration).await;
612
613        let mut lease2 = limit
614            .acquire_lease_redis(&conn, &key, Instant::now())
615            .await
616            .unwrap();
617
618        lease1.release_redis(&conn).await;
619        lease2.release_redis(&conn).await;
620    }
621}
throttle/limit.rs

throttle/
limit.rs
