From a5ac802687d2a606f83da287bb6c8eca50db9277 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tuomas=20M=C3=A4kinen?=
 <1947505+tuommaki@users.noreply.github.com>
Date: Thu, 21 Mar 2024 12:29:30 +0200
Subject: [PATCH] Improve CID allocations (#159)

- Allocating CIDs sequentially decreases probability for conflicts.
- Using Set for CID allocations provides better API.
---
 crates/node/src/vmm/qemu.rs | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/crates/node/src/vmm/qemu.rs b/crates/node/src/vmm/qemu.rs
index b2bfc3dd..1ba78daa 100644
--- a/crates/node/src/vmm/qemu.rs
+++ b/crates/node/src/vmm/qemu.rs
@@ -11,11 +11,14 @@ use rand::{distributions::Alphanumeric, Rng};
 use serde_json::json;
 use std::{
     any::Any,
-    collections::HashMap,
+    collections::{BTreeSet, HashMap},
     fs::File,
     path::Path,
     process::{Child, Command, Stdio},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    },
     time::{Duration, Instant},
 };
 use tokio::{
@@ -64,7 +67,8 @@ pub struct QEMUVMHandle {
 
 pub struct Qemu {
     config: Arc<Config>,
-    cid_allocations: Vec<u32>,
+    next_cid: AtomicU32,
+    cid_allocations: BTreeSet<u32>,
     vm_registry: HashMap<u32, QEMUVMHandle>,
 }
 
@@ -72,33 +76,32 @@ impl Qemu {
     pub fn new(config: Arc<Config>) -> Self {
         Qemu {
             config,
-            cid_allocations: vec![],
+            next_cid: AtomicU32::new(4),
+            cid_allocations: Default::default(),
             vm_registry: HashMap::new(),
         }
     }
 
     fn allocate_cid(&mut self) -> u32 {
         loop {
-            let cid = rand::random::<u32>();
+            let cid = self.next_cid.fetch_add(1, Ordering::Relaxed);
 
             if cid < 3 {
                 // CIDs 0, 1 and 2 are reserved.
                 continue;
             }
 
-            if self.cid_allocations.iter().any(|&x| x == cid) {
+            if !self.cid_allocations.insert(cid) {
                 // Generated CID found from existing allocations.
                 continue;
             };
 
-            self.cid_allocations.push(cid);
             return cid;
         }
     }
 
     fn release_cid(&mut self, cid: u32) {
-        if let Some(idx) = self.cid_allocations.iter().position(|&x| x == cid) {
-            self.cid_allocations.remove(idx);
+        if self.cid_allocations.remove(&cid) {
             self.vm_registry.remove(&cid);
         }
     }