new allocator service

Cargo.toml

@@ -41,7 +41,7 @@ num-traits = "0.2"
 tempfile = "3.5"
 
 [workspace]
-members = [".", "samples/command_executer", "driver-bindings", "eif_loader", "enclave_build", "vsock_proxy"]
+members = [".", "samples/command_executer", "driver-bindings", "eif_loader", "enclave_build", "vsock_proxy","allocator"]
 
 [features]
 default = []

allocator/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "allocator"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+log = "0.4.22"
+serde = {version = "1.0.217", features = ["derive"] }
+serde_yaml = "0.9.34"
+thiserror = "2.0.9"
+lazy_static = "1.4"
+
+
+[dev-dependencies]
+tempfile = "3.2"

allocator/src/configuration.rs

@@ -0,0 +1,47 @@
+use serde::Deserialize;
+use crate::resources;
+use crate::error::Error;
+
+//deserializing from allocator.yaml file
+#[derive(Debug, PartialEq, Deserialize,Clone)]
+#[serde(deny_unknown_fields)]
+#[serde(untagged)]
+pub enum ResourcePool {
+    CpuCount { memory_mib: usize , cpu_count: usize},
+    CpuPool { cpu_pool: String, memory_mib: usize },
+}
+pub fn get_resource_pool_from_config()  -> Result<Vec<ResourcePool>, Box<dyn std::error::Error>> {
+    //config file deserializing
+    let f = std::fs::File::open("/etc/nitro_enclaves/allocator.yaml")?;
+    let pool: Vec<ResourcePool> =  match serde_yaml::from_reader(f) {
+        Ok(pool) => pool,
+        Err(_) => {return Err(Box::new(Error::ConfigFileCorruption));},//error messages use anyhow
+    };
+    if pool.len() > 4 {
+       eprintln!("{}",Error::MoreResourcePoolThanSupported);
+    }
+    Ok(pool)
+}
+pub fn get_current_allocated_cpu_pool() -> Result<Option<std::collections::BTreeSet::<usize>>, Box<dyn std::error::Error>> {
+    let f = std::fs::read_to_string("/sys/module/nitro_enclaves/parameters/ne_cpus")?;
+    if f.trim().is_empty() {
+        return Ok(None);
+    }
+    let cpu_list = resources::cpu::parse_cpu_list(&f[..])?;
+    Ok(Some(cpu_list))
+}
+//clears everything in a numa node.
+pub fn clear_everything_in_numa_node() -> Result<(), Box<dyn std::error::Error>> {//change the name
+    match get_current_allocated_cpu_pool()?{
+		Some(cpu_list) => {
+		//find numa by one of cpuids
+		let numa = resources::cpu::get_numa_node_for_cpu(cpu_list.clone().into_iter().next().unwrap())?;
+		//release everything
+		let _ = resources::huge_pages::release_all_huge_pages(numa)?;
+		let _ = resources::cpu::deallocate_cpu_set(&cpu_list);
+		}
+		None => {}  
+  	};
+    Ok(())
+}
+

allocator/src/error.rs

@@ -0,0 +1,14 @@
+#[derive(thiserror::Error, Debug)]
+pub enum Error
+{
+	#[error(transparent)]
+	ParseInt(#[from] std::num::ParseIntError),
+	#[error(transparent)]
+	TryFromInt(#[from] std::num::TryFromIntError),
+	#[error(transparent)]
+	Allocation(#[from] super::resources::Error),
+	#[error("Invalid config file. This might happened due to old config file or config file corruption. See release notes :")]
+	ConfigFileCorruption,
+	#[error("WARNING! Requested resource pool is more than supported. Supported Enclave number is 4")]
+	MoreResourcePoolThanSupported,
+}

allocator/src/main.rs

@@ -0,0 +1,34 @@
+mod resources;
+mod error;
+mod configuration;
+
+
+fn main()  -> Result<(), Box<dyn std::error::Error>> {
+	let _ = configuration::clear_everything_in_numa_node();
+
+	match configuration::get_resource_pool_from_config() {
+    	Ok(pool) => {			
+			let numa_node = match resources::Allocation::allocate_by_cpu_pools(pool.clone()){
+				Ok(numa) => numa,
+				Err(e) =>{
+					eprintln!("Allocation failed: {}",e);
+					return Err(Box::new(e));
+				},//proper error messages				
+			};
+      		match resources::Allocation::allocate_by_cpu_count(pool,numa_node) {
+				Ok(_) => {},
+				Err(e) => {
+					let _ = configuration::clear_everything_in_numa_node();
+					eprintln!(" Allocation failed: {}",e);
+					return Err(Box::new(e));
+				}
+			} //check if allocation successful or not, if not clear what you allocated previously
+    	}
+    	Err(e) => {
+			eprintln!("Allocation failed: {}",e);
+      		return Err(e);
+    	}
+
+  	};
+  	Ok(())
+}