common/handle_map/src/map.rs - beto-core - Git at Google

 // Copyright 2024 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.

 use crate::guard::{Downcast, LookupByHandle, ReadGuard, ReadWriteGuard};
 use crate::shard::{
     BoxedAnyProvider, HandleMapShard, ShardAllocationError, ShardValue, ValueProvider,
 };
 use crate::Handle;

 use core::any::Any;
 use core::convert::Infallible;
 use core::fmt::Debug;
 use core::marker::PhantomData;
 use core::ops::{Deref, DerefMut};
 use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};

 /// 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;

 /// Errors which may be raised while attempting to allocate
 /// a handle from contents given by a (fallible) value-provider.
 #[derive(Debug)]
 pub enum HandleMapTryAllocateError<E: Debug> {
     /// The call to the value-provider for the allocation failed.
     ValueProviderFailed(E),
     /// We couldn't reserve a spot for the allocation, because
     /// the handle-map was full.
     HandleMapFull,
 }

 /// 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> {
     /// Internal type-erased implementation. This allows most of the code to be shared for
     /// code-size reasons.
     erased: ErasedHandleMap,
     /// Keep `T` around so that we can downcast values from the erased implementation.
     _t: PhantomData<T>,
 }

 impl<T: Any + Send + Sync + 'static> HandleMap<T> {
     /// Creates a new handle-map with the given `HandleMapDimensions`.
     pub fn with_dimensions(dimensions: HandleMapDimensions) -> Self {
         Self {
             erased: ErasedHandleMap::with_dimensions(dimensions),
             _t: PhantomData,
         }
     }

     /// 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`.
     ///
     /// If you want the passed closure to be able to possibly fail, see
     /// [`Self::try_allocate`] instead.
     pub fn allocate(
         &self,
         initial_value_provider: impl FnOnce() -> T,
     ) -> Result<Handle, HandleMapFullError> {
         let value_provider =
             BoxedAnyProvider(move || Ok::<_, Infallible>(initial_value_provider()));
         self.erased
             .try_allocate_impl(value_provider)
             .map_err(|e| match e {
                 HandleMapTryAllocateError::ValueProviderFailed(never) => match never {},
                 HandleMapTryAllocateError::HandleMapFull => HandleMapFullError,
             })
     }

     /// Attempts to allocate 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
     /// `HandleMapTryAllocateError::HandleMapFull` and the given `initial_value_provider` will not
     /// be run.
     ///
     /// The passed initial-value provider may also fail, in which case this will return the error
     /// wrapped in `HandleMapTryAllocateError::ValueProviderFailed`.
     ///
     /// If your initial-value provider is infallible, see [`Self::allocate`] instead.
     pub fn try_allocate<E: Debug>(
         &self,
         initial_value_provider: impl FnOnce() -> Result<T, E>,
     ) -> Result<Handle, HandleMapTryAllocateError<E>> {
         let value_provider = BoxedAnyProvider(initial_value_provider);
         self.erased.try_allocate_impl(value_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<impl Deref<Target = T> + '_, HandleNotPresentError> {
         self.erased
             .get(handle)
             .map(|read_guard| read_guard.map(Downcast::<T>::new()))
     }

     /// 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<impl DerefMut<Target = T> + '_, HandleNotPresentError> {
         self.erased
             .get_mut(handle)
             .map(|read_write_guard| read_write_guard.map(Downcast::<T>::new()))
     }

     /// 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> {
         // Values should always match T
         #[allow(clippy::expect_used)]
         let downcast = |boxed: ShardValue| {
             *boxed
                 .downcast::<T>()
                 .expect("Handle map value has invalid type")
         };
         self.erased.deallocate(handle).map(downcast)
     }

     /// Gets the actual number of elements stored in the entire map.
     pub fn get_current_allocation_count(&self) -> u32 {
         self.erased.get_current_allocation_count()
     }

     /// 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.erased.set_new_handle_id_counter(value)
     }
 }

 /// Internal type-erased handle map implementation.
 struct ErasedHandleMap {
     /// 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]>,

     /// 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 ErasedHandleMap {
     /// Creates a new handle-map with the given `HandleMapDimensions`.
     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),
         }
     }

     /// Try allocate using a type erased ValueProvider.
     fn try_allocate_impl<E: Debug>(
         &self,
         mut value_provider: impl ValueProvider<E>,
     ) -> Result<Handle, HandleMapTryAllocateError<E>> {
         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_allocate_result = self.get_shard_from_handle(new_handle).try_allocate(
                 new_handle,
                 value_provider,
                 &self.outstanding_allocations_counter,
                 self.dimensions.max_active_handles,
             );
             match shard_allocate_result {
                 Ok(_) => {
                     return Ok(new_handle);
                 }
                 Err(ShardAllocationError::ValueProviderFailed(e)) => {
                     return Err(HandleMapTryAllocateError::ValueProviderFailed(e))
                 }
                 Err(ShardAllocationError::ExceedsAllocationLimit) => {
                     return Err(HandleMapTryAllocateError::HandleMapFull);
                 }
                 Err(ShardAllocationError::EntryOccupied(thrown_back_provider)) => {
                     // We need to do the whole thing again with a new ID
                     value_provider = thrown_back_provider;
                 }
             }
         }
     }

     fn get(&self, handle: Handle) -> Result<ReadGuard<LookupByHandle>, HandleNotPresentError> {
         self.get_shard_from_handle(handle).get(handle)
     }

     fn get_mut(
         &self,
         handle: Handle,
     ) -> Result<ReadWriteGuard<LookupByHandle>, HandleNotPresentError> {
         self.get_shard_from_handle(handle).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`].
     fn deallocate(&self, handle: Handle) -> Result<ShardValue, HandleNotPresentError> {
         self.get_shard_from_handle(handle)
             .deallocate(handle, &self.outstanding_allocations_counter)
     }

     /// Lookup the shard containing the given handle.
     fn get_shard_from_handle(&self, handle: Handle) -> &HandleMapShard {
         let shard_index = handle.get_shard_index(self.dimensions.num_shards);
         #[allow(clippy::expect_used)]
         self.handle_map_shards
             .get(shard_index)
             // This should never occur due to `get_shard_index` respecting the given `num_shards`.
             .expect("shard_index is always in range")
     }

     /// Gets the actual number of elements stored in the entire map.
     pub fn get_current_allocation_count(&self) -> u32 {
         self.outstanding_allocations_counter.load(Ordering::Relaxed)
     }

     /// 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);
     }
 }
	// Copyright 2024 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.

	use crate::guard::{Downcast, LookupByHandle, ReadGuard, ReadWriteGuard};
	use crate::shard::{
	BoxedAnyProvider, HandleMapShard, ShardAllocationError, ShardValue, ValueProvider,
	};
	use crate::Handle;

	use core::any::Any;
	use core::convert::Infallible;
	use core::fmt::Debug;
	use core::marker::PhantomData;
	use core::ops::{Deref, DerefMut};
	use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};

	/// 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;

	/// Errors which may be raised while attempting to allocate
	/// a handle from contents given by a (fallible) value-provider.
	#[derive(Debug)]
	pub enum HandleMapTryAllocateError<E: Debug> {
	/// The call to the value-provider for the allocation failed.
	ValueProviderFailed(E),
	/// We couldn't reserve a spot for the allocation, because
	/// the handle-map was full.
	HandleMapFull,
	}

	/// 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> {
	/// Internal type-erased implementation. This allows most of the code to be shared for
	/// code-size reasons.
	erased: ErasedHandleMap,
	/// Keep `T` around so that we can downcast values from the erased implementation.
	_t: PhantomData<T>,
	}

	impl<T: Any + Send + Sync + 'static> HandleMap<T> {
	/// Creates a new handle-map with the given `HandleMapDimensions`.
	pub fn with_dimensions(dimensions: HandleMapDimensions) -> Self {
	Self {
	erased: ErasedHandleMap::with_dimensions(dimensions),
	_t: PhantomData,
	}
	}

	/// 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`.
	///
	/// If you want the passed closure to be able to possibly fail, see
	/// [`Self::try_allocate`] instead.
	pub fn allocate(
	&self,
	initial_value_provider: impl FnOnce() -> T,
	) -> Result<Handle, HandleMapFullError> {
	let value_provider =
	BoxedAnyProvider(move \|\| Ok::<_, Infallible>(initial_value_provider()));
	self.erased
	.try_allocate_impl(value_provider)
	.map_err(\|e\| match e {
	HandleMapTryAllocateError::ValueProviderFailed(never) => match never {},
	HandleMapTryAllocateError::HandleMapFull => HandleMapFullError,
	})
	}

	/// Attempts to allocate 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
	/// `HandleMapTryAllocateError::HandleMapFull` and the given `initial_value_provider` will not
	/// be run.
	///
	/// The passed initial-value provider may also fail, in which case this will return the error
	/// wrapped in `HandleMapTryAllocateError::ValueProviderFailed`.
	///
	/// If your initial-value provider is infallible, see [`Self::allocate`] instead.
	pub fn try_allocate<E: Debug>(
	&self,
	initial_value_provider: impl FnOnce() -> Result<T, E>,
	) -> Result<Handle, HandleMapTryAllocateError<E>> {
	let value_provider = BoxedAnyProvider(initial_value_provider);
	self.erased.try_allocate_impl(value_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<impl Deref<Target = T> + '_, HandleNotPresentError> {
	self.erased
	.get(handle)
	.map(\|read_guard\| read_guard.map(Downcast::<T>::new()))
	}

	/// 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<impl DerefMut<Target = T> + '_, HandleNotPresentError> {
	self.erased
	.get_mut(handle)
	.map(\|read_write_guard\| read_write_guard.map(Downcast::<T>::new()))
	}

	/// 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> {
	// Values should always match T
	#[allow(clippy::expect_used)]
	let downcast = \|boxed: ShardValue\| {
	*boxed
	.downcast::<T>()
	.expect("Handle map value has invalid type")
	};
	self.erased.deallocate(handle).map(downcast)
	}

	/// Gets the actual number of elements stored in the entire map.
	pub fn get_current_allocation_count(&self) -> u32 {
	self.erased.get_current_allocation_count()
	}

	/// 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.erased.set_new_handle_id_counter(value)
	}
	}

	/// Internal type-erased handle map implementation.
	struct ErasedHandleMap {
	/// 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]>,

	/// 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 ErasedHandleMap {
	/// Creates a new handle-map with the given `HandleMapDimensions`.
	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),
	}
	}

	/// Try allocate using a type erased ValueProvider.
	fn try_allocate_impl<E: Debug>(
	&self,
	mut value_provider: impl ValueProvider<E>,
	) -> Result<Handle, HandleMapTryAllocateError<E>> {
	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_allocate_result = self.get_shard_from_handle(new_handle).try_allocate(
	new_handle,
	value_provider,
	&self.outstanding_allocations_counter,
	self.dimensions.max_active_handles,
	);
	match shard_allocate_result {
	Ok(_) => {
	return Ok(new_handle);
	}
	Err(ShardAllocationError::ValueProviderFailed(e)) => {
	return Err(HandleMapTryAllocateError::ValueProviderFailed(e))
	}
	Err(ShardAllocationError::ExceedsAllocationLimit) => {
	return Err(HandleMapTryAllocateError::HandleMapFull);
	}
	Err(ShardAllocationError::EntryOccupied(thrown_back_provider)) => {
	// We need to do the whole thing again with a new ID
	value_provider = thrown_back_provider;
	}
	}
	}
	}

	fn get(&self, handle: Handle) -> Result<ReadGuard<LookupByHandle>, HandleNotPresentError> {
	self.get_shard_from_handle(handle).get(handle)
	}

	fn get_mut(
	&self,
	handle: Handle,
	) -> Result<ReadWriteGuard<LookupByHandle>, HandleNotPresentError> {
	self.get_shard_from_handle(handle).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`].
	fn deallocate(&self, handle: Handle) -> Result<ShardValue, HandleNotPresentError> {
	self.get_shard_from_handle(handle)
	.deallocate(handle, &self.outstanding_allocations_counter)
	}

	/// Lookup the shard containing the given handle.
	fn get_shard_from_handle(&self, handle: Handle) -> &HandleMapShard {
	let shard_index = handle.get_shard_index(self.dimensions.num_shards);
	#[allow(clippy::expect_used)]
	self.handle_map_shards
	.get(shard_index)
	// This should never occur due to `get_shard_index` respecting the given `num_shards`.
	.expect("shard_index is always in range")
	}

	/// Gets the actual number of elements stored in the entire map.
	pub fn get_current_allocation_count(&self) -> u32 {
	self.outstanding_allocations_counter.load(Ordering::Relaxed)
	}

	/// 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);
	}
	}