nearby/util/handle_map/src/lib.rs - beto-core - Git at Google

 // Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 //! A thread-safe implementation of a map for managing object handles,
 //! a safer alternative to raw pointers for FFI interop.

 #![forbid(unsafe_code)]
 #![deny(missing_docs)]

 use std::boxed::Box;
 use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
 use std::vec::Vec;

 pub mod declare_handle_map;
 mod guard;
 pub(crate) mod shard;

 #[cfg(test)]
 mod tests;

 pub use guard::{ObjectReadGuardImpl, ObjectReadWriteGuardImpl};

 use shard::{HandleMapShard, ShardAllocationError};

 /// An individual handle to be given out by a [`HandleMap`].
 /// This representation is untyped, and just a wrapper
 /// around a handle-id, in contrast to implementors of `HandleLike`.
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Handle {
     handle_id: u64,
 }

 impl From<&Handle> for Handle {
     fn from(handle: &Handle) -> Self {
         *handle
     }
 }

 impl Handle {
     /// Constructs a handle wrapping the given ID.
     ///
     /// No validity checks are done on the wrapped ID
     /// to ensure that the given ID is active in
     /// any specific handle-map, and the type
     /// of the handle is not represented.
     ///
     /// As a result, this method is only useful for
     /// allowing access to handles from an FFI layer.
     pub fn from_id(handle_id: u64) -> Self {
         Self { handle_id }
     }

     /// Gets the ID for this handle.
     ///
     /// Since the underlying handle is un-typed,`
     /// this method is only suitable for
     /// transmitting handles across an FFI layer.
     pub fn get_id(&self) -> u64 {
         self.handle_id
     }

     /// Derives the shard index from the handle id
     fn get_shard_index(&self, num_shards: u8) -> usize {
         (self.handle_id % (num_shards as u64)) as usize
     }
 }

 /// Error raised when attempting to allocate into a full handle-map.
 #[derive(Debug)]
 pub struct HandleMapFullError;

 /// Error raised when the entry for a given [`Handle`] doesn't exist.
 #[derive(Debug)]
 pub struct HandleNotPresentError;

 /// FFI-transmissible structure expressing the dimensions
 /// (max # of allocatable slots, number of shards) of a handle-map
 /// to be used upon initialization.
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub struct HandleMapDimensions {
     /// The number of shards which are employed
     /// by the associated handle-map.
     pub num_shards: u8,
     /// The maximum number of active handles which may be
     /// stored within the associated handle-map.
     pub max_active_handles: u32,
 }

 /// A thread-safe mapping from "handle"s [like pointers, but safer]
 /// to underlying structures, supporting allocations, reads, writes,
 /// and deallocations of objects behind handles.
 pub struct HandleMap<T: Send + Sync> {
     /// The dimensions of this handle-map
     dimensions: HandleMapDimensions,

     /// The individually-lockable "shards" of the handle-map,
     /// among which the keys will be roughly uniformly-distributed.
     handle_map_shards: Box<[HandleMapShard<T>]>,

     /// An atomically-incrementing counter which tracks the
     /// next handle ID which allocations will attempt to use.
     new_handle_id_counter: AtomicU64,

     /// An atomic integer roughly tracking the number of
     /// currently-outstanding allocated entries in this
     /// handle-map among all [`HandleMapShard`]s.
     outstanding_allocations_counter: AtomicU32,
 }

 impl<T: Send + Sync> HandleMap<T> {
     /// Creates a new handle-map with the given `HandleMapDimensions`.
     pub fn with_dimensions(dimensions: HandleMapDimensions) -> Self {
         let mut handle_map_shards = Vec::with_capacity(dimensions.num_shards as usize);
         for _ in 0..dimensions.num_shards {
             handle_map_shards.push(HandleMapShard::default());
         }
         let handle_map_shards = handle_map_shards.into_boxed_slice();
         Self {
             dimensions,
             handle_map_shards,
             new_handle_id_counter: AtomicU64::new(0),
             outstanding_allocations_counter: AtomicU32::new(0),
         }
     }
 }

 impl<T: Send + Sync> HandleMap<T> {
     /// Allocates a new object within the given handle-map, returning
     /// a handle to the location it was stored at. This operation
     /// may fail if attempting to allocate over the `dimensions.max_active_handles`
     /// limit imposed on the handle-map, in which case this method
     /// will return a `HandleMapFullError`.
     pub fn allocate(
         &self,
         initial_value_provider: impl FnOnce() -> T,
     ) -> Result<Handle, HandleMapFullError> {
         let mut initial_value_provider = initial_value_provider;
         loop {
             // Increment the new-handle-ID counter using relaxed memory ordering,
             // since the only invariant that we want to enforce is that concurrently-running
             // threads always get distinct new handle-ids.
             let new_handle_id = self.new_handle_id_counter.fetch_add(1, Ordering::Relaxed);
             let new_handle = Handle::from_id(new_handle_id);
             let shard_index = new_handle.get_shard_index(self.dimensions.num_shards);

             // Now, check the shard to see if we can actually allocate into it.
             let shard_allocate_result = self.handle_map_shards[shard_index].try_allocate(
                 new_handle,
                 initial_value_provider,
                 &self.outstanding_allocations_counter,
                 self.dimensions.max_active_handles,
             );
             match shard_allocate_result {
                 Ok(_) => {
                     return Ok(new_handle);
                 }
                 Err(ShardAllocationError::ExceedsAllocationLimit) => {
                     return Err(HandleMapFullError);
                 }
                 Err(ShardAllocationError::EntryOccupied(thrown_back_provider)) => {
                     // We need to do the whole thing again with a new ID
                     initial_value_provider = thrown_back_provider;
                 }
             }
         }
     }

     /// Gets a read-only reference to an object within the given handle-map,
     /// if the given handle is present. Otherwise, returns [`HandleNotPresentError`].
     pub fn get(&self, handle: Handle) -> Result<ObjectReadGuardImpl<T>, HandleNotPresentError> {
         let shard_index = handle.get_shard_index(self.dimensions.num_shards);
         self.handle_map_shards[shard_index].get(handle)
     }

     /// Gets a read+write reference to an object within the given handle-map,
     /// if the given handle is present. Otherwise, returns [`HandleNotPresentError`].
     pub fn get_mut(
         &self,
         handle: Handle,
     ) -> Result<ObjectReadWriteGuardImpl<T>, HandleNotPresentError> {
         let shard_index = handle.get_shard_index(self.dimensions.num_shards);
         self.handle_map_shards[shard_index].get_mut(handle)
     }

     /// Removes the object pointed to by the given handle in
     /// the handle-map, returning the removed object if it
     /// exists. Otherwise, returns [`HandleNotPresentError`].
     pub fn deallocate(&self, handle: Handle) -> Result<T, HandleNotPresentError> {
         let shard_index = handle.get_shard_index(self.dimensions.num_shards);
         self.handle_map_shards[shard_index]
             .deallocate(handle, &self.outstanding_allocations_counter)
     }

     /// Gets the actual number of elements stored in the entire map.
     /// Only suitable for single-threaded sections of tests.
     #[cfg(test)]
     pub(crate) fn len(&self) -> usize {
         self.handle_map_shards.iter().map(|s| s.len()).sum()
     }

     /// Sets the new-handle-id counter to the given value.
     /// Only suitable for tests.
     #[cfg(test)]
     pub(crate) fn set_new_handle_id_counter(&mut self, value: u64) {
         self.new_handle_id_counter = AtomicU64::new(value);
     }
 }

 /// Externally-facing trait for things which behave like handle-map handles
 /// with a globally-defined handle-map for the type.
 pub trait HandleLike: Sized {
     /// The underlying object type pointed-to by this handle
     type Object: Send + Sync;

     /// Tries to allocate a new handle using the given provider
     /// to construct the underlying stored object as a new
     /// entry into the global handle table for this type.
     fn allocate(
         initial_value_provider: impl FnOnce() -> Self::Object,
     ) -> Result<Self, HandleMapFullError>;

     /// Gets a RAII read-guard on the contents behind this handle.
     fn get(&self) -> Result<ObjectReadGuardImpl<Self::Object>, HandleNotPresentError>;

     /// Gets a RAII read-write guard on the contents behind this handle.
     fn get_mut(&self) -> Result<ObjectReadWriteGuardImpl<Self::Object>, HandleNotPresentError>;

     /// Deallocates the contents behind this handle.
     fn deallocate(self) -> Result<Self::Object, HandleNotPresentError>;
 }
	// Copyright 2023 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	//! A thread-safe implementation of a map for managing object handles,
	//! a safer alternative to raw pointers for FFI interop.

	#![forbid(unsafe_code)]
	#![deny(missing_docs)]

	use std::boxed::Box;
	use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
	use std::vec::Vec;

	pub mod declare_handle_map;
	mod guard;
	pub(crate) mod shard;

	#[cfg(test)]
	mod tests;

	pub use guard::{ObjectReadGuardImpl, ObjectReadWriteGuardImpl};

	use shard::{HandleMapShard, ShardAllocationError};

	/// An individual handle to be given out by a [`HandleMap`].
	/// This representation is untyped, and just a wrapper
	/// around a handle-id, in contrast to implementors of `HandleLike`.
	#[derive(Clone, Copy, PartialEq, Eq, Hash)]
	pub struct Handle {
	handle_id: u64,
	}

	impl From<&Handle> for Handle {
	fn from(handle: &Handle) -> Self {
	*handle
	}
	}

	impl Handle {
	/// Constructs a handle wrapping the given ID.
	///
	/// No validity checks are done on the wrapped ID
	/// to ensure that the given ID is active in
	/// any specific handle-map, and the type
	/// of the handle is not represented.
	///
	/// As a result, this method is only useful for
	/// allowing access to handles from an FFI layer.
	pub fn from_id(handle_id: u64) -> Self {
	Self { handle_id }
	}

	/// Gets the ID for this handle.
	///
	/// Since the underlying handle is un-typed,`
	/// this method is only suitable for
	/// transmitting handles across an FFI layer.
	pub fn get_id(&self) -> u64 {
	self.handle_id
	}

	/// Derives the shard index from the handle id
	fn get_shard_index(&self, num_shards: u8) -> usize {
	(self.handle_id % (num_shards as u64)) as usize
	}
	}

	/// Error raised when attempting to allocate into a full handle-map.
	#[derive(Debug)]
	pub struct HandleMapFullError;

	/// Error raised when the entry for a given [`Handle`] doesn't exist.
	#[derive(Debug)]
	pub struct HandleNotPresentError;

	/// FFI-transmissible structure expressing the dimensions
	/// (max # of allocatable slots, number of shards) of a handle-map
	/// to be used upon initialization.
	#[repr(C)]
	#[derive(Clone, Copy)]
	pub struct HandleMapDimensions {
	/// The number of shards which are employed
	/// by the associated handle-map.
	pub num_shards: u8,
	/// The maximum number of active handles which may be
	/// stored within the associated handle-map.
	pub max_active_handles: u32,
	}

	/// A thread-safe mapping from "handle"s [like pointers, but safer]
	/// to underlying structures, supporting allocations, reads, writes,
	/// and deallocations of objects behind handles.
	pub struct HandleMap<T: Send + Sync> {
	/// The dimensions of this handle-map
	dimensions: HandleMapDimensions,

	/// The individually-lockable "shards" of the handle-map,
	/// among which the keys will be roughly uniformly-distributed.
	handle_map_shards: Box<[HandleMapShard<T>]>,

	/// An atomically-incrementing counter which tracks the
	/// next handle ID which allocations will attempt to use.
	new_handle_id_counter: AtomicU64,

	/// An atomic integer roughly tracking the number of
	/// currently-outstanding allocated entries in this
	/// handle-map among all [`HandleMapShard`]s.
	outstanding_allocations_counter: AtomicU32,
	}

	impl<T: Send + Sync> HandleMap<T> {
	/// Creates a new handle-map with the given `HandleMapDimensions`.
	pub fn with_dimensions(dimensions: HandleMapDimensions) -> Self {
	let mut handle_map_shards = Vec::with_capacity(dimensions.num_shards as usize);
	for _ in 0..dimensions.num_shards {
	handle_map_shards.push(HandleMapShard::default());
	}
	let handle_map_shards = handle_map_shards.into_boxed_slice();
	Self {
	dimensions,
	handle_map_shards,
	new_handle_id_counter: AtomicU64::new(0),
	outstanding_allocations_counter: AtomicU32::new(0),
	}
	}
	}

	impl<T: Send + Sync> HandleMap<T> {
	/// Allocates a new object within the given handle-map, returning
	/// a handle to the location it was stored at. This operation
	/// may fail if attempting to allocate over the `dimensions.max_active_handles`
	/// limit imposed on the handle-map, in which case this method
	/// will return a `HandleMapFullError`.
	pub fn allocate(
	&self,
	initial_value_provider: impl FnOnce() -> T,
	) -> Result<Handle, HandleMapFullError> {
	let mut initial_value_provider = initial_value_provider;
	loop {
	// Increment the new-handle-ID counter using relaxed memory ordering,
	// since the only invariant that we want to enforce is that concurrently-running
	// threads always get distinct new handle-ids.
	let new_handle_id = self.new_handle_id_counter.fetch_add(1, Ordering::Relaxed);
	let new_handle = Handle::from_id(new_handle_id);
	let shard_index = new_handle.get_shard_index(self.dimensions.num_shards);

	// Now, check the shard to see if we can actually allocate into it.
	let shard_allocate_result = self.handle_map_shards[shard_index].try_allocate(
	new_handle,
	initial_value_provider,
	&self.outstanding_allocations_counter,
	self.dimensions.max_active_handles,
	);
	match shard_allocate_result {
	Ok(_) => {
	return Ok(new_handle);
	}
	Err(ShardAllocationError::ExceedsAllocationLimit) => {
	return Err(HandleMapFullError);
	}
	Err(ShardAllocationError::EntryOccupied(thrown_back_provider)) => {
	// We need to do the whole thing again with a new ID
	initial_value_provider = thrown_back_provider;
	}
	}
	}
	}

	/// Gets a read-only reference to an object within the given handle-map,
	/// if the given handle is present. Otherwise, returns [`HandleNotPresentError`].
	pub fn get(&self, handle: Handle) -> Result<ObjectReadGuardImpl<T>, HandleNotPresentError> {
	let shard_index = handle.get_shard_index(self.dimensions.num_shards);
	self.handle_map_shards[shard_index].get(handle)
	}

	/// Gets a read+write reference to an object within the given handle-map,
	/// if the given handle is present. Otherwise, returns [`HandleNotPresentError`].
	pub fn get_mut(
	&self,
	handle: Handle,
	) -> Result<ObjectReadWriteGuardImpl<T>, HandleNotPresentError> {
	let shard_index = handle.get_shard_index(self.dimensions.num_shards);
	self.handle_map_shards[shard_index].get_mut(handle)
	}

	/// Removes the object pointed to by the given handle in
	/// the handle-map, returning the removed object if it
	/// exists. Otherwise, returns [`HandleNotPresentError`].
	pub fn deallocate(&self, handle: Handle) -> Result<T, HandleNotPresentError> {
	let shard_index = handle.get_shard_index(self.dimensions.num_shards);
	self.handle_map_shards[shard_index]
	.deallocate(handle, &self.outstanding_allocations_counter)
	}

	/// Gets the actual number of elements stored in the entire map.
	/// Only suitable for single-threaded sections of tests.
	#[cfg(test)]
	pub(crate) fn len(&self) -> usize {
	self.handle_map_shards.iter().map(\|s\| s.len()).sum()
	}

	/// Sets the new-handle-id counter to the given value.
	/// Only suitable for tests.
	#[cfg(test)]
	pub(crate) fn set_new_handle_id_counter(&mut self, value: u64) {
	self.new_handle_id_counter = AtomicU64::new(value);
	}
	}

	/// Externally-facing trait for things which behave like handle-map handles
	/// with a globally-defined handle-map for the type.
	pub trait HandleLike: Sized {
	/// The underlying object type pointed-to by this handle
	type Object: Send + Sync;

	/// Tries to allocate a new handle using the given provider
	/// to construct the underlying stored object as a new
	/// entry into the global handle table for this type.
	fn allocate(
	initial_value_provider: impl FnOnce() -> Self::Object,
	) -> Result<Self, HandleMapFullError>;

	/// Gets a RAII read-guard on the contents behind this handle.
	fn get(&self) -> Result<ObjectReadGuardImpl<Self::Object>, HandleNotPresentError>;

	/// Gets a RAII read-write guard on the contents behind this handle.
	fn get_mut(&self) -> Result<ObjectReadWriteGuardImpl<Self::Object>, HandleNotPresentError>;

	/// Deallocates the contents behind this handle.
	fn deallocate(self) -> Result<Self::Object, HandleNotPresentError>;
	}