[PATCH 04/22] rust/bql: Add BqlGuard to provide BQL context

[Zhao Liu <zhao1.liu@intel.com>]

There's some case requiring BQL context, for example, accessing
BqlCell/BqlRefCell or operation on IRQ in lockless IO.

Note though BqlCell/BqlRefCell ensures BQL won't be released during
reference exists, they don't create BQL context if BQL is not locked,
instead, they just assert and panic when BQL is not locked.

Therefore, it's necessary to provide a way to create BQL context by hand
at Rust side. BqlGuard is suitable for this need.

Signed-off-by: Zhao Liu <zhao1.liu@intel.com>
---
 include/qemu/main-loop.h |  22 ++++++++-
 rust/bql/src/lib.rs      | 101 ++++++++++++++++++++++++++++++++++++++-
 stubs/iothread-lock.c    |  11 +++++
 system/cpus.c            |  10 ++++
 4 files changed, 142 insertions(+), 2 deletions(-)

diff --git a/include/qemu/main-loop.h b/include/qemu/main-loop.h
index 0d55c636b21a..7dd7c8211f02 100644
--- a/include/qemu/main-loop.h
+++ b/include/qemu/main-loop.h
@@ -250,7 +250,7 @@ AioContext *iohandler_get_aio_context(void);
 /**
  * rust_bql_mock_lock:
  *
- * Called from Rust doctests to make bql_lock() return true.
+ * Called from Rust doctests to make bql_locked() return true.
  * Do not touch.
  */
 void rust_bql_mock_lock(void);
@@ -368,6 +368,16 @@ bool qemu_in_main_thread(void);
 #define bql_lock() bql_lock_impl(__FILE__, __LINE__)
 void bql_lock_impl(const char *file, int line);
 
+/**
+ * @rust_bql_lock: A wrapper over bql_lock().
+ *
+ * This function is used to allow bindgen to generate bql_lock()
+ * binding.
+ *
+ * Do not call this function directly! Use bql_lock() instead.
+ */
+void rust_bql_lock(void);
+
 /**
  * bql_unlock: Unlock the Big QEMU Lock (BQL).
  *
@@ -383,6 +393,16 @@ void bql_lock_impl(const char *file, int line);
  */
 void bql_unlock(void);
 
+/**
+ * @rust_bql_unlock: A wrapper over bql_unlock().
+ *
+ * This function is used to allow bindgen to generate bql_unlock()
+ * binding.
+ *
+ * Do not call this function directly! Use bql_unlock() instead.
+ */
+void rust_bql_unlock(void);
+
 /**
  * BQL_LOCK_GUARD
  *
diff --git a/rust/bql/src/lib.rs b/rust/bql/src/lib.rs
index ef08221e9c1a..48cc76a7557c 100644
--- a/rust/bql/src/lib.rs
+++ b/rust/bql/src/lib.rs
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 mod bindings;
-use bindings::{bql_block_unlock, bql_locked, rust_bql_mock_lock};
+use bindings::{bql_block_unlock, bql_locked, rust_bql_lock, rust_bql_mock_lock, rust_bql_unlock};
 
 mod cell;
 pub use cell::*;
@@ -27,3 +27,102 @@ pub fn block_unlock(increase: bool) {
         bql_block_unlock(increase);
     }
 }
+
+// this function is private since user should get BQL context via BqlGuard.
+fn lock() {
+    // SAFETY: the function locks bql which lifetime is enough to cover
+    // the device's entire lifetime.
+    unsafe {
+        rust_bql_lock();
+    }
+}
+
+// this function is private since user should get BQL context via BqlGuard.
+fn unlock() {
+    // SAFETY: the function unlocks bql which lifetime is enough to
+    // cover the device's entire lifetime.
+    unsafe {
+        rust_bql_unlock();
+    }
+}
+
+/// An RAII guard to ensure a block of code runs within the BQL context.
+///
+/// It checks if the BQL is already locked at its creation:
+/// * If not, it locks the BQL and will release it when it is dropped.
+/// * If yes, it blocks BQL unlocking until its lifetime is end.
+#[must_use]
+pub struct BqlGuard {
+    locked: bool,
+}
+
+impl BqlGuard {
+    /// Creates a new `BqlGuard` to ensure BQL is locked during its
+    /// lifetime.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bql::{BqlCell, BqlGuard};
+    ///
+    /// fn foo() {
+    ///     let _guard = BqlGuard::new(); // BQL is locked
+    ///
+    ///     let c = BqlCell::new(5);
+    ///     assert_eq!(c.get(), 5);
+    /// } // BQL could be unlocked
+    /// ```
+    pub fn new() -> Self {
+        if !is_locked() {
+            lock();
+            Self { locked: true }
+        } else {
+            block_unlock(true);
+            Self { locked: false }
+        }
+    }
+}
+
+impl Default for BqlGuard {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Drop for BqlGuard {
+    fn drop(&mut self) {
+        if self.locked {
+            unlock();
+        } else {
+            block_unlock(false);
+        }
+    }
+}
+
+/// Executes a closure (function) within the BQL context.
+///
+/// This function creates a `BqlGuard`.
+///
+/// # Examples
+///
+/// ```should_panic
+/// use bql::{with_guard, BqlRefCell};
+///
+/// let c = BqlRefCell::new(5);
+///
+/// with_guard(|| {
+///     // BQL is locked
+///     let m = c.borrow();
+///
+///     assert_eq!(*m, 5);
+/// }); // BQL could be unlocked
+///
+/// let b = c.borrow(); // this causes a panic
+/// ```
+pub fn with_guard<F, R>(f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    let _guard = BqlGuard::new();
+    f()
+}
diff --git a/stubs/iothread-lock.c b/stubs/iothread-lock.c
index c89c9c7228f3..e2ebce565a06 100644
--- a/stubs/iothread-lock.c
+++ b/stubs/iothread-lock.c
@@ -23,6 +23,17 @@ void bql_unlock(void)
     assert(!bql_unlock_blocked);
 }
 
+void rust_bql_lock(void)
+{
+    rust_bql_mock_lock();
+}
+
+void rust_bql_unlock(void)
+{
+    bql_unlock();
+    bql_is_locked = false;
+}
+
 void bql_block_unlock(bool increase)
 {
     uint32_t new_value;
diff --git a/system/cpus.c b/system/cpus.c
index ef2d2f241faa..da31bda59444 100644
--- a/system/cpus.c
+++ b/system/cpus.c
@@ -578,6 +578,11 @@ void bql_lock_impl(const char *file, int line)
     bql_lock_fn(&bql, file, line);
 }
 
+void rust_bql_lock(void)
+{
+    bql_lock();
+}
+
 void bql_unlock(void)
 {
     g_assert(bql_locked());
@@ -585,6 +590,11 @@ void bql_unlock(void)
     qemu_mutex_unlock(&bql);
 }
 
+void rust_bql_unlock(void)
+{
+    bql_unlock();
+}
+
 void qemu_cond_wait_bql(QemuCond *cond)
 {
     qemu_cond_wait(cond, &bql);
-- 
2.34.1