mirror of
https://github.com/natsukagami/youmubot.git
synced 2025-04-19 08:48:54 +00:00
prelude: Implement ratelimit
This commit is contained in:
parent
7e327fd131
commit
12948ae99f
4 changed files with 108 additions and 0 deletions
36
Cargo.lock
generated
36
Cargo.lock
generated
|
|
@ -283,6 +283,16 @@ dependencies = [
|
|||
"miniz_oxide 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "flume"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
dependencies = [
|
||||
"futures 0.3.5 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"rand 0.7.3 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"spinning_top 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
|
|
@ -606,6 +616,14 @@ name = "linked-hash-map"
|
|||
version = "0.5.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
|
||||
[[package]]
|
||||
name = "lock_api"
|
||||
version = "0.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
dependencies = [
|
||||
"scopeguard 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "log"
|
||||
version = "0.4.8"
|
||||
|
|
@ -1061,6 +1079,11 @@ dependencies = [
|
|||
"winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
|
||||
[[package]]
|
||||
name = "sct"
|
||||
version = "0.6.0"
|
||||
|
|
@ -1213,6 +1236,14 @@ name = "spin"
|
|||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
|
||||
[[package]]
|
||||
name = "spinning_top"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
dependencies = [
|
||||
"lock_api 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "static_assertions"
|
||||
version = "1.1.0"
|
||||
|
|
@ -1853,6 +1884,7 @@ dependencies = [
|
|||
"anyhow 1.0.32 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"async-trait 0.1.40 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"chrono 0.4.11 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"flume 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"futures-util 0.3.5 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"reqwest 0.10.4 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
"serenity 0.9.0-rc.1 (registry+https://github.com/rust-lang/crates.io-index)",
|
||||
|
|
@ -1898,6 +1930,7 @@ dependencies = [
|
|||
"checksum encoding_rs 0.8.23 (registry+https://github.com/rust-lang/crates.io-index)" = "e8ac63f94732332f44fe654443c46f6375d1939684c17b0afb6cb56b0456e171"
|
||||
"checksum env_logger 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)" = "44533bbbb3bb3c1fa17d9f2e4e38bbbaf8396ba82193c4cb1b6445d711445d36"
|
||||
"checksum flate2 1.0.14 (registry+https://github.com/rust-lang/crates.io-index)" = "2cfff41391129e0a856d6d822600b8d71179d46879e310417eb9c762eb178b42"
|
||||
"checksum flume 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)" = "8f7d4f1d85cb8cafb73c01d872e0124a96e603f5b7633ce1eac70015f066e946"
|
||||
"checksum fnv 1.0.7 (registry+https://github.com/rust-lang/crates.io-index)" = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||
"checksum foreign-types 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
|
||||
"checksum foreign-types-shared 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
|
||||
|
|
@ -1934,6 +1967,7 @@ dependencies = [
|
|||
"checksum lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
|
||||
"checksum libc 0.2.70 (registry+https://github.com/rust-lang/crates.io-index)" = "3baa92041a6fec78c687fa0cc2b3fae8884f743d672cf551bed1d6dac6988d0f"
|
||||
"checksum linked-hash-map 0.5.3 (registry+https://github.com/rust-lang/crates.io-index)" = "8dd5a6d5999d9907cda8ed67bbd137d3af8085216c2ac62de5be860bd41f304a"
|
||||
"checksum lock_api 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "28247cc5a5be2f05fbcd76dd0cf2c7d3b5400cb978a28042abcd4fa0b3f8261c"
|
||||
"checksum log 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)" = "14b6052be84e6b71ab17edffc2eeabf5c2c3ae1fdb464aae35ac50c67a44e1f7"
|
||||
"checksum matches 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08"
|
||||
"checksum memchr 2.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "3728d817d99e5ac407411fa471ff9800a778d88a24685968b36824eaf4bee400"
|
||||
|
|
@ -1984,6 +2018,7 @@ dependencies = [
|
|||
"checksum rustls-native-certs 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a75ffeb84a6bd9d014713119542ce415db3a3e4748f0bfce1e1416cd224a23a5"
|
||||
"checksum ryu 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)" = "ed3d612bc64430efeb3f7ee6ef26d590dce0c43249217bddc62112540c7941e1"
|
||||
"checksum schannel 0.1.19 (registry+https://github.com/rust-lang/crates.io-index)" = "8f05ba609c234e60bee0d547fe94a4c7e9da733d1c962cf6e59efa4cd9c8bc75"
|
||||
"checksum scopeguard 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
|
||||
"checksum sct 0.6.0 (registry+https://github.com/rust-lang/crates.io-index)" = "e3042af939fca8c3453b7af0f1c66e533a15a86169e39de2657310ade8f98d3c"
|
||||
"checksum security-framework 0.4.4 (registry+https://github.com/rust-lang/crates.io-index)" = "64808902d7d99f78eaddd2b4e2509713babc3dc3c85ad6f4c447680f3c01e535"
|
||||
"checksum security-framework-sys 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "17bf11d99252f512695eb468de5516e5cf75455521e69dfe343f3b74e4748405"
|
||||
|
|
@ -1999,6 +2034,7 @@ dependencies = [
|
|||
"checksum smallvec 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "c7cb5678e1615754284ec264d9bb5b4c27d2018577fd90ac0ceb578591ed5ee4"
|
||||
"checksum socket2 0.3.12 (registry+https://github.com/rust-lang/crates.io-index)" = "03088793f677dce356f3ccc2edb1b314ad191ab702a5de3faf49304f7e104918"
|
||||
"checksum spin 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
|
||||
"checksum spinning_top 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "7e529d73e80d64b5f2631f9035113347c578a1c9c7774b83a2b880788459ab36"
|
||||
"checksum static_assertions 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
|
||||
"checksum syn 1.0.23 (registry+https://github.com/rust-lang/crates.io-index)" = "95b5f192649e48a5302a13f2feb224df883b98933222369e4b3b0fe2a5447269"
|
||||
"checksum tempfile 3.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "7a6e24d9338a0a5be79593e2fa15a648add6138caa803e2d5bc782c371732ca9"
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ tokio = { version = "0.2", features = ["time"] }
|
|||
youmubot-db = { path = "../youmubot-db" }
|
||||
reqwest = "0.10"
|
||||
chrono = "0.4"
|
||||
flume = "0.9"
|
||||
|
||||
[dependencies.serenity]
|
||||
version = "0.9.0-rc.0"
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ pub mod args;
|
|||
pub mod hook;
|
||||
pub mod pagination;
|
||||
pub mod setup;
|
||||
pub mod ratelimit;
|
||||
|
||||
pub use announcer::{Announcer, AnnouncerHandler};
|
||||
pub use args::{Duration, UsernameArg};
|
||||
|
|
|
|||
70
youmubot-prelude/src/ratelimit.rs
Normal file
70
youmubot-prelude/src/ratelimit.rs
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
/// Provides a simple ratelimit lock (that only works in tokio)
|
||||
// use tokio::time::
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::Result;
|
||||
use flume::{bounded as channel, Receiver, Sender};
|
||||
use std::ops::Deref;
|
||||
|
||||
/// Holds the underlying `T` in a rate-limited way.
|
||||
pub struct Ratelimit<T> {
|
||||
inner: T,
|
||||
recv: Receiver<()>,
|
||||
send: Sender<()>,
|
||||
|
||||
wait_time: Duration,
|
||||
}
|
||||
|
||||
struct RatelimitGuard<'a, T> {
|
||||
inner: &'a T,
|
||||
send: &'a Sender<()>,
|
||||
wait_time: &'a Duration,
|
||||
}
|
||||
|
||||
impl<T> Ratelimit<T> {
|
||||
/// Create a new ratelimit with at most `count` uses in `wait_time`.
|
||||
pub fn new(inner: T, count: usize, wait_time: Duration) -> Self {
|
||||
let (send, recv) = channel(count);
|
||||
(0..count).for_each(|_| {
|
||||
send.send(()).ok();
|
||||
});
|
||||
Self {
|
||||
inner,
|
||||
send,
|
||||
recv,
|
||||
wait_time,
|
||||
}
|
||||
}
|
||||
|
||||
/// Borrow the inner `T`. You can only hol this reference `count` times in `wait_time`.
|
||||
/// The clock counts from the moment the ref is dropped.
|
||||
pub async fn borrow<'a>(&'a self) -> Result<impl Deref<Target = T> + 'a> {
|
||||
self.recv.recv_async().await?;
|
||||
eprintln!("lock accquired! {} left", self.recv.len());
|
||||
Ok(RatelimitGuard {
|
||||
inner: &self.inner,
|
||||
send: &self.send,
|
||||
wait_time: &self.wait_time,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T> Deref for RatelimitGuard<'a, T> {
|
||||
type Target = T;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T> Drop for RatelimitGuard<'a, T> {
|
||||
fn drop(&mut self) {
|
||||
let send = self.send.clone();
|
||||
let wait_time = self.wait_time.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::delay_for(wait_time).await;
|
||||
eprintln!("lock lifting!");
|
||||
send.send_async(()).await.ok();
|
||||
eprintln!("lock lifted!");
|
||||
});
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue