Another locks bug fix

Cargo.toml

@@ -3,7 +3,7 @@ members = ["codegen", "examples", "performance_measurement", "performance_measur
 
 [package]
 name = "worktable"
-version = "0.6.11"
+version = "0.6.12"
 edition = "2024"
 authors = ["Handy-caT"]
 license = "MIT"
@@ -16,13 +16,14 @@ perf_measurements = ["dep:performance_measurement", "dep:performance_measurement
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+worktable_codegen = { path = "codegen", version = "0.6.12" }
+
 eyre = "0.6.12"
 derive_more = { version = "1.0.0", features = ["from", "error", "display", "into"] }
 tokio = { version = "1", features = ["full"] }
 tracing = "0.1"
 rkyv = { version = "0.8.9", features = ["uuid-1"] }
 lockfree = { version = "0.5.1" }
-worktable_codegen = { path = "codegen", version = "0.6.11" }
 fastrand = "2.3.0"
 futures = "0.3.30"
 uuid = { version = "1.10.0", features = ["v4", "v7"] }

codegen/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "worktable_codegen"
-version = "0.6.11"
+version = "0.6.12"
 edition = "2024"
 license = "MIT"
 description = "WorkTable codegeneration crate"

codegen/src/worktable/generator/queries/in_place.rs

@@ -117,15 +117,15 @@ impl Generator {
                 by: #by_type,
             ) -> eyre::Result<()> {
                 let pk: #pk_type = by.into();
+                let lock = {
+                    #custom_lock
+                };
                 let link = self
                     .0
                     .pk_map
                     .get(&pk)
                     .map(|v| v.get().value)
                     .ok_or(WorkTableError::NotFound)?;
-                let lock = {
-                    #custom_lock
-                };
                 unsafe {
                     self.0
                         .data

codegen/src/worktable/generator/queries/locks.rs

@@ -132,7 +132,7 @@ impl Generator {
                 #[allow(clippy::mutable_key_type)]
                 let (locks, op_lock) = lock_guard.lock(lock_id);
                 drop(lock_guard);
-                futures::future::join_all(locks.iter().map(|l| l.as_ref()).collect::<Vec<_>>()).await;
+                futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
 
                 op_lock
             } else {
@@ -147,7 +147,7 @@ impl Generator {
                     drop(old_lock_guard);
                     drop(guard);
 
-                    futures::future::join_all(locks.iter().map(|l| l.as_ref()).collect::<Vec<_>>()).await;
+                    futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
                 }
 
                 op_lock
@@ -166,8 +166,7 @@ impl Generator {
                 #[allow(clippy::mutable_key_type)]
                 let (locks, op_lock) = lock_guard.#ident(lock_id);
                 drop(lock_guard);
-                futures::future::join_all(locks.iter().map(|l| l.as_ref()).collect::<Vec<_>>()).await;
-
+                futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
                 op_lock
             } else {
                 let mut lock = #lock_ident::new();
@@ -182,7 +181,7 @@ impl Generator {
                     drop(old_lock_guard);
                     drop(guard);
 
-                    futures::future::join_all(locks.iter().map(|l| l.as_ref()).collect::<Vec<_>>()).await;
+                    futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
                 }
 
                 op_lock

codegen/src/worktable/generator/queries/update.rs

@@ -518,6 +518,7 @@ impl Generator {
                     locks.insert(pk, op_lock);
                 }
 
+                let links: Vec<_> = self.0.indexes.#index.get(#by).map(|(_, l)| *l).collect();
                 let mut pk_to_unlock: std::collections::HashMap<_, std::sync::Arc<Lock>> = std::collections::HashMap::new();
                 let op_id = OperationId::Multi(uuid::Uuid::now_v7());
                 for link in links.into_iter() {
@@ -612,6 +613,11 @@ impl Generator {
                     #custom_lock
                 };
 
+                let link = self.0.indexes.#index
+                    .get(#by)
+                    .map(|kv| kv.get().value)
+                    .ok_or(WorkTableError::NotFound)?;
+
                 let op_id = OperationId::Single(uuid::Uuid::now_v7());
                 #size_check
                 #diff_process

src/lock/map.rs

@@ -49,14 +49,12 @@ where
         LockType: RowLock,
     {
         let mut set = self.map.write();
-        let remove = if let Some(lock) = set.get(key) {
-            !lock.read().await.is_locked()
-        } else {
-            false
-        };
-
-        if remove {
-            set.remove(key);
+        if let Some(lock) = set.get(key).cloned() {
+            if let Ok(guard) = lock.try_read() {
+                if !guard.is_locked() {
+                    set.remove(key);
+                }
+            }
         }
     }
 

src/lock/mod.rs

@@ -5,19 +5,20 @@ use std::future::Future;
 use std::hash::{Hash, Hasher};
 use std::pin::Pin;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use derive_more::From;
 use futures::task::AtomicWaker;
-
 pub use map::LockMap;
+use parking_lot::Mutex;
 pub use row_lock::RowLock;
 
 #[derive(Debug)]
 pub struct Lock {
     id: u16,
-    locked: AtomicBool,
-    waker: AtomicWaker,
+    locked: Arc<AtomicBool>,
+    wakers: Mutex<Vec<Arc<AtomicWaker>>>,
 }
 
 impl PartialEq for Lock {
@@ -38,32 +39,59 @@ impl Lock {
     pub fn new(id: u16) -> Self {
         Self {
             id,
-            locked: AtomicBool::from(true),
-            waker: AtomicWaker::new(),
+            locked: Arc::new(AtomicBool::from(true)),
+            wakers: Mutex::new(vec![]),
         }
     }
 
+    pub fn id(&self) -> u16 {
+        self.id
+    }
+
     pub fn unlock(&self) {
-        self.locked.store(false, Ordering::Release);
-        self.waker.wake()
+        self.locked.store(false, Ordering::Relaxed);
+        let guard = self.wakers.lock();
+        for w in guard.iter() {
+            w.wake()
+        }
     }
 
     pub fn lock(&self) {
-        self.locked.store(true, Ordering::Release);
-        self.waker.wake()
+        self.locked.store(true, Ordering::Relaxed);
     }
 
     pub fn is_locked(&self) -> bool {
-        self.locked.load(Ordering::Acquire)
+        self.locked.load(Ordering::Relaxed)
+    }
+
+    pub fn wait(&self) -> LockWait {
+        let mut guard = self.wakers.lock();
+        let waker = Arc::new(AtomicWaker::new());
+        guard.push(waker.clone());
+        LockWait {
+            locked: self.locked.clone(),
+            waker,
+        }
     }
 }
 
-impl Future for &Lock {
+#[derive(Debug)]
+pub struct LockWait {
+    locked: Arc<AtomicBool>,
+    waker: Arc<AtomicWaker>,
+}
+
+impl Future for LockWait {
     type Output = ();
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        self.as_ref().waker.register(cx.waker());
-        if self.locked.load(Ordering::Acquire) {
+        if !self.locked.load(Ordering::Relaxed) {
+            return Poll::Ready(());
+        }
+
+        self.waker.register(cx.waker());
+
+        if self.locked.load(Ordering::Relaxed) {
             Poll::Pending
         } else {
             Poll::Ready(())