Local cache lock (#5193)

apps/labrinth/src/database/redis/util.rs

@@ -1,10 +1,16 @@
 use std::fmt::Debug;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use deadpool_redis::PoolError;
 use derive_more::{Deref, DerefMut};
 use redis::{FromRedisValue, RedisResult, ToRedisArgs};
+use tokio::sync::Notify;
+use tokio::time::{Duration, timeout};
 use tracing::{Instrument, info_span};
 
+use crate::database::models::DatabaseError;
+
 #[derive(Debug, Clone, Deref, DerefMut)]
 pub struct InstrumentedPool {
     inner: deadpool_redis::Pool,
@@ -91,3 +97,83 @@ impl InstrumentedCmd {
         self.inner.query_async(con).instrument(span).await
     }
 }
+
+pub fn cache() -> (CacheWriter, CacheSubscriber) {
+    let shared = Arc::new(Shared::new());
+    (
+        CacheWriter {
+            shared: shared.clone(),
+        },
+        CacheSubscriber { shared },
+    )
+}
+
+pub struct CacheWriter {
+    shared: Arc<Shared>,
+}
+
+impl CacheWriter {
+    pub fn write(&self) {
+        self.shared.make_ready();
+    }
+}
+
+#[derive(Clone)]
+pub struct CacheSubscriber {
+    shared: Arc<Shared>,
+}
+
+impl CacheSubscriber {
+    pub async fn wait_timeout(
+        self,
+        duration: Duration,
+    ) -> Result<(), DatabaseError> {
+        timeout(duration, self.shared.wait()).await.map_err(|_| {
+            DatabaseError::LocalCacheTimeout {
+                released: 0,
+                total: 1,
+            }
+        })
+    }
+}
+
+struct Shared {
+    ready: AtomicBool,
+    // With this implementation's intrusive linked lists, the waiters are stored inline in the future
+    // so there's no heap allocation per waiter.
+    wakers: Notify,
+}
+
+impl Shared {
+    fn new() -> Self {
+        Self {
+            ready: AtomicBool::new(false),
+            wakers: Notify::new(),
+        }
+    }
+
+    fn make_ready(&self) {
+        self.ready.store(true, Ordering::Release);
+        self.wakers.notify_waiters();
+    }
+
+    async fn wait(&self) {
+        let ready = self.ready.load(Ordering::Acquire);
+
+        if ready {
+            return;
+        }
+
+        let notification = self.wakers.notified();
+        // Don't need to call `enable` as we use notify_waiters
+
+        // Prevent race where the writer set the ready bit and notified waiters between the load and registering the waiter
+        let ready = self.ready.load(Ordering::SeqCst);
+
+        if ready {
+            return;
+        }
+
+        notification.await;
+    }
+}