surrealcs_kernel/

allocator.rs

//! An allocator for mapping an index to a value T.
use std::collections::VecDeque;

use nanoservices_utils::errors::{NanoServiceError, NanoServiceErrorStatus};

/// An allocator for mapping an index to a value T.
///
/// # Fields
/// * `connection_pool` - A vector of optional values of type T that can be allocated and deallocated, and accessed by index.
/// * `free_connections` - A queue of indices that have been deallocated and are available for allocation.
/// * `allocated_connections` - A queue of indices that have been allocated and are currently in use.
#[derive(Debug)]
pub struct Allocator<T: Clone> {
	pub connection_pool: Vec<Option<T>>,
	pub free_connections: VecDeque<usize>,
	pub allocated_connections: VecDeque<usize>,
}

impl<T: Clone> Default for Allocator<T> {
	fn default() -> Self {
		Self::new()
	}
}

impl<T: Clone> Allocator<T> {
	/// The constructor for the Allocator struct.
	///
	/// # Returns
	/// A new instance of the Allocator struct.
	pub fn new() -> Self {
		Self {
			connection_pool: Vec::new(),
			free_connections: VecDeque::new(),
			allocated_connections: VecDeque::new(),
		}
	}

	/// Allocates a connection in the connection pool.
	///
	/// # Arguments
	/// * `connection` - The connection to be allocated.
	///
	/// # Returns
	/// The index of the allocated connection to be used to access the connection in the connection pool.
	pub fn allocate(&mut self, connection: T) -> usize {
		if let Some(index) = self.free_connections.pop_front() {
			self.connection_pool[index] = Some(connection);
			self.allocated_connections.push_back(index);
			index
		} else {
			self.connection_pool.push(Some(connection));
			self.allocated_connections.push_back(self.connection_pool.len() - 1);
			self.connection_pool.len() - 1
		}
	}

	/// Yields the next allocated index in the connection pool.
	///
	/// # Returns
	/// The next allocated index in the connection pool. (empty if thee is no allocated connection available)
	pub fn yield_next_allocated_index(&mut self) -> Option<usize> {
		// for some reason this line throws a warning, however, the index is not
		// unused
		#[allow(unused_assignments)]
		let mut index: Option<usize> = None;

		loop {
			index = self.allocated_connections.pop_front();
			if index.is_none() {
				break;
			}
			let i = index.unwrap();
			let connection = self.connection_pool[i].as_ref();
			match connection {
				Some(_) => {
					self.allocated_connections.push_back(i);
					break;
				}
				None => continue,
			}
		}
		index
	}

	/// Extracts a connection from the connection pool.
	///
	/// # Arguments
	/// * `index` - The index of the connection to be extracted.
	///
	/// # Returns
	/// The connection at the specified index.
	pub fn extract_connection(&self, index: usize) -> Result<(T, usize), NanoServiceError> {
		if index + 1 > self.connection_pool.len() {
			return Err(NanoServiceError::new(
				format!("Connection at index {index} is unavailable"),
				NanoServiceErrorStatus::Unknown,
			));
		}
		let connection = self.connection_pool[index].as_ref();
		match connection {
			Some(c) => Ok((c.clone(), index)),
			None => Err(NanoServiceError::new(
				format!("Connection at index {index} is deallocated"),
				NanoServiceErrorStatus::Unknown,
			)),
		}
	}

	/// Deallocates a connection in the connection pool.
	///
	/// # Arguments
	/// * `index` - The index of the connection to be deallocated.
	pub fn deallocate(&mut self, index: usize) -> Result<T, NanoServiceError> {
		if index + 1 > self.connection_pool.len() {
			return Err(NanoServiceError::new(
				format!("Connection at index {index} is unavailable"),
				NanoServiceErrorStatus::Unknown,
			));
		}
		let placeholder = self.connection_pool[index].take();
		if placeholder.is_none() {
			return Err(NanoServiceError::new(
				format!("Connection at index {index} is already deallocated"),
				NanoServiceErrorStatus::Unknown,
			));
		}
		self.free_connections.push_back(index);
		Ok(placeholder.unwrap())
	}

	/// Yields all the connections in the connection pool.
	///
	/// # Returns
	/// A vector of all the connections in the connection pool.
	pub fn yield_connections(&self) -> Vec<T> {
		let mut connections = Vec::new();
		for i in &self.allocated_connections {
			let connection = self.connection_pool[*i].as_ref();
			match connection {
				Some(c) => connections.push(c.clone()),
				None => continue,
			}
		}
		connections
	}
}

#[cfg(test)]
mod tests {

	use super::*;

	#[test]
	fn test_new() {
		let allocator = Allocator::<usize>::new();
		assert_eq!(allocator.connection_pool.len(), 0);
		assert_eq!(allocator.free_connections.len(), 0);
		assert_eq!(allocator.allocated_connections.len(), 0);
	}

	#[test]
	fn test_allocate() {
		let mut allocator = Allocator::<usize>::new();
		let index = allocator.allocate(5);
		assert_eq!(index, 0);
		assert_eq!(allocator.connection_pool[index], Some(5));
		assert_eq!(allocator.allocated_connections.len(), 1);
		assert_eq!(allocator.free_connections.len(), 0);
		assert_eq!(allocator.connection_pool.len(), 1);

		// ensures that there is only one connection in the pool
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));

		let index = allocator.allocate(6);
		assert_eq!(index, 1);
		assert_eq!(allocator.connection_pool[index], Some(6));
		assert_eq!(allocator.allocated_connections.len(), 2);
		assert_eq!(allocator.free_connections.len(), 0);

		// ensures that we cycle between the two connections
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
	}

	#[test]
	fn test_deallocate() {
		let mut allocator = Allocator::<usize>::new();
		let index = allocator.allocate(5);
		let index_two = allocator.allocate(6);

		assert_eq!(index, 0);
		assert_eq!(index_two, 1);
		assert_eq!(allocator.connection_pool[index_two], Some(6));
		assert_eq!(allocator.allocated_connections.len(), 2);
		assert_eq!(allocator.free_connections.len(), 0);

		// ensures that we cycle between the two connections
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));

		assert_eq!(Ok(5), allocator.deallocate(0));
		assert_eq!(allocator.connection_pool[0], None);
		assert_eq!(allocator.allocated_connections.len(), 2);
		assert_eq!(allocator.free_connections.len(), 1);
		assert_eq!(allocator.free_connections[0], 0);

		// only yields the second connection
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));

		// ensure that the first dead connection is recycled
		let index = allocator.allocate(7);
		assert_eq!(index, 0);
		assert_eq!(allocator.connection_pool[index], Some(7));

		// ensures that we cycle between the two connections
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));
		assert_eq!(allocator.yield_next_allocated_index(), Some(0));
		assert_eq!(allocator.yield_next_allocated_index(), Some(1));

		// test an allocation that is out of bounds
		let result = allocator.deallocate(20);
		assert!(result.is_err());
		if let Err(e) = result {
			assert_eq!(e.message, "Connection at index 20 is unavailable");
		}
	}

	#[test]
	fn test_double_deallocate() {
		let mut allocator = Allocator::<usize>::new();
		let _ = allocator.allocate(5);
		let _ = allocator.allocate(6);

		assert_eq!(Ok(5), allocator.deallocate(0));
		assert_eq!(Ok(6), allocator.deallocate(1));

		assert_eq!(allocator.connection_pool[0], None);
		assert_eq!(allocator.connection_pool[1], None);
		assert_eq!(allocator.allocated_connections.len(), 2);
		assert_eq!(allocator.free_connections.len(), 2);
		assert_eq!(allocator.free_connections[0], 0);
		assert_eq!(allocator.free_connections[1], 1);

		// deallocate something that has already been deallocated
		let result = allocator.deallocate(0);
		assert!(result.is_err());
		if let Err(e) = result {
			assert_eq!(e.message, "Connection at index 0 is already deallocated");
		}
	}

	#[test]
	fn test_extract_connection() {
		let mut allocator = Allocator::<usize>::new();

		let outcome = allocator.extract_connection(20);
		assert!(outcome.is_err());
		if let Err(e) = outcome {
			assert_eq!(e.message, "Connection at index 20 is unavailable");
		}

		let index = allocator.allocate(5);
		let connection = allocator.extract_connection(index);
		let result = connection.unwrap();
		assert_eq!(result.0, 5);
		assert_eq!(result.1, 0);

		let _ = allocator.deallocate(index);
		let outcome = allocator.extract_connection(index);
		assert!(outcome.is_err());
		if let Err(e) = outcome {
			assert_eq!(e.message, "Connection at index 0 is deallocated");
		}
	}

	#[test]
	fn test_yield_connections() {
		let mut allocator = Allocator::<usize>::new();
		let _ = allocator.allocate(5);
		let _ = allocator.allocate(6);

		let connections = allocator.yield_connections();
		assert_eq!(connections.len(), 2);
		assert_eq!(connections[0], 5);
		assert_eq!(connections[1], 6);

		let _ = allocator.deallocate(0);
		let connections = allocator.yield_connections();
		assert_eq!(connections.len(), 1);
		assert_eq!(connections[0], 6);
	}
}