haste_repository/pg/
rate_limit.rs

1use std::{pin::Pin, sync::LazyLock};
2
3use crate::pg::{
4    PGConnection,
5    transaction::{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| match &e {
45                // `check_rate_limit` raises once the points pass the maximum.
46                // That is the limit doing its job, not a failure to check it,
47                // so surface it as `Exceeded` rather than a generic error.
48                sqlx::Error::Database(db_error)
49                    if db_error.message().starts_with("Rate limit exceeded") =>
50                {
51                    RateLimitError::Exceeded
52                }
53                _ => RateLimitError::Error(e.to_string()),
54            })?;
55
56    Ok(result.0)
57}
58
59async fn check_rate_limit_remote(
60    pg: PGConnection,
61    rate_key: &str,
62    max: i32,
63    points: i32,
64    window_in_seconds: i32,
65) -> Result<i32, haste_rate_limit::RateLimitError> {
66    match &pg {
67        PGConnection::Pool(_pool, _) => {
68            let tx = create_transaction(&pg, false)
69                .await
70                .map_err(|e| RateLimitError::Error(e.to_string()))?;
71            let res = {
72                let mut conn = tx.lock().await;
73                check_rate_limit_remote_with_executor(
74                    &mut **conn,
75                    rate_key,
76                    max,
77                    points,
78                    window_in_seconds,
79                )
80                .await?
81            };
82            commit_transaction(tx)
83                .await
84                .map_err(|e| RateLimitError::Error(e.to_string()))?;
85            Ok(res)
86        }
87        PGConnection::Transaction(tx, _, _) => {
88            let mut tx = tx.lock().await;
89            check_rate_limit_remote_with_executor(
90                &mut **tx,
91                rate_key,
92                max,
93                points,
94                window_in_seconds,
95            )
96            .await
97        }
98    }
99}
100
101async fn check_rate_limit(
102    connection: PGConnection,
103    rate_key: &str,
104    max: i32,
105    points: i32,
106    window_in_seconds: i32,
107) -> Result<i32, haste_rate_limit::RateLimitError> {
108    if let Some(current) = MEMORY.get(rate_key).await {
109        let cloned_key = rate_key.to_string();
110        let connection_clone = connection.clone();
111
112        tokio::spawn(async move {
113            let result = check_rate_limit_remote(
114                connection_clone,
115                &cloned_key,
116                max,
117                points,
118                window_in_seconds,
119            )
120            .await;
121
122            if let Ok(points) = result {
123                MEMORY
124                    .insert(cloned_key, RateLimitState::Count(points))
125                    .await;
126            } else if let Err(e) = result {
127                match e {
128                    RateLimitError::Exceeded => {
129                        MEMORY.insert(cloned_key, RateLimitState::Max).await;
130                    }
131                    RateLimitError::Error(e) => {
132                        println!("Error checking rate limit: {e:?}");
133                    }
134                }
135            }
136        });
137
138        match current {
139            RateLimitState::Count(current) => {
140                let current_score = current + points;
141
142                if current_score > max {
143                    Err(RateLimitError::Exceeded)
144                } else {
145                    MEMORY
146                        .insert(rate_key.to_string(), RateLimitState::Count(current_score))
147                        .await;
148                    Ok(current_score)
149                }
150            }
151            RateLimitState::Max => Err(RateLimitError::Exceeded),
152        }
153    } else {
154        let result =
155            check_rate_limit_remote(connection, rate_key, max, points, window_in_seconds).await?;
156
157        MEMORY
158            .insert(rate_key.to_string(), RateLimitState::Count(result))
159            .await;
160
161        Ok(result)
162    }
163}
164
165impl RateLimit for PGConnection {
166    /// Returns the current points after the operation.
167    /// Note use of box and pin so can satisfy dynamic dispatch requirements.
168    fn check<'a>(
169        &'a self,
170        rate_key: &'a str,
171        max: i32,
172        points: i32,
173        window_in_seconds: i32,
174    ) -> Pin<Box<dyn Future<Output = Result<i32, haste_rate_limit::RateLimitError>> + Send + 'a>>
175    {
176        let connection = self.clone();
177        Box::pin(async move {
178            let res =
179                check_rate_limit(connection, rate_key, max, points, window_in_seconds).await?;
180            Ok(res)
181        })
182    }
183
184    /// Goes straight to Postgres, skipping [`MEMORY`]. `check_rate_limit` takes
185    /// a per-key advisory transaction lock, so concurrent callers are counted
186    /// one at a time with no lost increments. The in-memory cache is what makes
187    /// [`RateLimit::check`] fast but imprecise under a burst; this trades that
188    /// for an exact count on every call.
189    fn check_uncached<'a>(
190        &'a self,
191        rate_key: &'a str,
192        max: i32,
193        points: i32,
194        window_in_seconds: i32,
195    ) -> Pin<Box<dyn Future<Output = Result<i32, haste_rate_limit::RateLimitError>> + Send + 'a>>
196    {
197        let connection = self.clone();
198        Box::pin(async move {
199            check_rate_limit_remote(connection, rate_key, max, points, window_in_seconds).await
200        })
201    }
202}
haste_repository/pg/rate_limit.rs

haste_repository/pg/
rate_limit.rs
