78 lines
1.5 KiB
Rust
78 lines
1.5 KiB
Rust
use std::{
|
|
marker::PhantomData,
|
|
sync::{
|
|
atomic::{AtomicPtr, Ordering},
|
|
Arc,
|
|
},
|
|
};
|
|
|
|
pub struct AtomicArc<T> {
|
|
inner: AtomicPtr<T>,
|
|
_phantom: PhantomData<Arc<T>>,
|
|
}
|
|
|
|
impl<T> Drop for AtomicArc<T> {
|
|
fn drop(&mut self) {
|
|
unsafe { Arc::decrement_strong_count(self.load()) }
|
|
}
|
|
}
|
|
|
|
impl<T> AtomicArc<T> {
|
|
fn load(&self) -> *const T {
|
|
self.inner.load(Ordering::Relaxed)
|
|
}
|
|
|
|
pub fn new(inner: Arc<T>) -> Self {
|
|
Self {
|
|
inner: AtomicPtr::new(Arc::into_raw(inner) as *mut T),
|
|
_phantom: PhantomData,
|
|
}
|
|
}
|
|
|
|
pub fn get(&self) -> Arc<T> {
|
|
unsafe {
|
|
let ptr = self.load();
|
|
Arc::increment_strong_count(ptr);
|
|
Arc::from_raw(ptr)
|
|
}
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn swap(&self, new: Arc<T>) -> Arc<T> {
|
|
unsafe {
|
|
let old = self
|
|
.inner
|
|
.swap(Arc::into_raw(new) as *mut T, Ordering::Relaxed);
|
|
|
|
Arc::from_raw(old)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::Arc;
|
|
|
|
use super::AtomicArc;
|
|
|
|
#[test]
|
|
fn test_basic() {
|
|
let atomic_arc = AtomicArc::new(Arc::new(1));
|
|
|
|
let first = atomic_arc.get();
|
|
assert_eq!(*first, 1);
|
|
|
|
_ = atomic_arc.swap(Arc::new(2));
|
|
|
|
assert_eq!(*first, 1);
|
|
|
|
Arc::into_inner(first).unwrap();
|
|
|
|
let second = atomic_arc.get();
|
|
|
|
std::mem::drop(atomic_arc);
|
|
|
|
Arc::into_inner(second).unwrap();
|
|
}
|
|
}
|