fix txn lock bugs and add more assertions

Author: MarinPostma

libsql-wal/src/shared_wal.rs

@@ -104,41 +104,38 @@ impl<IO: Io> SharedWal<IO> {
             match tx {
                 Transaction::Write(_) => unreachable!("already in a write transaction"),
                 Transaction::Read(read_tx) => {
-                    {
-                        let mut reserved = self.wal_lock.reserved.lock();
-                        match *reserved {
-                            // we have already reserved the slot, go ahead and try to acquire
-                            Some(id) if id == read_tx.conn_id => {
-                                tracing::trace!("taking reserved slot");
-                                reserved.take();
-                                let lock = self.wal_lock.tx_id.lock_blocking();
+                    let mut reserved = self.wal_lock.reserved.lock();
+                    match *reserved {
+                        // we have already reserved the slot, go ahead and try to acquire
+                        Some(id) if id == read_tx.conn_id => {
+                            tracing::trace!("taking reserved slot");
+                            reserved.take();
+                            let lock = self.wal_lock.tx_id.lock_blocking();
+                            assert!(lock.is_none());
+                            let write_tx = self.acquire_write(read_tx, lock, reserved)?;
+                            *tx = Transaction::Write(write_tx);
+                            return Ok(());
+                        }
+                        None => {
+                            let lock = self.wal_lock.tx_id.lock_blocking();
+                            if lock.is_none() && self.wal_lock.waiters.is_empty() {
                                 let write_tx = self.acquire_write(read_tx, lock, reserved)?;
                                 *tx = Transaction::Write(write_tx);
                                 return Ok(());
                             }
-                            _ => (),
                         }
+                        _ => (),
                     }
 
-                    let lock = self.wal_lock.tx_id.lock_blocking();
-                    match *lock {
-                        None if self.wal_lock.waiters.is_empty() => {
-                            let write_tx =
-                                self.acquire_write(read_tx, lock, self.wal_lock.reserved.lock())?;
-                            *tx = Transaction::Write(write_tx);
-                            return Ok(());
-                        }
-                        Some(_) | None => {
-                            tracing::trace!(
-                                "txn currently held by another connection, registering to wait queue"
-                            );
-                            let parker = crossbeam::sync::Parker::new();
-                            let unparker = parker.unparker().clone();
-                            self.wal_lock.waiters.push((unparker, read_tx.conn_id));
-                            drop(lock);
-                            parker.park();
-                        }
-                    }
+                    tracing::trace!(
+                        "txn currently held by another connection, registering to wait queue"
+                    );
+
+                    let parker = crossbeam::sync::Parker::new();
+                    let unparker = parker.unparker().clone();
+                    self.wal_lock.waiters.push((unparker, read_tx.conn_id));
+                    drop(reserved);
+                    parker.park();
                 }
             }
         }
@@ -150,6 +147,8 @@ impl<IO: Io> SharedWal<IO> {
         mut tx_id_lock: async_lock::MutexGuard<Option<u64>>,
         mut reserved: MutexGuard<Option<u64>>,
     ) -> Result<WriteTransaction<IO::File>> {
+        assert!(reserved.is_none() || *reserved == Some(read_tx.conn_id));
+        assert!(tx_id_lock.is_none());
         // we read two fields in the header. There is no risk that a transaction commit in
         // between the two reads because this would require that:
         // 1) there would be a running txn

libsql-wal/src/transaction.rs

@@ -283,6 +283,8 @@ impl<F> WriteTransaction<F> {
         let Self {
             wal_lock, read_tx, ..
         } = self;
+        // always acquire lock in this order: reserved, then tx_id
+        let mut reserved = wal_lock.reserved.lock();
         let mut lock = wal_lock.tx_id.lock_blocking();
         match *lock {
             Some(lock_id) if lock_id == read_tx.id => {
@@ -291,7 +293,7 @@ impl<F> WriteTransaction<F> {
             _ => (),
         }
 
-        if let Some(id) = *wal_lock.reserved.lock() {
+        if let Some(id) = *reserved {
             tracing::trace!("tx already reserved by {id}");
             return read_tx;
         }
@@ -304,7 +306,7 @@ impl<F> WriteTransaction<F> {
                 }
                 crossbeam::deque::Steal::Success((unparker, id)) => {
                     tracing::trace!("waking up {id}");
-                    wal_lock.reserved.lock().replace(id);
+                    reserved.replace(id);
                     unparker.unpark();
                     break;
                 }