tocket 0.2.1

A implemetation of 'Token Bucket' algorithm
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

pub mod error;
pub mod whitelist;

mod codec;
mod message;
mod processing;

pub use error::DistributedStorageError;
pub use whitelist::WhitelistStrategy;

use crate::distributed::codec::Codec;
use crate::distributed::message::Message;
use crate::{InMemoryStorage, Storage, TokenBucketAlgorithm};

use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::sync::mpsc;
use tokio_util::udp::UdpFramed;
use tracing::Instrument;

type AcquireTx = mpsc::UnboundedSender<u32>;
type AcquireRx = mpsc::UnboundedReceiver<u32>;

/// A distributed storage that under the hood stores the state in the local `InMemoryStorage`
/// and sends messages to the rest of the distributed storages via UDP messages on each tokens acquiring,
/// according to the strategy used.
///
/// Useful when you have multiple application instances with shared state
/// but don't want to run additional storage (e.g. Redis).
///
/// # Available strategies:
/// - [`WhitelistStrategy`]
///
/// # Example
/// See usage examples in strategies above.
///
/// [`WhitelistStrategy`]: crate::distributed::whitelist::WhitelistStrategy
pub struct DistributedStorage {
    tx: AcquireTx,
    storage: Arc<InMemoryStorage>,
    listen_addr: SocketAddr,
}

impl DistributedStorage {
    /// Creates a distributed storage with the given strategy
    /// and starts a background task that will listen a UDP socket.
    ///
    /// # Errors
    ///
    /// Will return `Err` if failed to resolve listen address.
    pub async fn serve<A, S>(
        rps_limit: u32,
        listen_addr: A,
        strategy: S,
    ) -> Result<Self, DistributedStorageError>
    where
        A: ToSocketAddrs,
        S: Strategy + Send + 'static,
    {
        let listen_addr = listen_addr.to_socket_addrs()?.collect::<Vec<_>>();
        let socket = UdpSocket::bind(listen_addr.as_slice()).await?;
        let listen_addr = socket.local_addr()?;

        let storage = Arc::new(InMemoryStorage::new(rps_limit));
        let (tx, rx) = mpsc::unbounded_channel();
        tokio::spawn(
            processing::process(socket, strategy, Arc::clone(&storage), rx)
                .instrument(tracing::Span::current()),
        );

        Ok(Self {
            tx,
            storage,
            listen_addr,
        })
    }

    /// Get listen address.
    pub fn listen_addr(&self) -> SocketAddr {
        self.listen_addr
    }
}

impl Storage for DistributedStorage {
    type Error = DistributedStorageError;

    fn try_acquire(&self, alg: TokenBucketAlgorithm, permits: u32) -> Result<(), Self::Error> {
        self.storage.try_acquire(alg, permits)?;
        self.tx
            .send(permits)
            .expect("sending permits to background task failed, this is a bug");
        Ok(())
    }
}

#[async_trait::async_trait]
pub trait Strategy: private::Sealed {
    async fn on_acquire(
        &mut self,
        permits: u32,
        framed: &mut UdpFramed<Codec>,
    ) -> Result<(), DistributedStorageError>;

    async fn on_msg_recv(
        &mut self,
        msg: Message,
        source: SocketAddr,
        storage: &InMemoryStorage,
        framed: &mut UdpFramed<Codec>,
    ) -> Result<(), DistributedStorageError>;
}

mod private {
    use super::*;

    pub trait Sealed {}

    impl Sealed for WhitelistStrategy {}
}