From 6bdd8852efe4dee48b479617a5921279bf759f78 Mon Sep 17 00:00:00 2001
From: meh <no-reply@pathscale.com>
Date: Mon, 3 Aug 2026 01:23:43 +0700
Subject: [PATCH] test: pin multi-row primary-key lock order

---
 tests/worktable/lock_order.rs | 124 ++++++++++++++++++++++++++++++++++
 tests/worktable/mod.rs        |   1 +
 2 files changed, 125 insertions(+)
 create mode 100644 tests/worktable/lock_order.rs

diff --git a/tests/worktable/lock_order.rs b/tests/worktable/lock_order.rs
new file mode 100644
index 0000000..d1c342c
--- /dev/null
+++ b/tests/worktable/lock_order.rs
@@ -0,0 +1,124 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use worktable::prelude::*;
+use worktable::worktable;
+
+worktable!(
+    name: LockOrder,
+    columns: {
+        id: u64 primary_key autoincrement,
+        group_a: u64,
+        group_b: u64,
+        value: u64,
+    },
+    indexes: {
+        group_a_idx: group_a,
+        group_b_idx: group_b,
+    },
+    queries: {
+        update: {
+            ValueByGroupA(value) by group_a,
+            ValueByGroupB(value) by group_b,
+        }
+    }
+);
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn multi_row_update_locks_in_primary_key_order_not_index_order() {
+    const ROWS: u64 = 128;
+    let (table, use_group_a, blocker_id) = loop {
+        let table = Arc::new(LockOrderWorkTable::default());
+        for id in 0..ROWS {
+            table
+                .insert(LockOrderRow {
+                    id,
+                    group_a: 1,
+                    group_b: 1,
+                    value: 0,
+                })
+                .unwrap();
+        }
+
+        let order_a: Vec<_> = table
+            .0
+            .indexes
+            .group_a_idx
+            .get(&1)
+            .map(|(_, link)| table.0.data.select_non_ghosted(link.0).unwrap().id)
+            .collect();
+        let order_b: Vec<_> = table
+            .0
+            .indexes
+            .group_b_idx
+            .get(&1)
+            .map(|(_, link)| table.0.data.select_non_ghosted(link.0).unwrap().id)
+            .collect();
+        if order_a[0] != 0 {
+            break (table, true, order_a[0]);
+        }
+        if order_b[0] != 0 {
+            break (table, false, order_b[0]);
+        }
+    };
+
+    // Hold the first row exposed by the chosen secondary index. An
+    // index-ordered update blocks here before touching primary key zero. A
+    // primary-key-ordered update registers its lock for key zero first.
+    let blocker_pk = LockOrderPrimaryKey(blocker_id);
+    let blocker = Arc::new(Lock::new(u16::MAX));
+    let mut blocker_state = LockOrderLock::new();
+    blocker_state.value_lock = Some(blocker.clone());
+    table
+        .0
+        .lock_manager
+        .insert(blocker_pk.clone(), Arc::new(tokio::sync::RwLock::new(blocker_state)));
+
+    let update_table = table.clone();
+    let update = tokio::spawn(async move {
+        if use_group_a {
+            update_table
+                .update_value_by_group_a(ValueByGroupAQuery { value: 1 }, 1)
+                .await
+        } else {
+            update_table
+                .update_value_by_group_b(ValueByGroupBQuery { value: 1 }, 1)
+                .await
+        }
+    });
+
+    let blocker_state = table.0.lock_manager.get(&blocker_pk).unwrap();
+    tokio::time::timeout(Duration::from_secs(2), async {
+        loop {
+            let installed_id = blocker_state
+                .read()
+                .await
+                .value_lock
+                .as_ref()
+                .expect("the update locks value")
+                .id();
+            if installed_id != blocker.id() {
+                break;
+            }
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("the update did not reach the deliberately blocked row");
+
+    assert!(
+        table.0.lock_manager.get(&LockOrderPrimaryKey(0)).is_some(),
+        "the update followed secondary-index order instead of locking primary key zero first"
+    );
+
+    blocker.unlock();
+    tokio::time::timeout(Duration::from_secs(2), update)
+        .await
+        .expect("the update did not resume after the predecessor lock was released")
+        .unwrap()
+        .unwrap();
+
+    for id in 0..ROWS {
+        assert_eq!(table.select(id).unwrap().value, 1);
+    }
+}
diff --git a/tests/worktable/mod.rs b/tests/worktable/mod.rs
index b4db304..ba6d2c2 100644
--- a/tests/worktable/mod.rs
+++ b/tests/worktable/mod.rs
@@ -8,6 +8,7 @@ mod delete;
 mod float;
 mod in_place;
 mod index;
+mod lock_order;
 mod nid;
 mod option;
 mod tuple_primary_key;