diff --git a/Cargo.lock b/Cargo.lock index 4271295..c0e15f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/youmubot-prelude/Cargo.toml b/youmubot-prelude/Cargo.toml index c4ea399..fbb99ba 100644 --- a/youmubot-prelude/Cargo.toml +++ b/youmubot-prelude/Cargo.toml @@ -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" diff --git a/youmubot-prelude/src/lib.rs b/youmubot-prelude/src/lib.rs index b979bae..2cb006d 100644 --- a/youmubot-prelude/src/lib.rs +++ b/youmubot-prelude/src/lib.rs @@ -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}; diff --git a/youmubot-prelude/src/ratelimit.rs b/youmubot-prelude/src/ratelimit.rs new file mode 100644 index 0000000..d117198 --- /dev/null +++ b/youmubot-prelude/src/ratelimit.rs @@ -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 { + inner: T, + recv: Receiver<()>, + send: Sender<()>, + + wait_time: Duration, +} + +struct RatelimitGuard<'a, T> { + inner: &'a T, + send: &'a Sender<()>, + wait_time: &'a Duration, +} + +impl Ratelimit { + /// 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 + '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!"); + }); + } +}