From 80b2c98a68d4c68f372f849e8d26f7cae5867000 Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@rivosinc.com>
Date: Thu, 21 Nov 2024 15:32:45 +0000
Subject: [PATCH] vm-allocator: Add an allocator KVM memory slots

Add an allocator for KVM memory slots, a finite resource in the kernel
- these need to be manipulated across different crates with a common
dependency on this crate.

Signed-off-by: Rob Bradford <rbradford@rivosinc.com>
---
 vm-allocator/src/lib.rs         |  2 ++
 vm-allocator/src/memory_slot.rs | 39 +++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 vm-allocator/src/memory_slot.rs

diff --git a/vm-allocator/src/lib.rs b/vm-allocator/src/lib.rs
index 801eff9faa..d560482889 100644
--- a/vm-allocator/src/lib.rs
+++ b/vm-allocator/src/lib.rs
@@ -21,3 +21,5 @@ pub use crate::gsi::GsiAllocator;
 #[cfg(target_arch = "x86_64")]
 pub use crate::gsi::GsiApic;
 pub use crate::system::SystemAllocator;
+mod memory_slot;
+pub use memory_slot::MemorySlotAllocator;
diff --git a/vm-allocator/src/memory_slot.rs b/vm-allocator/src/memory_slot.rs
new file mode 100644
index 0000000000..63a41d9c53
--- /dev/null
+++ b/vm-allocator/src/memory_slot.rs
@@ -0,0 +1,39 @@
+// Copyright © 2024 Rivos Inc
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::{Arc, Mutex};
+
+/// Allocator for KVM memory slots
+pub struct MemorySlotAllocator {
+    next_memory_slot: Arc<AtomicU32>,
+    memory_slot_free_list: Arc<Mutex<Vec<u32>>>,
+}
+
+impl MemorySlotAllocator {
+    /// Next free memory slot
+    pub fn next_memory_slot(&self) -> u32 {
+        if let Some(slot_id) = self.memory_slot_free_list.lock().unwrap().pop() {
+            return slot_id;
+        }
+        self.next_memory_slot.fetch_add(1, Ordering::SeqCst)
+    }
+
+    /// Release memory slot for reuse
+    pub fn free_memory_slot(&mut self, slot: u32) {
+        self.memory_slot_free_list.lock().unwrap().push(slot)
+    }
+
+    /// Instantiate struct
+    pub fn new(
+        next_memory_slot: Arc<AtomicU32>,
+        memory_slot_free_list: Arc<Mutex<Vec<u32>>>,
+    ) -> Self {
+        Self {
+            next_memory_slot,
+            memory_slot_free_list,
+        }
+    }
+}