traq_ws_bot/
utils.rs

1use std::sync::Arc;
2
3use tokio::sync::Semaphore;
4
5use crate::events::common::Message;
6#[cfg(feature = "openapi")]
7use crate::openapi;
8
9/// メッセージに `target_user_id` に対するメンションが含まれているかどうかを返す
10pub fn is_mentioned_message(message: &Message, target_user_id: &str) -> bool {
11    message
12        .embedded
13        .iter()
14        .any(|emb| emb.id.as_str() == target_user_id)
15}
16
17/// BOT の access token から簡単な API Client を作成する
18pub fn create_client(bot_access_token: impl Into<String>) -> reqwest::Client {
19    let mut headers = reqwest::header::HeaderMap::new();
20    let authorization_token = format!("Bearer {}", bot_access_token.into());
21    headers.insert(
22        reqwest::header::AUTHORIZATION,
23        reqwest::header::HeaderValue::from_str(&authorization_token).unwrap(),
24    );
25    reqwest::Client::builder()
26        .default_headers(headers)
27        .build()
28        .unwrap()
29}
30
31/// BOT の access token から openapi の configuration を作成する
32#[cfg(feature = "openapi")]
33pub fn create_configuration(bot_access_token: impl Into<String>) -> openapi::apis::configuration::Configuration {
34    openapi::apis::configuration::Configuration {
35        bearer_access_token: Some(bot_access_token.into()),
36        ..Default::default()
37    }
38}
39
40/// interval 間に最大 max_count 回しか実行されないようにすることができる struct
41///
42/// **Default:** `interval` は 5 秒, `max_count` は 5 回
43///
44/// # Example
45/// ```
46/// use std::time::Duration;
47/// use traq_ws_bot::utils::RateLimiter;
48///
49/// let mut limiter = RateLimiter::new(5, Duration::from_secs(5)); // RateLimiter::default() と同じ
50/// loop {
51///    if limiter.try_acquire() {
52///        // 5秒間に5回しか実行されない
53///        println!("Hello");
54///    }
55/// }
56/// ```
57pub struct RateLimiter {
58    interval: std::time::Duration,
59    semaphore: Arc<Semaphore>,
60}
61impl RateLimiter {
62    /// interval 間に最大 max_count 回しか実行されないようにすることができる struct を作成する
63    pub fn new(max_count: usize, interval: std::time::Duration) -> Self {
64        Self {
65            interval,
66            semaphore: Arc::new(Semaphore::new(max_count)),
67        }
68    }
69
70    /// lock を取得するまで待機する
71    ///
72    /// # Example
73    /// ```
74    /// use std::time::Duration;
75    /// use traq_ws_bot::utils::RateLimiter;
76    ///
77    /// let mut limiter = RateLimiter::new(5, Duration::from_secs(5)); // RateLimiter::default() と同じ
78    /// loop {
79    ///     await limiter.acquire();
80    ///     // 5秒間に5回しか実行されない
81    ///     println!("Hello");
82    /// }
83    pub async fn acquire(&self) {
84        let semaphore = self.semaphore.clone();
85        let permit = semaphore.acquire_owned().await.unwrap();
86        let interval = self.interval;
87
88        tokio::spawn( async move {
89            tokio::time::sleep(interval).await;
90            drop(permit);
91        });
92    }
93
94    /// lock を取得できたら true を返す
95    ///
96    /// # Example
97    /// ```
98    /// use std::time::Duration;
99    /// use traq_ws_bot::utils::RateLimiter;
100    ///
101    /// let mut limiter = RateLimiter::new(5, Duration::from_secs(5)); // RateLimiter::default() と同じ
102    /// loop {
103    ///     if limiter.try_acquire() {
104    ///        // 5秒間に5回しか実行されない
105    ///       println!("Hello");
106    ///    }
107    /// }
108    pub fn try_acquire(&self) -> bool {
109        let semaphore = self.semaphore.clone();
110        match semaphore.try_acquire_owned() {
111            Ok(permit) => {
112                let interval = self.interval;
113                tokio::spawn( async move {
114                    tokio::time::sleep(interval).await;
115                    drop(permit);
116                });
117                true
118            }
119            Err(_) => false,
120        }
121    }
122}
123impl Default for RateLimiter {
124    fn default() -> Self {
125        Self::new(5, std::time::Duration::from_secs(5))
126    }
127}
traq_ws_bot/utils.rs

traq_ws_bot/
utils.rs
