korrosync 0.4.0

A KOReader Sync Server
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

use std::time::Duration;

use governor::{clock::QuantaInstant, middleware::NoOpMiddleware};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tower_governor::{
    GovernorLayer, governor::GovernorConfigBuilder, key_extractor::PeerIpKeyExtractor,
};

use crate::config::RateLimit;

pub fn rate_limiter_layer<RespBody>(
    shutdown_token: CancellationToken,
    config: &RateLimit,
) -> (
    GovernorLayer<PeerIpKeyExtractor, NoOpMiddleware<QuantaInstant>, RespBody>,
    JoinHandle<()>,
) {
    let governor_conf = GovernorConfigBuilder::default()
        .per_second(config.per_second)
        .burst_size(config.burst_size)
        .finish()
        .unwrap();

    let governor_limiter = governor_conf.limiter().clone();
    let interval = Duration::from_secs(60);

    let cleanup_task = tokio::spawn(async move {
        // separate background task to clean up
        loop {
            tokio::select! {
                _ = shutdown_token.cancelled() => {
                    tracing::info!("Rate limiter cleanup task shutting down");
                    break;
                }
                _ = tokio::time::sleep(interval) => {
                    tracing::debug!("rate limiting storage size: {}", governor_limiter.len());
                    governor_limiter.retain_recent();
                }
            }
        }
    });

    (GovernorLayer::new(governor_conf), cleanup_task)
}