haste_repository/pg/
rate_limit.rs

1use std::{pin::Pin, sync::LazyLock};
2
3use crate::pg::{
4    PGConnection,
5    utilities::{commit_transaction, create_transaction},
6};
7use haste_rate_limit::{RateLimit, RateLimitError};
8use moka::future::{Cache, CacheBuilder};
9use sqlx::PgExecutor;
10
11#[derive(Clone, Copy)]
12enum RateLimitState {
13    Count(i32),
14    Max,
15}
16
17static MEMORY: LazyLock<Cache<String, RateLimitState>> = LazyLock::new(
18    // Cache entries live for 30 seconds, after which they will be automatically evicted.
19    || {
20        CacheBuilder::new(10_000)
21            .time_to_idle(std::time::Duration::from_secs(30))
22            .build()
23    },
24);
25
26async fn check_rate_limit_remote_with_executor<'a, 'e, E>(
27    executor: E,
28    rate_key: &'a str,
29    max: i32,
30    points: i32,
31    window_in_seconds: i32,
32) -> Result<i32, haste_rate_limit::RateLimitError>
33where
34    E: PgExecutor<'e>,
35{
36    let result =
37        sqlx::query_as::<_, (i32,)>("SELECT check_rate_limit($1, $2, $3, $4) as current_limit")
38            .bind(rate_key)
39            .bind(max)
40            .bind(points)
41            .bind(window_in_seconds)
42            .fetch_one(executor)
43            .await
44            .map_err(|e| RateLimitError::Error(e.to_string()))?;
45
46    Ok(result.0)
47}
48
49async fn check_rate_limit_remote(
50    pg: PGConnection,
51    rate_key: &str,
52    max: i32,
53    points: i32,
54    window_in_seconds: i32,
55) -> Result<i32, haste_rate_limit::RateLimitError> {
56    match &pg {
57        PGConnection::Pool(_pool, _) => {
58            let tx = create_transaction(&pg, true)
59                .await
60                .map_err(|e| RateLimitError::Error(e.to_string()))?;
61            let res = {
62                let mut conn = tx.lock().await;
63                check_rate_limit_remote_with_executor(
64                    &mut **conn,
65                    rate_key,
66                    max,
67                    points,
68                    window_in_seconds,
69                )
70                .await?
71            };
72            commit_transaction(tx)
73                .await
74                .map_err(|e| RateLimitError::Error(e.to_string()))?;
75            Ok(res)
76        }
77        PGConnection::Transaction(tx, _) => {
78            let mut tx = tx.lock().await;
79            check_rate_limit_remote_with_executor(
80                &mut **tx,
81                rate_key,
82                max,
83                points,
84                window_in_seconds,
85            )
86            .await
87        }
88    }
89}
90
91async fn check_rate_limit(
92    connection: PGConnection,
93    rate_key: &str,
94    max: i32,
95    points: i32,
96    window_in_seconds: i32,
97) -> Result<i32, haste_rate_limit::RateLimitError> {
98    if let Some(current) = MEMORY.get(rate_key).await {
99        let cloned_key = rate_key.to_string();
100        let connection_clone = connection.clone();
101
102        tokio::spawn(async move {
103            let result = check_rate_limit_remote(
104                connection_clone,
105                &cloned_key,
106                max,
107                points,
108                window_in_seconds,
109            )
110            .await;
111
112            if let Ok(points) = result {
113                MEMORY
114                    .insert(cloned_key, RateLimitState::Count(points))
115                    .await;
116            } else if let Err(e) = result {
117                match e {
118                    RateLimitError::Exceeded => {
119                        MEMORY.insert(cloned_key, RateLimitState::Max).await;
120                    }
121                    RateLimitError::Error(e) => {
122                        println!("Error checking rate limit: {e:?}");
123                    }
124                }
125            }
126        });
127
128        match current {
129            RateLimitState::Count(current) => {
130                let current_score = current + points;
131
132                if current_score > max {
133                    Err(RateLimitError::Exceeded)
134                } else {
135                    MEMORY
136                        .insert(rate_key.to_string(), RateLimitState::Count(current_score))
137                        .await;
138                    Ok(current_score)
139                }
140            }
141            RateLimitState::Max => Err(RateLimitError::Exceeded),
142        }
143    } else {
144        let result =
145            check_rate_limit_remote(connection, rate_key, max, points, window_in_seconds).await?;
146
147        MEMORY
148            .insert(rate_key.to_string(), RateLimitState::Count(result))
149            .await;
150
151        Ok(result)
152    }
153}
154
155impl RateLimit for PGConnection {
156    /// Returns the current points after the operation.
157    /// Note use of box and pin so can satisfy dynamic dispatch requirements.
158    fn check<'a>(
159        &'a self,
160        rate_key: &'a str,
161        max: i32,
162        points: i32,
163        window_in_seconds: i32,
164    ) -> Pin<Box<dyn Future<Output = Result<i32, haste_rate_limit::RateLimitError>> + Send + 'a>>
165    {
166        let connection = self.clone();
167        Box::pin(async move {
168            let res =
169                check_rate_limit(connection, rate_key, max, points, window_in_seconds).await?;
170            Ok(res)
171        })
172    }
173}
haste_repository/pg/rate_limit.rs

haste_repository/pg/
rate_limit.rs
