2 files changed, 127 insertions, 0 deletions
diff --git a/third_party/rust/oneshot-uniffi/tests/helpers/mod.rs b/third_party/rust/oneshot-uniffi/tests/helpers/mod.rs
new file mode 100644
index 0000000000..1b145396e7
--- /dev/null
+++ b/third_party/rust/oneshot-uniffi/tests/helpers/mod.rs
@@ -0,0 +1,63 @@
+#![allow(dead_code)]
+
+extern crate alloc;
+
+#[cfg(not(loom))]
+use alloc::sync::Arc;
+#[cfg(not(loom))]
+use core::sync::atomic::{AtomicUsize, Ordering::SeqCst};
+#[cfg(loom)]
+use loom::sync::{
+    atomic::{AtomicUsize, Ordering::SeqCst},
+    Arc,
+};
+
+#[cfg(loom)]
+pub mod waker;
+
+pub fn maybe_loom_model(test: impl Fn() + Sync + Send + 'static) {
+    #[cfg(loom)]
+    loom::model(test);
+    #[cfg(not(loom))]
+    test();
+}
+
+pub struct DropCounter<T> {
+    drop_count: Arc<AtomicUsize>,
+    value: Option<T>,
+}
+
+pub struct DropCounterHandle(Arc<AtomicUsize>);
+
+impl<T> DropCounter<T> {
+    pub fn new(value: T) -> (Self, DropCounterHandle) {
+        let drop_count = Arc::new(AtomicUsize::new(0));
+        (
+            Self {
+                drop_count: drop_count.clone(),
+                value: Some(value),
+            },
+            DropCounterHandle(drop_count),
+        )
+    }
+
+    pub fn value(&self) -> &T {
+        self.value.as_ref().unwrap()
+    }
+
+    pub fn into_value(mut self) -> T {
+        self.value.take().unwrap()
+    }
+}
+
+impl DropCounterHandle {
+    pub fn count(&self) -> usize {
+        self.0.load(SeqCst)
+    }
+}
+
+impl<T> Drop for DropCounter<T> {
+    fn drop(&mut self) {
+        self.drop_count.fetch_add(1, SeqCst);
+    }
+}
diff --git a/third_party/rust/oneshot-uniffi/tests/helpers/waker.rs b/third_party/rust/oneshot-uniffi/tests/helpers/waker.rs
new file mode 100644
index 0000000000..2e3f1bee19
--- /dev/null
+++ b/third_party/rust/oneshot-uniffi/tests/helpers/waker.rs
@@ -0,0 +1,64 @@
+//! Creates a Waker that can be observed from tests.
+
+use std::mem::forget;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::Arc;
+use std::task::{RawWaker, RawWakerVTable, Waker};
+
+#[derive(Default)]
+pub struct WakerHandle {
+    clone_count: AtomicU32,
+    drop_count: AtomicU32,
+    wake_count: AtomicU32,
+}
+
+impl WakerHandle {
+    pub fn clone_count(&self) -> u32 {
+        self.clone_count.load(Ordering::Relaxed)
+    }
+
+    pub fn drop_count(&self) -> u32 {
+        self.drop_count.load(Ordering::Relaxed)
+    }
+
+    pub fn wake_count(&self) -> u32 {
+        self.wake_count.load(Ordering::Relaxed)
+    }
+}
+
+pub fn waker() -> (Waker, Arc<WakerHandle>) {
+    let waker_handle = Arc::new(WakerHandle::default());
+    let waker_handle_ptr = Arc::into_raw(waker_handle.clone());
+    let raw_waker = RawWaker::new(waker_handle_ptr as *const _, waker_vtable());
+    (unsafe { Waker::from_raw(raw_waker) }, waker_handle)
+}
+
+pub(super) fn waker_vtable() -> &'static RawWakerVTable {
+    &RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw)
+}
+
+unsafe fn clone_raw(data: *const ()) -> RawWaker {
+    let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
+    handle.clone_count.fetch_add(1, Ordering::Relaxed);
+    forget(handle.clone());
+    forget(handle);
+    RawWaker::new(data, waker_vtable())
+}
+
+unsafe fn wake_raw(data: *const ()) {
+    let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
+    handle.wake_count.fetch_add(1, Ordering::Relaxed);
+    handle.drop_count.fetch_add(1, Ordering::Relaxed);
+}
+
+unsafe fn wake_by_ref_raw(data: *const ()) {
+    let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
+    handle.wake_count.fetch_add(1, Ordering::Relaxed);
+    forget(handle)
+}
+
+unsafe fn drop_raw(data: *const ()) {
+    let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
+    handle.drop_count.fetch_add(1, Ordering::Relaxed);
+    drop(handle)
+}