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

//! Rate-limiting for futures.
//!
//! This crate hooks the
//! [ratelimit_meter](https://crates.io/crates/ratelimit_meter) crate up
//! to futures v0.1 (the same version supported by Tokio right now).
//!
//! # Usage & mechanics of rate limiting with futures
//!
//! To use this crate's Future type, use the provided `Ratelimit::new`
//! function. It takes a direct rate limiter (an in-memory rate limiter
//! implementation), and returns a Future that can be chained to the
//! actual work that you mean to perform:
//!
//! ```rust
//! # extern crate ratelimit_meter;
//! # extern crate ratelimit_futures;
//! # extern crate futures;
//! use futures::prelude::*;
//! use futures::future::{self, FutureResult};
//! use ratelimit_meter::{DirectRateLimiter, LeakyBucket};
//! use ratelimit_futures::Ratelimit;
//! use std::num::NonZeroU32;
//!
//! let mut lim = DirectRateLimiter::<LeakyBucket>::per_second(NonZeroU32::new(1).unwrap());
//! {
//!     let mut lim = lim.clone();
//!     let ratelimit_future = Ratelimit::new(&mut lim);
//!     let future_of_3 = ratelimit_future.and_then(|_| {
//!         Ok(3)
//!     });
//! }
//! {
//!     let mut lim = lim.clone();
//!     let ratelimit_future = Ratelimit::new(&mut lim);
//!     let future_of_4 = ratelimit_future.and_then(|_| {
//!         Ok(4)
//!     });
//! }
//! // 1 second will pass before both futures resolve.
//! ```
//!
//! In this example, we're constructing futures that can each start work
//! only once the (shared) rate limiter says that it's ok to start.
//!
//! You can probably guess the mechanics of using these rate-limiting
//! futures:
//!
//! * Chain your work to them using `.and_then`.
//! * Construct and a single rate limiter for the work that needs to count
//!   against that rate limit. You can share them using their `Clone`
//!   trait.
//! * Rate-limiting futures will wait as long as it takes to arrive at a
//!   point where code is allowed to proceed. If the shared rate limiter
//!   already allowed another piece of code to proceed, the wait time will
//!   be extended.

use futures::{Async, Future, Poll};
use futures_timer::Delay;
use ratelimit_meter::{algorithms::Algorithm, DirectRateLimiter, NonConformance};
use std::io;

/// The rate-limiter as a future.
pub struct Ratelimit<'a, A: Algorithm>
where
    <A as Algorithm>::NegativeDecision: NonConformance,
{
    delay: Delay,
    limiter: &'a mut DirectRateLimiter<A>,
    first_time: bool,
}

impl<'a, A: Algorithm> Ratelimit<'a, A>
where
    <A as Algorithm>::NegativeDecision: NonConformance,
{
    /// Check if the rate-limiter would allow a request through.
    fn check(&mut self) -> Result<(), ()> {
        match self.limiter.check() {
            Ok(()) => Ok(()),
            Err(nc) => {
                let earliest = nc.earliest_possible();
                self.delay.reset_at(earliest);
                Err(())
            }
        }
    }

    /// Creates a new future that resolves successfully as soon as the
    /// rate limiter allows it.
    pub fn new(limiter: &'a mut DirectRateLimiter<A>) -> Self {
        Ratelimit {
            delay: Delay::new(Default::default()),
            first_time: true,
            limiter,
        }
    }
}

impl<'a, A: Algorithm> Future for Ratelimit<'a, A>
where
    <A as Algorithm>::NegativeDecision: NonConformance,
{
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        if self.first_time {
            // First time we run, let's check the rate-limiter and set
            // up a delay if we can't proceed:
            self.first_time = false;
            if self.check().is_ok() {
                return Ok(Async::Ready(()));
            }
        }
        match self.delay.poll() {
            // Timer says we should check the rate-limiter again, do
            // it and reset the delay otherwise.
            Ok(Async::Ready(_)) => match self.check() {
                Ok(_) => Ok(Async::Ready(())),
                Err(_) => {
                    self.delay.poll()?;
                    Ok(Async::NotReady)
                }
            },

            // timer isn't yet ready, let's wait:
            Ok(Async::NotReady) => Ok(Async::NotReady),

            // something went wrong:
            Err(e) => Err(e),
        }
    }
}