[PATCH 2/3] rust: sync: add SetOnce::try_get_or_populate()

[Alice Ryhl <aliceryhl@google.com>]

The SetOnce::populate() method does not internally synchronize callers
that fail to populate the value with the successful call. This means
that naive loops using as_ref() and populate() can lead to spinning on
the initialization, which is best avoided. Thus, provide a helper that
avoids this issue using a user-provided lock.

One potential alternative is to change populate() so that the failing
caller actually does synchronize with the successful call to populate().
However, this is somewhat tricky:

* There are users of SetOnce that construct it in const context, and we
  currently don't have the ability to do that for most locks, so we
  cannot easily add a lock to SetOnce.
* Just spinning on the atomic is undesirable unless we disable
  preemption in the success path. If we do disable preemption, then that
  raises complications for handling the PREEMPT_RT case.
* It also raises questions about deadlocks if populate() is called from
  irqs.

By using a user-provided lock, we do not have to worry about these
issues inside SetOnce.

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 rust/kernel/sync/set_once.rs | 43 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/rust/kernel/sync/set_once.rs b/rust/kernel/sync/set_once.rs
index a78f8c8e87db..d9cc598a8d78 100644
--- a/rust/kernel/sync/set_once.rs
+++ b/rust/kernel/sync/set_once.rs
@@ -2,11 +2,18 @@
 
 //! A container that can be initialized at most once.
 
-use super::atomic::{
-    ordering::{Acquire, Relaxed, Release},
-    Atomic,
-};
 use core::{cell::UnsafeCell, mem::MaybeUninit};
+use kernel::sync::{
+    atomic::{
+        ordering::{
+            Acquire,
+            Relaxed,
+            Release, //
+        },
+        Atomic, //
+    },
+    lock, //
+};
 
 /// A container that can be populated at most once. Thread safe.
 ///
@@ -104,6 +111,34 @@ pub fn populate(&self, value: T) -> Result<&T, T> {
         }
     }
 
+    /// Get the value, or populate it if it's missing.
+    ///
+    /// This method is useful to avoid spinning on the internal atomic state. If all writers call
+    /// this method with the same lock, then they are synchronized with each other and it's
+    /// guaranteed that no caller will attempt to invoke [`SetOnce::populate`] more than once.
+    pub fn try_get_or_populate<F, E, U, B>(&self, lock: &lock::Lock<U, B>, f: F) -> Result<&T, E>
+    where
+        B: lock::Backend,
+        F: FnOnce() -> Result<T, E>,
+    {
+        if let Some(value) = self.as_ref() {
+            return Ok(value);
+        }
+
+        let mut to_insert = f()?;
+        loop {
+            if let Some(value) = self.as_ref() {
+                return Ok(value);
+            }
+
+            let _guard = lock.lock();
+            match self.populate(to_insert) {
+                Ok(value) => return Ok(value),
+                Err(ret) => to_insert = ret,
+            }
+        }
+    }
+
     /// Get a copy of the contained object.
     ///
     /// Returns [`None`] if the [`SetOnce`] is empty.

-- 
2.55.0.229.g6434b31f56-goog