prelude: Implement ratelimit

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};

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<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!");
+        });
+    }
+}