diff --git a/src/lib.rs b/src/lib.rs index 572d722..43009e8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,7 @@ pub mod hex; pub mod inetaddress; pub mod io; pub mod memory; +pub mod oncearc; pub mod ringbuffer; pub mod str; pub mod sync; diff --git a/src/oncearc.rs b/src/oncearc.rs new file mode 100644 index 0000000..d47089b --- /dev/null +++ b/src/oncearc.rs @@ -0,0 +1,94 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::mem::transmute; +use std::ptr::null_mut; +use std::sync::atomic::{AtomicPtr, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +pub struct OnceArc { + /// A simple pointer for fast reads if non-null, otherwise falls through to atomic access. + /// This is written in init_once() but may not be instantly visible everywhere, but in this + /// case accesses still work via the atomic pointer. + fast_ptr: *mut T, + /// The real "authoritative" pointer is atomic. + ptr: AtomicPtr, +} + +impl OnceArc { + #[inline] + pub fn new() -> Self { + Self { fast_ptr: null_mut(), ptr: AtomicPtr::new(null_mut()) } + } + + /// Initialize the value of this OnceArc. + /// This will panic if it is called more than once. + #[inline] + pub fn init_once(&self, obj: Arc) { + let obj: *mut T = unsafe { transmute(Arc::into_raw(obj)) }; + if self + .ptr + .compare_exchange(null_mut(), obj, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + unsafe { + std::ptr::write_volatile(transmute(&self.fast_ptr as *const *mut T), obj); + } + } else { + panic!("OnceArc can only be initialized once"); + } + } + + /// Load this value, or return None if not yet initialized. + #[inline(always)] + pub fn load(&self) -> Option<&T> { + if !self.fast_ptr.is_null() { + Some(unsafe { &*self.fast_ptr }) + } else { + let ptr = self.ptr.load(Ordering::Acquire); + if ptr.is_null() { + None + } else { + Some(unsafe { &*ptr }) + } + } + } + + /// Load this value or busy wait (with a short delay) until available. + /// This should be used when the value is expected to be either available or initialized + /// pretty much immediately, such as in the case of two concurrent dependencies initializing + /// each other. + #[inline(always)] + pub fn load_wait(&self) -> &T { + if !self.fast_ptr.is_null() { + return unsafe { &*self.fast_ptr }; + } else { + loop { + let ptr = self.ptr.load(Ordering::Acquire); + if ptr.is_null() { + std::thread::sleep(Duration::from_millis(1)); + } else { + return unsafe { &*ptr }; + } + } + } + } +} + +unsafe impl Sync for OnceArc {} +unsafe impl Send for OnceArc where T: Send {} + +impl Drop for OnceArc { + fn drop(&mut self) { + let obj = self.ptr.load(Ordering::Acquire); + if !obj.is_null() { + unsafe { drop(Arc::from_raw(obj)) }; + } + } +}