From ac0ee47b926e0220eba3c1ce8a0a14db63df2b43 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 1 May 2026 12:04:04 -0500 Subject: [PATCH 01/35] remove generic actor cache. this was done for two reasons: 1) the vm and agent actor caches have slightly different requirements, specifically around keys (actorid vs vmid) 2) the scheduler needs direct control of the cache so the cache can be updated in response to messages it forwards. being possibly up to 1 second out of date in this data could cause problems. In theory the scheduler could update the generic cache, but I don't think using a generic cache if we need to modify it directly anyway really makes sense. Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 345 +++++++++++++++++++------- odorobo/src/utils/actor_cache.rs | 174 ------------- odorobo/src/utils/mod.rs | 1 - 3 files changed, 250 insertions(+), 270 deletions(-) delete mode 100644 odorobo/src/utils/actor_cache.rs diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 11d6a62..df83778 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,12 +1,13 @@ use std::ops::ControlFlow; +use std::sync::Arc; +use std::time::Duration; -use async_trait::async_trait; use kameo::prelude::*; use libp2p::futures::TryStreamExt; +use tracing::trace; +use ulid::Ulid; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; -use crate::utils::actor_cache::ActorCache; -use crate::utils::actor_cache::ActorCacheUpdater; use crate::utils::actor_names::VM; use crate::messages::vm::*; use crate::messages::agent::*; @@ -17,12 +18,42 @@ use stable_eyre::eyre::OptionExt; use stable_eyre::{Report, Result, eyre::eyre}; use tracing::info_span; use tracing::{info, warn}; +use dashmap::DashMap; +use tokio::task::JoinHandle; + + +#[derive(Debug, Clone)] +pub struct CachedAgentActor { + pub actor_ref: RemoteActorRef, + pub metadata: AgentStatus, +} + +#[derive(Debug, Clone)] +pub struct CachedVMActor { + pub actor_ref: RemoteActorRef, + pub metadata: GetVMInfoReply, +} + #[derive(RemoteActor)] pub struct SchedulerActor { - pub agent_actor_cache: ActorCache, - pub vm_actor_cache: ActorCache + pub agent_data_cache: Arc>, + pub agent_keepalive_tasks: Arc>>, + + // todo: we might need a better way to store this. + // we are likely going to want to store vms even if we don't know their actorid (ex: actor hasn't been started or is shutdown) + // but we also might want ot be able to store them without a ulid, possibly + // so we might need a vec of vms and then to just store maps/indexes of actorid and ulid to vector index + // and then like a freelist or something. + // i dont really love that option either though cause it feels overkill. + // maybe we sure just be using a proper database entirely? + // idk. will figure it out later. + pub vm_actorid_ulid_map: Arc>, + pub vm_data_cache: Arc>, + pub vm_keepalive_tasks: Arc>>, + + pub cache_actor_finder: Option>, } // todo: this might need to be a runtime thing but this makes it easy to write for now and could easily be switched out later. @@ -31,20 +62,191 @@ static VCPU_OVERPROVISIONMENT_DENOMINATOR: u32 = 1; impl SchedulerActor { - async fn lookup_by_actor_id( + async fn lookup_agent_by_actor_id( &mut self, actor_id: &ActorId, ) -> Option> { - self.agent_actor_cache.data_cache.get(actor_id).map(|data| data.actor_ref.clone()) + self.agent_data_cache.get(actor_id).map(|data| data.actor_ref.clone()) } - async fn lookup_by_hostname( + async fn lookup_agent_by_hostname( &mut self, hostname: &str, ) -> Option> { - self.agent_actor_cache.data_cache.iter().find(|data| data.metadata.hostname == hostname).map(|data| data.actor_ref.clone()) + self.agent_data_cache.iter().find(|data| data.metadata.hostname == hostname).map(|data| data.actor_ref.clone()) + } + + + // someone should likely give caleb a firm talking to about code duplication due to this section, but things are just different enough that trying to make them one function requires usage of a lot of generics which feels even worse. so i dont know what to do. cappy please fix. i hate this. + async fn vm_actor_finder( + parent_actor_ref: RemoteActorRef, + vm_actorid_ulid_map: Arc>, + data_cache: Arc>, + keepalive_tasks: Arc>> + ) -> Result<(), Report> { + + while let Some(vm_actor) = RemoteActorRef::::lookup_all(VM).try_next().await? { + if !keepalive_tasks.contains_key(&vm_actor.id()) { + trace!(?vm_actor, "starting vm_updater_task"); + + parent_actor_ref.link_remote(&vm_actor).await?; + + let vm_actor_id = vm_actor.id(); + + let vm_actorid_ulid_map_clone = Arc::clone(&vm_actorid_ulid_map); + let data_cache_clone = Arc::clone(&data_cache); + let updater_task = tokio::spawn(async move { + Self::vm_updater_task( + vm_actor, + vm_actorid_ulid_map_clone, + data_cache_clone + ).await; + }); + + keepalive_tasks.insert( + vm_actor_id, + updater_task + ); + } + } + + Ok(()) + + } + + async fn vm_updater_task( + actor_ref: RemoteActorRef, + vm_actorid_ulid_map: Arc>, + data_cache: Arc>, + ) { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + let mut fails = 0; + loop { + if let Ok(metadata) = actor_ref.ask(&GetVMInfo {vmid: None}).await { + let vmid = metadata.vmid; + + vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that + + data_cache.insert( + vmid, + CachedVMActor { + actor_ref: actor_ref.clone(), + metadata: metadata + } + ); + + fails = 0; + } else { + fails += 1; + } + + if fails > 5 { + // todo: possibly better error handling + warn!(?actor_ref, "can no longer reach agent actor.") + } + + interval.tick().await; + } + } + + async fn agent_actor_finder( + parent_actor_ref: RemoteActorRef, + data_cache: Arc>, + keepalive_tasks: Arc>>, + ) -> Result<(), Report> { + info!("running agent_actor_finder"); + while let Some(agent_actor) = RemoteActorRef::::lookup_all(AGENT).try_next().await? { + if !keepalive_tasks.contains_key(&agent_actor.id()) { + trace!(?agent_actor, "starting agent_updater_task"); + + parent_actor_ref.link_remote(&agent_actor).await?; + + let agent_actor_id = agent_actor.id(); + + let data_cache_clone = Arc::clone(&data_cache); + let updater_task = tokio::spawn(async move { + Self::agent_updater_task( + agent_actor, + data_cache_clone + ).await; + }); + + keepalive_tasks.insert( + agent_actor_id, + updater_task + ); + } + } + + Ok(()) + } + + async fn agent_updater_task( + actor_ref: RemoteActorRef, + data_cache: Arc>, + ) { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + let mut fails = 0; + loop { + if let Ok(metadata) = actor_ref.ask(&GetAgentStatus).await { + data_cache.insert( + actor_ref.id(), + CachedAgentActor { + actor_ref: actor_ref.clone(), + metadata: metadata + } + ); + + fails = 0; + } else { + fails += 1; + } + + if fails > 5 { + // todo: possibly better error handling + warn!(?actor_ref, "can no longer reach agent actor.") + } + + interval.tick().await; + } + } + + fn start_actor_finder(&mut self, actor_ref: RemoteActorRef) { + let agent_data_cache_arc_clone = Arc::clone(&self.agent_data_cache); + let agent_keepalive_tasks_arc_clone = Arc::clone(&self.agent_keepalive_tasks); + + let vm_actorid_ulid_map_arc_clone = Arc::clone(&self.vm_actorid_ulid_map); + let vm_data_cache_arc_clone = Arc::clone(&self.vm_data_cache); + let vm_keepalive_tasks_arc_clone = Arc::clone(&self.vm_keepalive_tasks); + + self.cache_actor_finder = Some( + tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + loop { + let vm_join_handle = Self::vm_actor_finder( + actor_ref.clone(), + Arc::clone(&vm_actorid_ulid_map_arc_clone), + Arc::clone(&vm_data_cache_arc_clone), + Arc::clone(&vm_keepalive_tasks_arc_clone), + ); + + let agent_join_handle = Self::agent_actor_finder( + actor_ref.clone(), + Arc::clone(&agent_data_cache_arc_clone), + Arc::clone(&agent_keepalive_tasks_arc_clone), + ); + + // intentionally ignoring results because we want to keep finding actors even if an attempt fails + let _ = tokio::join!(vm_join_handle, agent_join_handle); + + interval.tick().await; + } + }) + ); + } + /// current scheduling algo info: /// this is vaguely based on https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ /// when a vm is attempted to be scheduled, we loop through every agent and score it based on some rules @@ -62,13 +264,12 @@ impl SchedulerActor { // todo: this arguably could be done as map-reduce. is that better? let span = info_span!("schedule_agent"); span.in_scope(|| { - for agent in self.agent_actor_cache.data_cache.iter() { + for agent in self.agent_data_cache.iter() { let mut agent_score = 0u32; let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; - if agent.metadata.used_vcpus >= agent_max_vcpus { continue; } @@ -114,80 +315,6 @@ impl SchedulerActor { -#[derive(Copy, Clone)] -struct AgentActorCacheUpdater; - -#[derive(Debug, Clone)] -pub struct CachedAgentActor { - pub actor_ref: RemoteActorRef, - pub metadata: AgentStatus, -} - -#[async_trait] -impl ActorCacheUpdater for AgentActorCacheUpdater { - async fn get_actor_refs(&self) -> Result>> { - let mut agent_actors_lookup = RemoteActorRef::::lookup_all(AGENT); - let mut actor_ref_vec = Vec::new(); - - while let Some(agent_actor) = agent_actors_lookup.try_next().await? { - actor_ref_vec.push(agent_actor); - } - - Ok(actor_ref_vec) - } - - async fn on_update(&self, actor_ref: &RemoteActorRef, previous_value: Option) -> Result { - let output_actor_ref = match previous_value { - Some(value) => value.actor_ref, - _ => actor_ref.clone(), - }; - - Ok(CachedAgentActor { - actor_ref: output_actor_ref, - metadata: actor_ref.ask(&GetAgentStatus).await? - }) - } -} - - -// todo: this code is really bad, and we should not have effectively two copies of ths same thing. -#[derive(Copy, Clone)] -struct VMActorCacheUpdater; - -#[derive(Debug, Clone)] -pub struct CachedVMActor { - pub actor_ref: RemoteActorRef, - pub metadata: GetVMInfoReply, -} - -#[async_trait] -impl ActorCacheUpdater for VMActorCacheUpdater { - async fn get_actor_refs(&self) -> Result>> { - let mut agent_actors_lookup = RemoteActorRef::::lookup_all(VM); - let mut actor_ref_vec = Vec::new(); - - while let Some(agent_actor) = agent_actors_lookup.try_next().await? { - actor_ref_vec.push(agent_actor); - } - - Ok(actor_ref_vec) - } - - async fn on_update(&self, actor_ref: &RemoteActorRef, previous_value: Option) -> Result { - let output_actor_ref = match previous_value { - Some(value) => value.actor_ref, - _ => actor_ref.clone(), - }; - - Ok(CachedVMActor { - actor_ref: output_actor_ref, - metadata: actor_ref.ask(&GetVMInfo {vmid: None}).await? - }) - } -} - - - impl Actor for SchedulerActor { type Args = (); type Error = Report; @@ -197,29 +324,57 @@ impl Actor for SchedulerActor { info!(?peer_id, "Scheduler Actor started!"); - Ok(Self { - agent_actor_cache: ActorCache::new(actor_ref.clone(), AgentActorCacheUpdater)?, - vm_actor_cache: ActorCache::new(actor_ref, VMActorCacheUpdater)? - }) + let mut scheduler_actor = SchedulerActor { + agent_data_cache: Arc::new(DashMap::new()), + agent_keepalive_tasks: Arc::new(DashMap::new()), + vm_actorid_ulid_map: Arc::new(DashMap::new()), + vm_data_cache: Arc::new(DashMap::new()), + vm_keepalive_tasks: Arc::new(DashMap::new()), + cache_actor_finder: None, + }; + + scheduler_actor.start_actor_finder(actor_ref.into_remote_ref().await); + + Ok(scheduler_actor) } async fn on_link_died( &mut self, actor_ref: WeakActorRef, - id: ActorId, + actor_id: ActorId, reason: ActorStopReason, ) -> Result, Self::Error> { - warn!("Linked actor {id:?} died with reason {reason:?}"); + warn!(?actor_id, ?reason, "Linked actor died"); + // check that scheduler actor is still alive. let Some(_) = actor_ref.upgrade() else { return Ok(ControlFlow::Break(ActorStopReason::Killed)); }; - self.agent_actor_cache.on_link_died(id).await; - self.vm_actor_cache.on_link_died(id).await; - info!(vm_actor_cache=?self.agent_actor_cache.data_cache, agent_actor_cache=?self.vm_actor_cache.data_cache, "data caches post actor removal"); + if let Some((_, keepalive_task)) = self.agent_keepalive_tasks.remove(&actor_id) { + trace!(?actor_id, "Aborting agent keepalive task"); + keepalive_task.abort(); + }; + + + self.agent_data_cache.remove(&actor_id); + + // todo: attempt vm migration or restart or whatever on agent death. + + if let Some((_, keepalive_task)) = self.vm_keepalive_tasks.remove(&actor_id) { + trace!(?actor_id, "Aborting vm keepalive task"); + keepalive_task.abort(); + }; + + if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&actor_id) { + // todo: we likely should keep a copy of the VirtualMachine manifest in the cache. + // instead of removing the vm entirely, we should just modify the status to shutdown or crashed or something. + trace!(?actor_id, ?vmid, "Removing vm from vm_data_cache"); + self.vm_data_cache.remove(&vmid); + } + Ok(ControlFlow::Continue(())) } @@ -300,7 +455,7 @@ impl Message for SchedulerActor { ) -> Self::Reply { let mut vms = Vec::new(); - for agent in self.agent_actor_cache.data_cache.iter() { + for agent in self.agent_data_cache.iter() { vms.extend_from_slice(agent.metadata.vms.as_slice()); } diff --git a/odorobo/src/utils/actor_cache.rs b/odorobo/src/utils/actor_cache.rs deleted file mode 100644 index eb4b45b..0000000 --- a/odorobo/src/utils/actor_cache.rs +++ /dev/null @@ -1,174 +0,0 @@ -use std::{marker::PhantomData, sync::Arc, time::Duration}; - -use async_trait::async_trait; -use dashmap::DashMap; -use kameo::{prelude::*}; -use tokio::task::JoinHandle; -use stable_eyre::{Report, Result}; -use tracing::{info, instrument, trace}; - -use std::fmt; - -// TODO: refactor to use derive macro, but I (caleb) don't know how to write a derive macro. -// -// The best way to write this would be using a derive similar to kameo -// so you would create a struct with #[derive(ActorCache)]. -// Then you would set the types and then write get_actor_refs and on_update. -// This would combine everything into one struct and make it a lot easier to work with. -// -// similar to: https://github.com/tqwewe/kameo/blob/1d498c0566b613b9afe6d54965c4b191c84432e0/src/actor.rs#L122 -// -// -// other things we might want: -// - a default get_actor_refs that just finds all actor_refs with a specific actor string. -// - change get_actor_refs to use an iterator - - -#[async_trait] -pub trait ActorCacheUpdater: Sync + Send + Copy + 'static { - // todo: this could probably be better if it was an iterator, but I am lazy and don't want to right now. - async fn get_actor_refs(&self) -> Result>>; - async fn on_update(&self, actor_ref: &RemoteActorRef, previous_value: Option) -> Result; -} - - -#[derive(Debug)] -pub struct ActorCache { - parent_actor_ref: ActorRef, - pub data_cache: Arc>, - keepalive_tasks: Arc>>, - actor_finder: Option>, - - child_actor_type: PhantomData -} - -// todo: impl Drop to automatically kill all the keepalive_tasks and the actor_finder task. - -impl ActorCache { - pub fn new( - parent_actor_ref: ActorRef, - updater: impl ActorCacheUpdater - ) -> Result { - - let data_cache = Arc::new(DashMap::new()); - let keepalive_tasks = Arc::new(DashMap::new()); - - let actor_cache = ActorCache { - parent_actor_ref: parent_actor_ref.clone(), - data_cache: data_cache, - keepalive_tasks: keepalive_tasks, - actor_finder: None, - - child_actor_type: PhantomData - }; - - actor_cache.start_actor_finder(parent_actor_ref, updater); - - Ok(actor_cache) - } - - /// run this function inside of the on_link_died of the ParentActor - pub async fn on_link_died( - &self, - id: ActorId - ) { - info!("removing agent actor from cache {id:?}"); - - if let Some(actor_keepalive_task) = self.keepalive_tasks.remove(&id) { - trace!("Aborting keepalive task for agent {id:?}"); - actor_keepalive_task.1.abort(); - }; - - self.data_cache.remove(&id); - } - - fn start_actor_finder( - &self, - parent_actor_ref: ActorRef, - updater: impl ActorCacheUpdater - ) { - let keepalive_tasks_clone = Arc::clone(&self.keepalive_tasks); - let data_cache_clone = Arc::clone(&self.data_cache); - - tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_secs(1)); - loop { - let _ = Self::actor_finder( - parent_actor_ref.clone(), - Arc::clone(&keepalive_tasks_clone), - Arc::clone(&data_cache_clone), - updater - ).await; - - interval.tick().await; - } - }); - } - - async fn actor_finder( - parent_actor_ref: ActorRef, - keepalive_tasks: Arc>>, - data_cache: Arc>, - updater: impl ActorCacheUpdater - ) -> Result<(), Report> { - let actor_refs = updater.get_actor_refs().await?; - - info!(?actor_refs, "running actor_finder"); - - for actor_ref in actor_refs { - if !keepalive_tasks.contains_key(&actor_ref.id()) { - trace!(?actor_ref, "starting updater_task"); - - parent_actor_ref.link_remote(&actor_ref).await?; - - let actor_ref_clone = actor_ref.clone(); - let data_cache_clone = Arc::clone(&data_cache); - let updater_task = tokio::spawn(async move { - Self::updater_task( - actor_ref_clone, - data_cache_clone, - updater - ).await; - }); - - keepalive_tasks.insert( - actor_ref.id(), - updater_task - ); - } - } - - Ok(()) - } - - #[instrument(skip_all)] - async fn updater_task( - actor_ref: RemoteActorRef, - data_cache: Arc>, - updater: impl ActorCacheUpdater - ) { - let mut interval = tokio::time::interval(Duration::from_secs(1)); - - loop { - let actor_id = actor_ref.id(); - - let mut previous_value_option = None; - - - - if let Some(data_ref) = data_cache.get(&actor_id) { - previous_value_option = Some(data_ref.clone()); - } - - - if let Ok(update) = updater.on_update(&actor_ref, previous_value_option).await { - data_cache.insert( - actor_id, - update.clone() - ); - } - - interval.tick().await; - } - } -} diff --git a/odorobo/src/utils/mod.rs b/odorobo/src/utils/mod.rs index 01c7125..819d3ef 100644 --- a/odorobo/src/utils/mod.rs +++ b/odorobo/src/utils/mod.rs @@ -1,5 +1,4 @@ pub mod actor_names; -pub mod actor_cache; use aide::OperationIo; use stable_eyre::{Result, Report}; From 09c41f063c2e0a7f56852175c3ecbc7e8af5a565 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 1 May 2026 12:31:31 -0500 Subject: [PATCH 02/35] improve scheduling algorithm to use estimated future values. Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index df83778..1ae3119 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -252,39 +252,39 @@ impl SchedulerActor { /// when a vm is attempted to be scheduled, we loop through every agent and score it based on some rules /// there are hard rules that will simply throw out an agent entirely. /// otherwise, we take whatever the best agent we can find is. - /// - /// additionally, because caleb is way too performance brained, he used integer math for the entire scoring algorithm just so we didnt have to convert to floats. async fn schedule_agent( &mut self, msg: &CreateVM ) -> Result, Report> { let mut best_agent = None; - let mut best_agent_score = 0u32; + let mut best_agent_score = 0.0f32; // todo: this arguably could be done as map-reduce. is that better? let span = info_span!("schedule_agent"); - span.in_scope(|| { + span.in_scope(|| { for agent in self.agent_data_cache.iter() { - let mut agent_score = 0u32; + let mut agent_score = 0.0f32; let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; + // todo: do we care about VMData.max_vcpus? + let agent_used_vcpus = agent.metadata.used_vcpus + msg.config.data.vcpus; - - if agent.metadata.used_vcpus >= agent_max_vcpus { + if agent_used_vcpus >= agent_max_vcpus { continue; } - agent_score += (agent_max_vcpus - agent.metadata.used_vcpus) * 1024 / agent_max_vcpus; + agent_score += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. let agent_max_ram = agent.metadata.ram; + let agent_used_ram = agent.metadata.used_ram + msg.config.data.memory; - if agent.metadata.used_ram >= agent_max_ram { + if agent_used_ram >= agent_max_ram { continue; } - agent_score += ((agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) * 1024 / agent_max_ram.as_u64()) as u32; + agent_score += (agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; // todo: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ From d0dd375d5fa19350a8e47fdba8eaac37462d7128 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 1 May 2026 13:36:11 -0500 Subject: [PATCH 03/35] committing before lunch so i dont have code just sitting on my machine Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 36 +++++++++++++++++---------- odorobo/src/ch_driver/actor.rs | 12 +++++---- odorobo/src/messages/vm.rs | 3 +-- 3 files changed, 31 insertions(+), 20 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 1ae3119..78569b2 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -387,20 +387,26 @@ impl Message for SchedulerActor { type Reply = Result; async fn handle(&mut self, msg: CreateVM, _ctx: &mut Context) -> Self::Reply { - loop { - let target_agent = self.schedule_agent(&msg).await?; - - match target_agent.ask(&msg).await { - Ok(reply) => { - return Ok(reply) - }, - Err(err) => { - warn!( - "CreateVM forwarding failed, trying again: {err}" - ); - } - } + let target_agent = self.schedule_agent(&msg).await?; + + // we add to cache first, because we want to make sure future requests assume this vm exists. if the message fails, we clean it up afterward. + + // massive problem for me to fix later: the metadata can be auto updated. so this change won't be persisted and we need to continue to know this vm is likely scheduled on this device. + // we likely need to create a new vec of likey? vms or something and then match against that too during scheduling for example. + + if let Some(mut cached_data) = self.agent_data_cache.get_mut(&target_agent.id()) { + cached_data.metadata.vms.push(msg.vmid); + } else { + return Err(eyre!("target agent is not in data cache")) + } + + let reply = target_agent.ask(&msg).await; + + if reply.is_err() { + // remove from caches } + + Ok(reply?) } } @@ -416,6 +422,7 @@ impl Message for SchedulerActor { tracing::trace!(?vm, "DeleteVM"); if let Some(vm) = vm { vm.tell(&msg).send()?; + // todo: update cache Ok(DeleteVMReply) } else { Err(eyre!("VM not found")) @@ -435,6 +442,9 @@ impl Message for SchedulerActor { tracing::trace!(?vm, "ShutdownVM"); if let Some(vm) = vm { vm.tell(&msg).send()?; + + // todo: update cache + Ok(ShutdownVMReply) } else { Err(eyre!("VM not found")) diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index 414b00b..8889b6d 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -28,16 +28,17 @@ pub struct VMActor { /// path to the Cloud Hypervisor socket, in /run/odorobo/vms//ch.sock pub vm_instance: VMInstance, pub migration_state: Option, + pub manifest: VirtualMachine } impl Actor for VMActor { - // tuple of VM ID and optional config - type Args = (ulid::Ulid, Option); + // tuple of VM ID and manifest + type Args = (ulid::Ulid, VirtualMachine); type Error = Report; #[tracing::instrument(skip_all)] - async fn on_start((vmid, vm_config): Self::Args, actor_ref: ActorRef) -> Result { - let mut vminstance = VMInstance::spawn(&vmid.to_string(), vm_config.map(VmConfig::from), None).await?; + async fn on_start((vmid, vm_manifest): Self::Args, actor_ref: ActorRef) -> Result { + let mut vminstance = VMInstance::spawn(&vmid.to_string(), Some(VmConfig::from(vm_manifest.clone())), None).await?; // Take the child process out so we can watch for unexpected death. // destroy() handles a missing child_process gracefully. @@ -69,6 +70,7 @@ impl Actor for VMActor { vmid, vm_instance: vminstance, migration_state: None, + manifest: vm_manifest }) } @@ -157,7 +159,7 @@ impl Message for VMActor { ) -> Self::Reply { GetVMInfoReply { vmid: self.vmid, - config: self.vm_instance.vm_config.clone(), + config: self.manifest.clone(), // we likely dont want to send the entire manifest on every update, but some of this data is required and this is easier for now. } } } diff --git a/odorobo/src/messages/vm.rs b/odorobo/src/messages/vm.rs index 512a793..4e979c8 100644 --- a/odorobo/src/messages/vm.rs +++ b/odorobo/src/messages/vm.rs @@ -1,5 +1,4 @@ //! VM-related messages -use cloud_hypervisor_client::models::VmConfig; use kameo::prelude::*; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -99,5 +98,5 @@ pub struct GetVMInfo { #[derive(Serialize, Deserialize, Reply, Debug, Clone)] pub struct GetVMInfoReply { pub vmid: Ulid, - pub config: Option, + pub config: VirtualMachine, } From 444950203e219b418f5178955d16eaf0f1ec13e8 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 1 May 2026 17:58:06 -0500 Subject: [PATCH 04/35] continue to work on scheduler actor caches and placement groups Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 66 ++++++++++++++++++++------- odorobo/src/ch_driver/actor.rs | 6 +-- odorobo/src/messages/vm.rs | 3 +- 3 files changed, 55 insertions(+), 20 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 78569b2..6ff9c3c 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -26,11 +26,14 @@ use tokio::task::JoinHandle; pub struct CachedAgentActor { pub actor_ref: RemoteActorRef, pub metadata: AgentStatus, + /// this is a list of all VMs that may be on an agent. it is used for rules such as affinity to make sure we don't schedule things in ways that arent allowed + /// We don't know for a fact these VMs are scheduled due to latency and boot up delay, but they may be scheduled. + pub extended_vm_list: Vec, } #[derive(Debug, Clone)] pub struct CachedVMActor { - pub actor_ref: RemoteActorRef, + pub actor_ref: Option>, pub metadata: GetVMInfoReply, } @@ -49,6 +52,9 @@ pub struct SchedulerActor { // i dont really love that option either though cause it feels overkill. // maybe we sure just be using a proper database entirely? // idk. will figure it out later. + // + // new related problem: i just realized vmids/ulids and actorids dont have to be unique. + // if a vm is migrating from one actor to another, there might be two actors with the same vmid. pub vm_actorid_ulid_map: Arc>, pub vm_data_cache: Arc>, pub vm_keepalive_tasks: Arc>>, @@ -130,7 +136,7 @@ impl SchedulerActor { data_cache.insert( vmid, CachedVMActor { - actor_ref: actor_ref.clone(), + actor_ref: Some(actor_ref.clone()), metadata: metadata } ); @@ -189,13 +195,25 @@ impl SchedulerActor { let mut fails = 0; loop { if let Ok(metadata) = actor_ref.ask(&GetAgentStatus).await { - data_cache.insert( - actor_ref.id(), - CachedAgentActor { - actor_ref: actor_ref.clone(), - metadata: metadata - } - ); + if data_cache.contains_key(&actor_ref.id()) { + data_cache.alter( + &actor_ref.id(), + |_, mut v| { + v.metadata = metadata; + + v + } + ); + } else { + data_cache.insert( + actor_ref.id(), + CachedAgentActor { + actor_ref: actor_ref.clone(), + metadata, + extended_vm_list: vec![] + } + ); + } fails = 0; } else { @@ -292,7 +310,7 @@ impl SchedulerActor { // todo (future): possibly keep a percent of agents completely empty, to be able to be converted to dedis automatically. - // they would have their agent score set to 1, so they can be scheduled to if there is no other avaliable agents. + // they would have their agent score set to like f32::MIN, so they can be scheduled to if there is no other avaliable agents. // rough pseudo code to implement this: // if agent.metadata.vms.len() == 0 && hash(agent.config.hostname) % total_chance < threshold { // agent_score = 1; @@ -390,21 +408,37 @@ impl Message for SchedulerActor { let target_agent = self.schedule_agent(&msg).await?; // we add to cache first, because we want to make sure future requests assume this vm exists. if the message fails, we clean it up afterward. - - // massive problem for me to fix later: the metadata can be auto updated. so this change won't be persisted and we need to continue to know this vm is likely scheduled on this device. - // we likely need to create a new vec of likey? vms or something and then match against that too during scheduling for example. - if let Some(mut cached_data) = self.agent_data_cache.get_mut(&target_agent.id()) { - cached_data.metadata.vms.push(msg.vmid); + cached_data.extended_vm_list.push(msg.vmid); } else { return Err(eyre!("target agent is not in data cache")) } + + self.vm_data_cache.insert( + msg.vmid, + CachedVMActor { + actor_ref: None, + metadata: GetVMInfoReply { + vmid: msg.vmid, + config: Some(msg.config.clone()) + } + }); + + let reply = target_agent.ask(&msg).await; + + // remove from caches if we fail to schedule if reply.is_err() { - // remove from caches + self.agent_data_cache.alter(&target_agent.id(), |_,mut v| { + v.extended_vm_list.retain(|&vmid| vmid != msg.vmid); + + v + }); + self.vm_data_cache.remove(&msg.vmid); } + Ok(reply?) } diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index 8889b6d..b3a9f92 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -28,17 +28,17 @@ pub struct VMActor { /// path to the Cloud Hypervisor socket, in /run/odorobo/vms//ch.sock pub vm_instance: VMInstance, pub migration_state: Option, - pub manifest: VirtualMachine + pub manifest: Option } impl Actor for VMActor { // tuple of VM ID and manifest - type Args = (ulid::Ulid, VirtualMachine); + type Args = (ulid::Ulid, Option); type Error = Report; #[tracing::instrument(skip_all)] async fn on_start((vmid, vm_manifest): Self::Args, actor_ref: ActorRef) -> Result { - let mut vminstance = VMInstance::spawn(&vmid.to_string(), Some(VmConfig::from(vm_manifest.clone())), None).await?; + let mut vminstance = VMInstance::spawn(&vmid.to_string(), vm_manifest.clone().map(VmConfig::from), None).await?; // Take the child process out so we can watch for unexpected death. // destroy() handles a missing child_process gracefully. diff --git a/odorobo/src/messages/vm.rs b/odorobo/src/messages/vm.rs index 4e979c8..8318a5a 100644 --- a/odorobo/src/messages/vm.rs +++ b/odorobo/src/messages/vm.rs @@ -1,4 +1,5 @@ //! VM-related messages +use cloud_hypervisor_client::models::VmConfig; use kameo::prelude::*; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -98,5 +99,5 @@ pub struct GetVMInfo { #[derive(Serialize, Deserialize, Reply, Debug, Clone)] pub struct GetVMInfoReply { pub vmid: Ulid, - pub config: VirtualMachine, + pub config: Option, } From 7a635181a56282c35707f2f70ef62d2e6d8fee66 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Tue, 5 May 2026 14:23:07 -0500 Subject: [PATCH 05/35] continue to work on scheduler actor (still not finished) Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 236 ++++++++++++++++++-------- odorobo/src/types.rs | 12 +- 2 files changed, 168 insertions(+), 80 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 6ff9c3c..9ed9536 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,13 +1,17 @@ use std::ops::ControlFlow; use std::sync::Arc; use std::time::Duration; +use std::cmp::Ordering; +use ahash::AHashSet; +use dashmap::mapref::multiple::RefMulti; use kameo::prelude::*; use libp2p::futures::TryStreamExt; use tracing::trace; use ulid::Ulid; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; +use crate::types::AffinityStrictness; use crate::utils::actor_names::VM; use crate::messages::vm::*; use crate::messages::agent::*; @@ -26,9 +30,9 @@ use tokio::task::JoinHandle; pub struct CachedAgentActor { pub actor_ref: RemoteActorRef, pub metadata: AgentStatus, - /// this is a list of all VMs that may be on an agent. it is used for rules such as affinity to make sure we don't schedule things in ways that arent allowed + /// this is a set of all VMs that may be on an agent. it is used for rules such as affinity to make sure we don't schedule things in ways that arent allowed /// We don't know for a fact these VMs are scheduled due to latency and boot up delay, but they may be scheduled. - pub extended_vm_list: Vec, + pub extended_vm_set: AHashSet, } #[derive(Debug, Clone)] @@ -38,23 +42,35 @@ pub struct CachedVMActor { } - +// todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. +// especially when we haven't fully made decisions about some other things. #[derive(RemoteActor)] pub struct SchedulerActor { pub agent_data_cache: Arc>, pub agent_keepalive_tasks: Arc>>, // todo: we might need a better way to store this. - // we are likely going to want to store vms even if we don't know their actorid (ex: actor hasn't been started or is shutdown) - // but we also might want ot be able to store them without a ulid, possibly + // we 100% need a way to store vms even if we don't know their actorid (ex: actor hasn't been started or is shutdown) + // we also might want to be able to store them without a ulid, possibly // so we might need a vec of vms and then to just store maps/indexes of actorid and ulid to vector index // and then like a freelist or something. // i dont really love that option either though cause it feels overkill. // maybe we sure just be using a proper database entirely? // idk. will figure it out later. // - // new related problem: i just realized vmids/ulids and actorids dont have to be unique. + // new related problem: i just realized vmid, actorid pairs dont have to be unique. // if a vm is migrating from one actor to another, there might be two actors with the same vmid. + // + // additional context (05/05/2026): we almost may need a way to store them without a ulid, due to how CH migration works. + // the question becomes if we want to abstract CH migration away entirely from the scheduler. + // we could also possibly ignore it for the non-HA scheduler. + // I (caleb) want to ask cappy (and possibly Lea) about these problems. + // + // the best solution for at least some of this is almost certainly having an external reliable DB (such as etcd) to store some of these things permanently. + // we will need that specifically for what VMs are supposed to be running, because if a large percentage of the cluster goes down, including the manager, we need a way to recover. + // and i dont think leaving that on dashboard which could have high latency is a good idea. + // alternatively we could have the other manager nodes try to keep track of that data, but i think we are going to run into issues with keeping the state consistent between all nodes. + // we may need to make some architecture designs about db consistency vs uptime vs speed in that situation, and im not doing that on my own. pub vm_actorid_ulid_map: Arc>, pub vm_data_cache: Arc>, pub vm_keepalive_tasks: Arc>>, @@ -128,7 +144,7 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(metadata) = actor_ref.ask(&GetVMInfo {vmid: None}).await { + if let Ok(metadata) = actor_ref.ask(&GetVMInfo {vmid: None}).await { // todo: replace with kameo stream let vmid = metadata.vmid; vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that @@ -194,13 +210,15 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(metadata) = actor_ref.ask(&GetAgentStatus).await { + if let Ok(metadata) = actor_ref.ask(&GetAgentStatus).await { // todo: replace with kameo stream if data_cache.contains_key(&actor_ref.id()) { data_cache.alter( &actor_ref.id(), |_, mut v| { v.metadata = metadata; + v.extended_vm_set.extend(v.metadata.vms.iter()); + v } ); @@ -210,7 +228,7 @@ impl SchedulerActor { CachedAgentActor { actor_ref: actor_ref.clone(), metadata, - extended_vm_list: vec![] + extended_vm_set: AHashSet::new() } ); } @@ -261,74 +279,143 @@ impl SchedulerActor { } }) ); - } - - /// current scheduling algo info: - /// this is vaguely based on https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ - /// when a vm is attempted to be scheduled, we loop through every agent and score it based on some rules - /// there are hard rules that will simply throw out an agent entirely. - /// otherwise, we take whatever the best agent we can find is. - async fn schedule_agent( + /// Determine the best agent to schedule a specific VM creation request to. + /// + /// Rough explanation of the algorithm: + /// Loop through every known agent. + /// Go through a set of rules to determine if the VM can be scheduled on this agent at all, and an affinity score and a general score. + /// + /// Based on these scores, pick the best agent. + /// First the affinity score is used, because these are things the customer specifically wanted. + /// If the affinity score is tied, we use the general score as a tie breaker. + /// The general score uses things like resource utilization to not over load any specific agent. + /// + /// + /// Affinity rules are roughly based on https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ + /// + /// todo: + /// - the cache likely needs to be updated automatically when a new vm is scheduled for info like used resources, because otherwise we have to deal with latency on that data we are using + /// and then if someone tries to schedule lets say 10 VMs in a batch, we could end up scheduling them all to the same agent because the metadata hasn't updated. + /// - there are a few solutions for this but they all kinda suck, mostly due to also making sure we deal with latency properly. I am ignoring the issue for now. + fn schedule_agent( &mut self, msg: &CreateVM ) -> Result, Report> { - let mut best_agent = None; - let mut best_agent_score = 0.0f32; + // todo: this could likely be better idiomatic rust. + // I suspect there is a map-reduce operation that does the exact scoring thing I am trying to do. + // I also assume there is a better function for the and_then + let best_agent = self.agent_data_cache.iter() + .map(|agent| (agent.actor_ref.clone(), score_agent(msg, &agent))) + .reduce(|best, new| if new.1 > best.1 { new } else { best }) + .and_then(|best| if best.1 == AgentScore::REJECTED { None } else { Some(best.0) }); - // todo: this arguably could be done as map-reduce. is that better? - let span = info_span!("schedule_agent"); - span.in_scope(|| { - for agent in self.agent_data_cache.iter() { - let mut agent_score = 0.0f32; + best_agent.ok_or_eyre("No valid agents found.") + } +} - let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; - // todo: do we care about VMData.max_vcpus? - let agent_used_vcpus = agent.metadata.used_vcpus + msg.config.data.vcpus; +#[derive(Debug, Clone, Copy, PartialEq)] +struct AgentScore { + general: f32, + affinity: i64 +} - if agent_used_vcpus >= agent_max_vcpus { - continue; - } +impl AgentScore { + pub const REJECTED: Self = Self { + general: f32::NEG_INFINITY, + affinity: i64::MIN + }; +} - agent_score += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; +impl Default for AgentScore { + fn default() -> Self { + Self { + general: 0.0, + affinity: 0 + } + } +} +impl PartialOrd for AgentScore { + fn partial_cmp(&self, other: &Self) -> Option { + let affinity_cmp = self.affinity.cmp(&other.affinity); - // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. - let agent_max_ram = agent.metadata.ram; - let agent_used_ram = agent.metadata.used_ram + msg.config.data.memory; + if affinity_cmp != Ordering::Equal { + return Some(affinity_cmp); + } - if agent_used_ram >= agent_max_ram { - continue; - } + self.general.partial_cmp(&other.general) + } +} + + +fn score_agent( + msg: &CreateVM, + agent: &RefMulti<'_, ActorId, CachedAgentActor> +) -> AgentScore { + let mut score = AgentScore::default(); - agent_score += (agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; + let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; + // todo: do we care about VMData.max_vcpus? + let agent_used_vcpus = agent.metadata.used_vcpus + msg.config.data.vcpus; + if agent_used_vcpus >= agent_max_vcpus { + return AgentScore::REJECTED; + } + + score.general += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; - // todo: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ + // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. + let agent_max_ram = agent.metadata.ram; + let agent_used_ram = agent.metadata.used_ram + msg.config.data.memory; + if agent_used_ram >= agent_max_ram { + return AgentScore::REJECTED; + } - // todo (future): possibly keep a percent of agents completely empty, to be able to be converted to dedis automatically. - // they would have their agent score set to like f32::MIN, so they can be scheduled to if there is no other avaliable agents. - // rough pseudo code to implement this: - // if agent.metadata.vms.len() == 0 && hash(agent.config.hostname) % total_chance < threshold { - // agent_score = 1; - // } + score.general += (agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; + // todo: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ - info!(agent=?agent.value(), score=agent_score); - if agent_score > best_agent_score { - best_agent = Some(agent.actor_ref.clone()); - best_agent_score = agent_score; + if let Some(affinity_rules) = &msg.config.affinity { + for rule in affinity_rules { + let mut meets_requirements = false; + + for requirement in &rule.requirements { + let requirement_outcome = true; // todo: set this based on requirement + + if requirement_outcome { + meets_requirements = true; + break; } } - }); - best_agent.ok_or_eyre("No valid agents found.") + let follows_rule = meets_requirements ^ rule.inverse; + + match (rule.strictness, follows_rule) { + (AffinityStrictness::Required, false) => return AgentScore::REJECTED, + (AffinityStrictness::Required, true) => {}, // specifically do nothing + (AffinityStrictness::Preferred { weight }, follows_rule) => { + score.affinity += follows_rule as i64 * weight; + }, + } + } } + + + + // todo (future): possibly keep a percent of agents completely empty, to be able to be converted to dedis automatically. + // they would have their agent score set to like f32::MIN, so they can be scheduled to if there is no other available agents. + // rough pseudo code to implement this: + // if agent.metadata.vms.len() == 0 && hash(agent.config.hostname) % total_chance < threshold { + // agent_score = 1; + // } + + score } @@ -376,10 +463,8 @@ impl Actor for SchedulerActor { keepalive_task.abort(); }; - self.agent_data_cache.remove(&actor_id); - // todo: attempt vm migration or restart or whatever on agent death. if let Some((_, keepalive_task)) = self.vm_keepalive_tasks.remove(&actor_id) { trace!(?actor_id, "Aborting vm keepalive task"); @@ -387,13 +472,24 @@ impl Actor for SchedulerActor { }; if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&actor_id) { + // todo: this solution definitely isn't optimal. we likely should be storing which agent actor, this specific actor id is scheduled on and specifically removing it from that one. + // especially because things get iffy with migration, but im ignoring that for the moment. + for mut agent in self.agent_data_cache.iter_mut() { + agent.extended_vm_set.remove(&vmid); + } + // todo: we likely should keep a copy of the VirtualMachine manifest in the cache. // instead of removing the vm entirely, we should just modify the status to shutdown or crashed or something. + // + // another potential issue is that the link dying doesn't guarantee that the vm is dead. if there is a networking partition, things get iffy. trace!(?actor_id, ?vmid, "Removing vm from vm_data_cache"); self.vm_data_cache.remove(&vmid); } + // todo: attempt vm restarts if necessary. + + Ok(ControlFlow::Continue(())) } } @@ -405,40 +501,40 @@ impl Message for SchedulerActor { type Reply = Result; async fn handle(&mut self, msg: CreateVM, _ctx: &mut Context) -> Self::Reply { - let target_agent = self.schedule_agent(&msg).await?; + let target_agent = self.schedule_agent(&msg)?; // we add to cache first, because we want to make sure future requests assume this vm exists. if the message fails, we clean it up afterward. if let Some(mut cached_data) = self.agent_data_cache.get_mut(&target_agent.id()) { - cached_data.extended_vm_list.push(msg.vmid); + cached_data.extended_vm_set.insert(msg.vmid); } else { return Err(eyre!("target agent is not in data cache")) } - + self.vm_data_cache.insert( msg.vmid, - CachedVMActor { - actor_ref: None, - metadata: GetVMInfoReply { - vmid: msg.vmid, + CachedVMActor { + actor_ref: None, + metadata: GetVMInfoReply { + vmid: msg.vmid, config: Some(msg.config.clone()) } }); - - + + let reply = target_agent.ask(&msg).await; - + // remove from caches if we fail to schedule if reply.is_err() { self.agent_data_cache.alter(&target_agent.id(), |_,mut v| { - v.extended_vm_list.retain(|&vmid| vmid != msg.vmid); - + v.extended_vm_set.remove(&msg.vmid); + v }); self.vm_data_cache.remove(&msg.vmid); } - + Ok(reply?) } @@ -454,9 +550,8 @@ impl Message for SchedulerActor { ) -> Self::Reply { let vm = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)).await?; tracing::trace!(?vm, "DeleteVM"); - if let Some(vm) = vm { + if let Some(vm) = vm { // don't update cache, because we rely on link dying and updater task to remove from cache once the VM is fully down. vm.tell(&msg).send()?; - // todo: update cache Ok(DeleteVMReply) } else { Err(eyre!("VM not found")) @@ -474,11 +569,8 @@ impl Message for SchedulerActor { ) -> Self::Reply { let vm = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)).await?; tracing::trace!(?vm, "ShutdownVM"); - if let Some(vm) = vm { + if let Some(vm) = vm { // don't update cache, because we rely on link dying and updater task to remove from cache once the VM is fully down. vm.tell(&msg).send()?; - - // todo: update cache - Ok(ShutdownVMReply) } else { Err(eyre!("VM not found")) diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index b526cf8..902708a 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -175,12 +175,14 @@ pub struct VirtualMachine { pub struct AffinityRule { pub strictness: AffinityStrictness, pub affinity_type: AffinityType, - pub direction: AffinityDirection, + /// if true, the outcome of the requirements is inverted + #[serde(default)] + pub inverse: bool, /// ORed together pub requirements: Vec } -#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] +#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone, Copy)] pub enum AffinityStrictness { Required, Preferred { weight: i64 } @@ -192,12 +194,6 @@ pub enum AffinityType { Agent } -#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] -pub enum AffinityDirection { - Normal, - Anti -} - #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub struct AffinityRequirement { pub key: String, From 11186cd562d7ea659623b6f550e77f7fb9329d67 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Tue, 5 May 2026 15:43:11 -0500 Subject: [PATCH 06/35] more work on scheduler affinity stuff. still not finished. Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 154 ++++++++++++++------------ odorobo/src/config.rs | 7 +- odorobo/src/types.rs | 5 +- 3 files changed, 89 insertions(+), 77 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 9ed9536..91a373f 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -12,6 +12,7 @@ use ulid::Ulid; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; use crate::types::AffinityStrictness; +use crate::types::AffinityType; use crate::utils::actor_names::VM; use crate::messages::vm::*; use crate::messages::agent::*; @@ -144,7 +145,7 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(metadata) = actor_ref.ask(&GetVMInfo {vmid: None}).await { // todo: replace with kameo stream + if let Ok(metadata) = actor_ref.ask(&GetVMInfo {vmid: None}).await { // todo: replace with kameo stream, (only send full data once, then send changes from there on) let vmid = metadata.vmid; vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that @@ -307,12 +308,93 @@ impl SchedulerActor { // I suspect there is a map-reduce operation that does the exact scoring thing I am trying to do. // I also assume there is a better function for the and_then let best_agent = self.agent_data_cache.iter() - .map(|agent| (agent.actor_ref.clone(), score_agent(msg, &agent))) + .map(|agent| (agent.actor_ref.clone(), self.score_agent(msg, &agent))) .reduce(|best, new| if new.1 > best.1 { new } else { best }) .and_then(|best| if best.1 == AgentScore::REJECTED { None } else { Some(best.0) }); best_agent.ok_or_eyre("No valid agents found.") } + + // this function intentionally only checks against the cache. this has some positives and negatives: + // positive: it will never trigger any network requests so its very fast, and having to do network requests for scoring whenever we want to schedule a vm is likely a bad idea + // negative: it technically has a delayed view of the cluster, meaning that some things that happened in the future, may not exist yet. so we need to be careful about how this is done so affinity rules are not accidentally broken. mostly this means, if we do anything that could affect the outcome of an affinity rule (ex: network request to an agent), we need to update the cache, before we do the action. + fn score_agent( + &self, + msg: &CreateVM, + agent: &RefMulti<'_, ActorId, CachedAgentActor> + ) -> AgentScore { + let mut score = AgentScore::default(); + + let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; + // todo: do we care about VMData.max_vcpus? + let agent_used_vcpus = agent.metadata.used_vcpus + msg.config.data.vcpus; + + if agent_used_vcpus >= agent_max_vcpus { + return AgentScore::REJECTED; + } + + score.general += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; + + + // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. + let agent_max_ram = agent.metadata.ram; + let agent_used_ram = agent.metadata.used_ram + msg.config.data.memory; + + if agent_used_ram >= agent_max_ram { + return AgentScore::REJECTED; + } + + score.general += (agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; + + + // todo: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ + + + if let Some(affinity_rules) = &msg.config.affinity { + for rule in affinity_rules { + let mut meets_requirements = false; + + let lhs_values: Vec = Vec::with_capacity(1); + + match &rule.affinity_type { + AffinityType::VirtualMachine(zone) => todo!(), + AffinityType::Agent => lhs_values.push(agent.metadata.), // todo: this should be object metadata, but I just realized that field isnt included in the data we have in the cache + } + + // in theory the next bit of this could would loop through all the lhs values and all the requirements and then actually do the computation, but I got busy with other work. + + for requirement in &rule.requirements { + let requirement_outcome = true; // todo: set this based on requirement + + if requirement_outcome { + meets_requirements = true; + break; + } + } + + let follows_rule = meets_requirements ^ rule.inverse; + + match (rule.strictness, follows_rule) { + (AffinityStrictness::Required, false) => return AgentScore::REJECTED, + (AffinityStrictness::Required, true) => {}, // specifically do nothing + (AffinityStrictness::Preferred { weight }, follows_rule) => { + score.affinity += follows_rule as i64 * weight; + }, + } + } + } + + + + // todo (future): possibly keep a percent of agents completely empty, to be able to be converted to dedis automatically. + // they would have their agent score set to like f32::MIN, so they can be scheduled to if there is no other available agents. + // rough pseudo code to implement this: + // if agent.metadata.vms.len() == 0 && hash(agent.config.hostname) % total_chance < threshold { + // agent_score = 1; + // } + + score + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -350,74 +432,6 @@ impl PartialOrd for AgentScore { } -fn score_agent( - msg: &CreateVM, - agent: &RefMulti<'_, ActorId, CachedAgentActor> -) -> AgentScore { - let mut score = AgentScore::default(); - - let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; - // todo: do we care about VMData.max_vcpus? - let agent_used_vcpus = agent.metadata.used_vcpus + msg.config.data.vcpus; - - if agent_used_vcpus >= agent_max_vcpus { - return AgentScore::REJECTED; - } - - score.general += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; - - - // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. - let agent_max_ram = agent.metadata.ram; - let agent_used_ram = agent.metadata.used_ram + msg.config.data.memory; - - if agent_used_ram >= agent_max_ram { - return AgentScore::REJECTED; - } - - score.general += (agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; - - - // todo: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ - - - if let Some(affinity_rules) = &msg.config.affinity { - for rule in affinity_rules { - let mut meets_requirements = false; - - for requirement in &rule.requirements { - let requirement_outcome = true; // todo: set this based on requirement - - if requirement_outcome { - meets_requirements = true; - break; - } - } - - let follows_rule = meets_requirements ^ rule.inverse; - - match (rule.strictness, follows_rule) { - (AffinityStrictness::Required, false) => return AgentScore::REJECTED, - (AffinityStrictness::Required, true) => {}, // specifically do nothing - (AffinityStrictness::Preferred { weight }, follows_rule) => { - score.affinity += follows_rule as i64 * weight; - }, - } - } - } - - - - // todo (future): possibly keep a percent of agents completely empty, to be able to be converted to dedis automatically. - // they would have their agent score set to like f32::MIN, so they can be scheduled to if there is no other available agents. - // rough pseudo code to implement this: - // if agent.metadata.vms.len() == 0 && hash(agent.config.hostname) % total_chance < threshold { - // agent_score = 1; - // } - - score -} - impl Actor for SchedulerActor { diff --git a/odorobo/src/config.rs b/odorobo/src/config.rs index 8b261ea..acb5287 100644 --- a/odorobo/src/config.rs +++ b/odorobo/src/config.rs @@ -144,13 +144,10 @@ pub struct Config { /// The number of VCPUs reserved for the agent. Defaults to 2. #[serde(default = "default_reserved_vcpus")] pub reserved_vcpus: u32, - /// this is just arbitrary data that will be shown but does no config - /// Arbitrary labels that can be used + + /// Arbitrary data that the infra team can set for notes for themselves. odorobo does not directly use these, but does include them in that can be used #[serde(default)] pub labels: AHashMap, - /// Arbitrary annotations that can be used - #[serde(default)] - pub annotations: AHashMap, #[serde(default)] pub network: NetworkConfig, diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index 902708a..95e53bb 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -190,10 +190,12 @@ pub enum AffinityStrictness { #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub enum AffinityType { - VirtualMachine, + VirtualMachine(Zone), Agent } +pub type Zone = String; + #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub struct AffinityRequirement { pub key: String, @@ -208,7 +210,6 @@ pub enum MetadataTable { Annotation } -// todo: possibly replace with std::ops #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub enum Operator { In, From 2202c17db1da66401a728c5adc11c77ddb0525dc Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Thu, 14 May 2026 17:36:09 -0500 Subject: [PATCH 07/35] theoretically write scheduler affinity scoring. need to test this. Signed-off-by: Caleb Jones --- odorobo/src/actors/agent_actor.rs | 3 +- odorobo/src/actors/scheduler_actor.rs | 135 ++++++++++++++++++++------ odorobo/src/ch_driver/actor.rs | 1 + odorobo/src/messages/agent.rs | 5 +- odorobo/src/types.rs | 11 ++- 5 files changed, 116 insertions(+), 39 deletions(-) diff --git a/odorobo/src/actors/agent_actor.rs b/odorobo/src/actors/agent_actor.rs index 3745c2d..6fd7b30 100644 --- a/odorobo/src/actors/agent_actor.rs +++ b/odorobo/src/actors/agent_actor.rs @@ -299,7 +299,8 @@ impl Message for AgentActor { ram: self.memory, vms: self.vms.keys().copied().collect(), used_vcpus: vcpus_used_by_vms + self.config.reserved_vcpus, - used_ram: ByteSize::b(ram_used_by_vms) + used_ram: ByteSize::b(ram_used_by_vms), + metadata: self.metadata.clone() } } } diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 91a373f..72ace12 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,3 +1,4 @@ +use std::iter; use std::ops::ControlFlow; use std::sync::Arc; use std::time::Duration; @@ -11,8 +12,13 @@ use tracing::trace; use ulid::Ulid; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; +use crate::types::AffinityRequirement; +use crate::types::AffinityRule; use crate::types::AffinityStrictness; use crate::types::AffinityType; +use crate::types::MetadataTable; +use crate::types::ObjectMetadata; +use crate::types::Operator; use crate::utils::actor_names::VM; use crate::messages::vm::*; use crate::messages::agent::*; @@ -30,7 +36,7 @@ use tokio::task::JoinHandle; #[derive(Debug, Clone)] pub struct CachedAgentActor { pub actor_ref: RemoteActorRef, - pub metadata: AgentStatus, + pub data: AgentStatus, /// this is a set of all VMs that may be on an agent. it is used for rules such as affinity to make sure we don't schedule things in ways that arent allowed /// We don't know for a fact these VMs are scheduled due to latency and boot up delay, but they may be scheduled. pub extended_vm_set: AHashSet, @@ -39,12 +45,20 @@ pub struct CachedAgentActor { #[derive(Debug, Clone)] pub struct CachedVMActor { pub actor_ref: Option>, - pub metadata: GetVMInfoReply, + pub data: GetVMInfoReply, } // todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. // especially when we haven't fully made decisions about some other things. +// +// todo: we should improve the cache to not have agents and vms send the full data on every update. +// I looked at kameo streams to make this better, but they aren't really intended for this kind of long term update use case. +// They use rust futures::stream which seems to be more intended for you have an iterator for example that will create data, but not like full on sending messages. +// This could likely be done pretty easily by having two get data messages. +// Option 1: One that creates a session and sends the full data and then only sends diffs after that. +// Option 2: we can have one that sends full data, and then another that only sends data that changes. +// Option 2 is easier to write and uses less compute, but uses more network bandwidth. #[derive(RemoteActor)] pub struct SchedulerActor { pub agent_data_cache: Arc>, @@ -96,7 +110,7 @@ impl SchedulerActor { &mut self, hostname: &str, ) -> Option> { - self.agent_data_cache.iter().find(|data| data.metadata.hostname == hostname).map(|data| data.actor_ref.clone()) + self.agent_data_cache.iter().find(|data| data.data.hostname == hostname).map(|data| data.actor_ref.clone()) } @@ -145,8 +159,8 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(metadata) = actor_ref.ask(&GetVMInfo {vmid: None}).await { // todo: replace with kameo stream, (only send full data once, then send changes from there on) - let vmid = metadata.vmid; + if let Ok(data) = actor_ref.ask(&GetVMInfo {vmid: None}).await { + let vmid = data.vmid; vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that @@ -154,7 +168,7 @@ impl SchedulerActor { vmid, CachedVMActor { actor_ref: Some(actor_ref.clone()), - metadata: metadata + data: data } ); @@ -211,14 +225,14 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(metadata) = actor_ref.ask(&GetAgentStatus).await { // todo: replace with kameo stream + if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { // todo: replace with kameo stream if data_cache.contains_key(&actor_ref.id()) { data_cache.alter( &actor_ref.id(), |_, mut v| { - v.metadata = metadata; + v.data = data; - v.extended_vm_set.extend(v.metadata.vms.iter()); + v.extended_vm_set.extend(v.data.vms.iter()); v } @@ -228,7 +242,7 @@ impl SchedulerActor { actor_ref.id(), CachedAgentActor { actor_ref: actor_ref.clone(), - metadata, + data: data, extended_vm_set: AHashSet::new() } ); @@ -325,9 +339,9 @@ impl SchedulerActor { ) -> AgentScore { let mut score = AgentScore::default(); - let agent_max_vcpus = agent.metadata.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; + let agent_max_vcpus = agent.data.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; // todo: do we care about VMData.max_vcpus? - let agent_used_vcpus = agent.metadata.used_vcpus + msg.config.data.vcpus; + let agent_used_vcpus = agent.data.used_vcpus + msg.config.data.vcpus; if agent_used_vcpus >= agent_max_vcpus { return AgentScore::REJECTED; @@ -337,42 +351,66 @@ impl SchedulerActor { // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. - let agent_max_ram = agent.metadata.ram; - let agent_used_ram = agent.metadata.used_ram + msg.config.data.memory; + let agent_max_ram = agent.data.ram; + let agent_used_ram = agent.data.used_ram + msg.config.data.memory; if agent_used_ram >= agent_max_ram { return AgentScore::REJECTED; } - score.general += (agent_max_ram.as_u64() - agent.metadata.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; - - - // todo: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ + score.general += (agent_max_ram.as_u64() - agent.data.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; + // roughly based on: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ if let Some(affinity_rules) = &msg.config.affinity { for rule in affinity_rules { - let mut meets_requirements = false; + let mut metadata_tables: Vec = Vec::with_capacity(1); + + match rule.affinity_type { + AffinityType::VirtualMachine => { + for vmid in &agent.extended_vm_set { + let Some(vm_data_cache_ref) = &self.vm_data_cache.get(vmid) else { + continue; + }; + + let Some(vm_manifest) = &vm_data_cache_ref.data.config else { + continue; + }; + + if let Some(metadata) = &vm_manifest.metadata { + metadata_tables.push(metadata.clone()); + } + } + }, + AffinityType::Agent => metadata_tables.push(agent.data.metadata.clone()), + }; - let lhs_values: Vec = Vec::with_capacity(1); + let mut follows_rule = false; - match &rule.affinity_type { - AffinityType::VirtualMachine(zone) => todo!(), - AffinityType::Agent => lhs_values.push(agent.metadata.), // todo: this should be object metadata, but I just realized that field isnt included in the data we have in the cache - } + for requirement in &rule.requirements { + let mut requirement_outcome = true; - // in theory the next bit of this could would loop through all the lhs values and all the requirements and then actually do the computation, but I got busy with other work. + for object_metadata in &metadata_tables { + let table = match requirement.table { + MetadataTable::Label => &object_metadata.labels, + MetadataTable::Annotation => &object_metadata.annotations, + }; - for requirement in &rule.requirements { - let requirement_outcome = true; // todo: set this based on requirement + let value_option = table.get(&requirement.key); + + if !evaluate_table_value(&value_option, requirement) { + requirement_outcome = false; + break; + } + } if requirement_outcome { - meets_requirements = true; - break; + follows_rule = true; + break } } - let follows_rule = meets_requirements ^ rule.inverse; + follows_rule ^= rule.inverse; match (rule.strictness, follows_rule) { (AffinityStrictness::Required, false) => return AgentScore::REJECTED, @@ -397,6 +435,39 @@ impl SchedulerActor { } } +fn evaluate_table_value( + value_option: &Option<&String>, + requirement: &AffinityRequirement +) -> bool { + let Some(value) = *value_option else { + return false; + }; + + match requirement.operator { + Operator::In => return requirement.values.iter().any(|e| *e == *value), + Operator::NotIn => return !requirement.values.iter().any(|e| *e == *value), + Operator::Lt | Operator::Gt => { + if requirement.values.len() != 1 { + return false; + } + + let Ok(value_number): Result = value.parse() else { + return false; + }; + + let Ok(requirement_value_number): Result = requirement.values[0].parse() else { + return false; + }; + + if requirement.operator == Operator::Lt { + return value_number < requirement_value_number; + } else { + return value_number > requirement_value_number; + } + }, + }; +} + #[derive(Debug, Clone, Copy, PartialEq)] struct AgentScore { general: f32, @@ -528,7 +599,7 @@ impl Message for SchedulerActor { msg.vmid, CachedVMActor { actor_ref: None, - metadata: GetVMInfoReply { + data: GetVMInfoReply { vmid: msg.vmid, config: Some(msg.config.clone()) } @@ -606,7 +677,7 @@ impl Message for SchedulerActor { let mut vms = Vec::new(); for agent in self.agent_data_cache.iter() { - vms.extend_from_slice(agent.metadata.vms.as_slice()); + vms.extend_from_slice(agent.data.vms.as_slice()); } Ok(AgentListVMsReply { vms }) diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index b3a9f92..8b82146 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -164,6 +164,7 @@ impl Message for VMActor { } } + #[remote_message] impl Message for VMActor { type Reply = MigrateVMReceiveReply; diff --git a/odorobo/src/messages/agent.rs b/odorobo/src/messages/agent.rs index 581761e..67a1c91 100644 --- a/odorobo/src/messages/agent.rs +++ b/odorobo/src/messages/agent.rs @@ -3,6 +3,8 @@ use kameo::Reply; use serde::{Deserialize, Serialize}; use ulid::Ulid; +use crate::types::ObjectMetadata; + #[derive(Serialize, Deserialize)] pub struct GetAgentStatus; @@ -16,5 +18,6 @@ pub struct AgentStatus { pub ram: ByteSize, pub used_vcpus: u32, pub used_ram: ByteSize, - pub vms: Vec + pub vms: Vec, + pub metadata: ObjectMetadata } diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index 95e53bb..2670ff1 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -166,7 +166,8 @@ pub struct VirtualMachine { /// Metadata pub metadata: Option, - /// List of Affinity rules for scheduling. These are ANDed / summed together depending on the strictness. + /// List of Affinity rules for scheduling. + /// Affinity rules are ANDed or summed based on strictness pub affinity: Option> } @@ -190,12 +191,12 @@ pub enum AffinityStrictness { #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub enum AffinityType { - VirtualMachine(Zone), + VirtualMachine, Agent } -pub type Zone = String; - +/// if there are several metadata tables, their results will be ANDed together +/// EX: if a requirement is checked against several VMs, it must pass all VMs for the requirement to be fulfilled. #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub struct AffinityRequirement { pub key: String, @@ -210,7 +211,7 @@ pub enum MetadataTable { Annotation } -#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] +#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone, PartialEq)] pub enum Operator { In, NotIn, From fedb58bb54b71f5b5c7b2f5b9aff0abcf237b422 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Mon, 6 Jul 2026 16:50:31 -0500 Subject: [PATCH 08/35] remove ipam actor Signed-off-by: Caleb Jones --- odorobo/src/actors/ip_management_actor.rs | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/odorobo/src/actors/ip_management_actor.rs b/odorobo/src/actors/ip_management_actor.rs index 9796b38..fd21d1a 100644 --- a/odorobo/src/actors/ip_management_actor.rs +++ b/odorobo/src/actors/ip_management_actor.rs @@ -21,18 +21,3 @@ fn test_calculate_mac_address() { let mac = calculate_mac_address(ip); assert_eq!(mac, [0x46, 0x59, 0x52, 168, 0x01, 0x01]); } - -/// HTTP REST API service -#[derive(RemoteActor)] -pub struct IPManagementActor; - -impl Actor for IPManagementActor { - type Args = (); - type Error = Report; - - async fn on_start(_state: Self::Args, _actor_ref: ActorRef) -> Result { - // if we need to like prep the router stuff - - Ok(Self) - } -} From 29c2a84cabe35a41f85f53a0a1dcb18b6eb9a847 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Mon, 6 Jul 2026 17:24:16 -0500 Subject: [PATCH 09/35] fix bug with getting all agent and vm actors in actor finders Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 72ace12..da92142 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -121,8 +121,11 @@ impl SchedulerActor { data_cache: Arc>, keepalive_tasks: Arc>> ) -> Result<(), Report> { + trace!("running vm_actor_finder"); - while let Some(vm_actor) = RemoteActorRef::::lookup_all(VM).try_next().await? { + let mut vm_actor_stream = RemoteActorRef::::lookup_all(VM); + + while let Some(vm_actor) = vm_actor_stream.try_next().await? { if !keepalive_tasks.contains_key(&vm_actor.id()) { trace!(?vm_actor, "starting vm_updater_task"); @@ -191,8 +194,11 @@ impl SchedulerActor { data_cache: Arc>, keepalive_tasks: Arc>>, ) -> Result<(), Report> { - info!("running agent_actor_finder"); - while let Some(agent_actor) = RemoteActorRef::::lookup_all(AGENT).try_next().await? { + trace!("running agent_actor_finder"); + + let mut agent_actor_stream = RemoteActorRef::::lookup_all(AGENT); + + while let Some(agent_actor) = agent_actor_stream.try_next().await? { if !keepalive_tasks.contains_key(&agent_actor.id()) { trace!(?agent_actor, "starting agent_updater_task"); From 2fde501a6b293cce79a1bd48431aba3c1a6c6913 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Tue, 7 Jul 2026 21:54:10 -0500 Subject: [PATCH 10/35] commiting stuff before cypress rebase. was in the middle of testing scheduler affinity Signed-off-by: Caleb Jones --- odorobo/src/actors/scheduler_actor.rs | 5 ++--- odoroboctl/src/cli.rs | 25 ++++++++++++++++++++++++- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index da92142..0f15c00 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,4 +1,3 @@ -use std::iter; use std::ops::ControlFlow; use std::sync::Arc; use std::time::Duration; @@ -13,7 +12,6 @@ use ulid::Ulid; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; use crate::types::AffinityRequirement; -use crate::types::AffinityRule; use crate::types::AffinityStrictness; use crate::types::AffinityType; use crate::types::MetadataTable; @@ -27,7 +25,6 @@ use crate::utils::actor_names::AGENT; use crate::utils::actor_names::vm_actor_id; use stable_eyre::eyre::OptionExt; use stable_eyre::{Report, Result, eyre::eyre}; -use tracing::info_span; use tracing::{info, warn}; use dashmap::DashMap; use tokio::task::JoinHandle; @@ -437,6 +434,8 @@ impl SchedulerActor { // agent_score = 1; // } + info!(?score, agent_id=?(agent.key()), "scored agent"); + score } } diff --git a/odoroboctl/src/cli.rs b/odoroboctl/src/cli.rs index b769271..69c530e 100644 --- a/odoroboctl/src/cli.rs +++ b/odoroboctl/src/cli.rs @@ -1,8 +1,10 @@ +use std::collections::BTreeMap; + use clap::{Parser, Subcommand}; use reqwest::{Client, Response}; use serde::{Deserialize}; use stable_eyre::Result; -use odorobo::types::{CreateVMRequest, VMData, VirtualMachine}; +use odorobo::types::{AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, CreateVMRequest, MetadataTable, ObjectMetadata, Operator, VMData, VirtualMachine}; use ulid::Ulid; use bytesize::ByteSize; @@ -107,6 +109,27 @@ pub async fn run_command(cli: Cli) -> Result<()> { image: "/var/lib/odorobo/f43.raw".to_string(), ..Default::default() }, + metadata: Some(ObjectMetadata { + labels: BTreeMap::new(), + annotations: BTreeMap::from([ + (String::from("distribution"), String::from("different")) + ]), + }), + affinity: Some(vec![ + AffinityRule { + strictness: AffinityStrictness::Required, + affinity_type: AffinityType::VirtualMachine, + inverse: false, + requirements: vec![ + AffinityRequirement { + key: String::from("distribution"), + table: MetadataTable::Annotation, + operator: Operator::NotIn, + values: vec![String::from("different")] + } + ] + } + ]), ..Default::default() }; From 674c5a9762fa0ba566235fc8c222a2641c1ed920 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 10 Jul 2026 02:16:41 -0500 Subject: [PATCH 11/35] cargo fmt Signed-off-by: Caleb Jones --- odorobo/src/actors/agent_actor.rs | 2 +- odorobo/src/actors/scheduler_actor.rs | 213 ++++++++++++-------------- odorobo/src/ch_driver/actor.rs | 5 +- odorobo/src/config.rs | 2 +- odorobo/src/messages/agent.rs | 3 +- odorobo/src/types.rs | 2 +- odoroboctl/src/cli.rs | 48 +++--- 7 files changed, 127 insertions(+), 148 deletions(-) diff --git a/odorobo/src/actors/agent_actor.rs b/odorobo/src/actors/agent_actor.rs index 4a8769b..ec53bc2 100644 --- a/odorobo/src/actors/agent_actor.rs +++ b/odorobo/src/actors/agent_actor.rs @@ -316,7 +316,7 @@ impl Message for AgentActor { vms: self.vms.keys().copied().collect(), used_vcpus: vcpus_used_by_vms + self.config.reserved_vcpus, used_ram: ByteSize::b(ram_used_by_vms), - metadata: self.metadata.clone() + metadata: self.metadata.clone(), } } } diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index bb15dd8..8539352 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,34 +1,33 @@ +use std::cmp::Ordering; use std::ops::ControlFlow; use std::sync::Arc; use std::time::Duration; -use std::cmp::Ordering; -use ahash::AHashSet; -use dashmap::mapref::multiple::RefMulti; -use kameo::prelude::*; -use libp2p::futures::TryStreamExt; -use tracing::trace; -use ulid::Ulid; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; +use crate::messages::agent::*; +use crate::messages::vm::*; +use crate::messages::{Ping, Pong}; use crate::types::AffinityRequirement; use crate::types::AffinityStrictness; use crate::types::AffinityType; use crate::types::MetadataTable; use crate::types::ObjectMetadata; use crate::types::Operator; -use crate::utils::actor_names::VM; -use crate::messages::vm::*; -use crate::messages::agent::*; -use crate::messages::{Ping, Pong}; use crate::utils::actor_names::AGENT; +use crate::utils::actor_names::VM; use crate::utils::actor_names::vm_actor_id; +use ahash::AHashSet; +use dashmap::DashMap; +use dashmap::mapref::multiple::RefMulti; +use kameo::prelude::*; +use libp2p::futures::TryStreamExt; use stable_eyre::eyre::OptionExt; use stable_eyre::{Report, Result, eyre::eyre}; -use tracing::{info, warn}; -use dashmap::DashMap; use tokio::task::JoinHandle; - +use tracing::trace; +use tracing::{info, warn}; +use ulid::Ulid; #[derive(Debug, Clone)] pub struct CachedAgentActor { @@ -45,7 +44,6 @@ pub struct CachedVMActor { pub data: GetVMInfoReply, } - // todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. // especially when we haven't fully made decisions about some other things. // @@ -99,23 +97,27 @@ impl SchedulerActor { &mut self, actor_id: &ActorId, ) -> Option> { - self.agent_data_cache.get(actor_id).map(|data| data.actor_ref.clone()) + self.agent_data_cache + .get(actor_id) + .map(|data| data.actor_ref.clone()) } async fn lookup_agent_by_hostname( &mut self, hostname: &str, ) -> Option> { - self.agent_data_cache.iter().find(|data| data.data.hostname == hostname).map(|data| data.actor_ref.clone()) + self.agent_data_cache + .iter() + .find(|data| data.data.hostname == hostname) + .map(|data| data.actor_ref.clone()) } - // someone should likely give caleb a firm talking to about code duplication due to this section, but things are just different enough that trying to make them one function requires usage of a lot of generics which feels even worse. so i dont know what to do. cappy please fix. i hate this. async fn vm_actor_finder( parent_actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, data_cache: Arc>, - keepalive_tasks: Arc>> + keepalive_tasks: Arc>>, ) -> Result<(), Report> { trace!("running vm_actor_finder"); @@ -132,22 +134,15 @@ impl SchedulerActor { let vm_actorid_ulid_map_clone = Arc::clone(&vm_actorid_ulid_map); let data_cache_clone = Arc::clone(&data_cache); let updater_task = tokio::spawn(async move { - Self::vm_updater_task( - vm_actor, - vm_actorid_ulid_map_clone, - data_cache_clone - ).await; + Self::vm_updater_task(vm_actor, vm_actorid_ulid_map_clone, data_cache_clone) + .await; }); - keepalive_tasks.insert( - vm_actor_id, - updater_task - ); + keepalive_tasks.insert(vm_actor_id, updater_task); } } Ok(()) - } async fn vm_updater_task( @@ -158,7 +153,7 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(data) = actor_ref.ask(&GetVMInfo {vmid: None}).await { + if let Ok(data) = actor_ref.ask(&GetVMInfo { vmid: None }).await { let vmid = data.vmid; vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that @@ -167,8 +162,8 @@ impl SchedulerActor { vmid, CachedVMActor { actor_ref: Some(actor_ref.clone()), - data: data - } + data: data, + }, ); fails = 0; @@ -204,16 +199,10 @@ impl SchedulerActor { let data_cache_clone = Arc::clone(&data_cache); let updater_task = tokio::spawn(async move { - Self::agent_updater_task( - agent_actor, - data_cache_clone - ).await; + Self::agent_updater_task(agent_actor, data_cache_clone).await; }); - keepalive_tasks.insert( - agent_actor_id, - updater_task - ); + keepalive_tasks.insert(agent_actor_id, updater_task); } } @@ -227,26 +216,24 @@ impl SchedulerActor { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; loop { - if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { // todo: replace with kameo stream + if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { + // todo: replace with kameo stream if data_cache.contains_key(&actor_ref.id()) { - data_cache.alter( - &actor_ref.id(), - |_, mut v| { - v.data = data; + data_cache.alter(&actor_ref.id(), |_, mut v| { + v.data = data; - v.extended_vm_set.extend(v.data.vms.iter()); + v.extended_vm_set.extend(v.data.vms.iter()); - v - } - ); + v + }); } else { data_cache.insert( actor_ref.id(), CachedAgentActor { actor_ref: actor_ref.clone(), data: data, - extended_vm_set: AHashSet::new() - } + extended_vm_set: AHashSet::new(), + }, ); } @@ -272,30 +259,28 @@ impl SchedulerActor { let vm_data_cache_arc_clone = Arc::clone(&self.vm_data_cache); let vm_keepalive_tasks_arc_clone = Arc::clone(&self.vm_keepalive_tasks); - self.cache_actor_finder = Some( - tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_secs(1)); - loop { - let vm_join_handle = Self::vm_actor_finder( - actor_ref.clone(), - Arc::clone(&vm_actorid_ulid_map_arc_clone), - Arc::clone(&vm_data_cache_arc_clone), - Arc::clone(&vm_keepalive_tasks_arc_clone), - ); + self.cache_actor_finder = Some(tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + loop { + let vm_join_handle = Self::vm_actor_finder( + actor_ref.clone(), + Arc::clone(&vm_actorid_ulid_map_arc_clone), + Arc::clone(&vm_data_cache_arc_clone), + Arc::clone(&vm_keepalive_tasks_arc_clone), + ); - let agent_join_handle = Self::agent_actor_finder( - actor_ref.clone(), - Arc::clone(&agent_data_cache_arc_clone), - Arc::clone(&agent_keepalive_tasks_arc_clone), - ); + let agent_join_handle = Self::agent_actor_finder( + actor_ref.clone(), + Arc::clone(&agent_data_cache_arc_clone), + Arc::clone(&agent_keepalive_tasks_arc_clone), + ); - // intentionally ignoring results because we want to keep finding actors even if an attempt fails - let _ = tokio::join!(vm_join_handle, agent_join_handle); + // intentionally ignoring results because we want to keep finding actors even if an attempt fails + let _ = tokio::join!(vm_join_handle, agent_join_handle); - interval.tick().await; - } - }) - ); + interval.tick().await; + } + })); } /// Determine the best agent to schedule a specific VM creation request to. @@ -316,17 +301,22 @@ impl SchedulerActor { /// - the cache likely needs to be updated automatically when a new vm is scheduled for info like used resources, because otherwise we have to deal with latency on that data we are using /// and then if someone tries to schedule lets say 10 VMs in a batch, we could end up scheduling them all to the same agent because the metadata hasn't updated. /// - there are a few solutions for this but they all kinda suck, mostly due to also making sure we deal with latency properly. I am ignoring the issue for now. - fn schedule_agent( - &mut self, - _msg: &CreateVM, - ) -> Result, Report> { + fn schedule_agent(&mut self, _msg: &CreateVM) -> Result, Report> { // todo: this could likely be better idiomatic rust. // I suspect there is a map-reduce operation that does the exact scoring thing I am trying to do. // I also assume there is a better function for the and_then - let best_agent = self.agent_data_cache.iter() + let best_agent = self + .agent_data_cache + .iter() .map(|agent| (agent.actor_ref.clone(), self.score_agent(msg, &agent))) .reduce(|best, new| if new.1 > best.1 { new } else { best }) - .and_then(|best| if best.1 == AgentScore::REJECTED { None } else { Some(best.0) }); + .and_then(|best| { + if best.1 == AgentScore::REJECTED { + None + } else { + Some(best.0) + } + }); best_agent.ok_or_eyre("No valid agents found.") } @@ -337,11 +327,12 @@ impl SchedulerActor { fn score_agent( &self, msg: &CreateVM, - agent: &RefMulti<'_, ActorId, CachedAgentActor> + agent: &RefMulti<'_, ActorId, CachedAgentActor>, ) -> AgentScore { let mut score = AgentScore::default(); - let agent_max_vcpus = agent.data.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR / VCPU_OVERPROVISIONMENT_DENOMINATOR; + let agent_max_vcpus = agent.data.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR + / VCPU_OVERPROVISIONMENT_DENOMINATOR; // todo: do we care about VMData.max_vcpus? let agent_used_vcpus = agent.data.used_vcpus + msg.config.data.vcpus; @@ -351,7 +342,6 @@ impl SchedulerActor { score.general += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; - // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. let agent_max_ram = agent.data.ram; let agent_used_ram = agent.data.used_ram + msg.config.data.memory; @@ -360,8 +350,8 @@ impl SchedulerActor { return AgentScore::REJECTED; } - score.general += (agent_max_ram.as_u64() - agent.data.used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; - + score.general += (agent_max_ram.as_u64() - agent.data.used_ram.as_u64()) as f32 + / agent_max_ram.as_u64() as f32; // roughly based on: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ if let Some(affinity_rules) = &msg.config.affinity { @@ -383,7 +373,7 @@ impl SchedulerActor { metadata_tables.push(metadata.clone()); } } - }, + } AffinityType::Agent => metadata_tables.push(agent.data.metadata.clone()), }; @@ -408,7 +398,7 @@ impl SchedulerActor { if requirement_outcome { follows_rule = true; - break + break; } } @@ -416,16 +406,14 @@ impl SchedulerActor { match (rule.strictness, follows_rule) { (AffinityStrictness::Required, false) => return AgentScore::REJECTED, - (AffinityStrictness::Required, true) => {}, // specifically do nothing + (AffinityStrictness::Required, true) => {} // specifically do nothing (AffinityStrictness::Preferred { weight }, follows_rule) => { score.affinity += follows_rule as i64 * weight; - }, + } } } } - - // todo (future): possibly keep a percent of agents completely empty, to be able to be converted to dedis automatically. // they would have their agent score set to like f32::MIN, so they can be scheduled to if there is no other available agents. // rough pseudo code to implement this: @@ -439,10 +427,7 @@ impl SchedulerActor { } } -fn evaluate_table_value( - value_option: &Option<&String>, - requirement: &AffinityRequirement -) -> bool { +fn evaluate_table_value(value_option: &Option<&String>, requirement: &AffinityRequirement) -> bool { let Some(value) = *value_option else { return false; }; @@ -468,20 +453,20 @@ fn evaluate_table_value( } else { return value_number > requirement_value_number; } - }, + } }; } #[derive(Debug, Clone, Copy, PartialEq)] struct AgentScore { general: f32, - affinity: i64 + affinity: i64, } impl AgentScore { pub const REJECTED: Self = Self { general: f32::NEG_INFINITY, - affinity: i64::MIN + affinity: i64::MIN, }; } @@ -489,7 +474,7 @@ impl Default for AgentScore { fn default() -> Self { Self { general: 0.0, - affinity: 0 + affinity: 0, } } } @@ -506,9 +491,6 @@ impl PartialOrd for AgentScore { } } - - - impl Actor for SchedulerActor { type Args = (); type Error = Report; @@ -545,7 +527,6 @@ impl Actor for SchedulerActor { return Ok(ControlFlow::Break(ActorStopReason::Killed)); }; - if let Some((_, keepalive_task)) = self.agent_keepalive_tasks.remove(&actor_id) { trace!(?actor_id, "Aborting agent keepalive task"); keepalive_task.abort(); @@ -553,7 +534,6 @@ impl Actor for SchedulerActor { self.agent_data_cache.remove(&actor_id); - if let Some((_, keepalive_task)) = self.vm_keepalive_tasks.remove(&actor_id) { trace!(?actor_id, "Aborting vm keepalive task"); keepalive_task.abort(); @@ -574,10 +554,8 @@ impl Actor for SchedulerActor { self.vm_data_cache.remove(&vmid); } - // todo: attempt vm restarts if necessary. - Ok(ControlFlow::Continue(())) } } @@ -585,14 +563,18 @@ impl Actor for SchedulerActor { impl Message for SchedulerActor { type Reply = Result; - async fn handle(&mut self, msg: CreateVM, _ctx: &mut Context) -> Self::Reply { + async fn handle( + &mut self, + msg: CreateVM, + _ctx: &mut Context, + ) -> Self::Reply { let target_agent = self.schedule_agent(&msg)?; // we add to cache first, because we want to make sure future requests assume this vm exists. if the message fails, we clean it up afterward. if let Some(mut cached_data) = self.agent_data_cache.get_mut(&target_agent.id()) { cached_data.extended_vm_set.insert(msg.vmid); } else { - return Err(eyre!("target agent is not in data cache")) + return Err(eyre!("target agent is not in data cache")); } self.vm_data_cache.insert( @@ -601,18 +583,16 @@ impl Message for SchedulerActor { actor_ref: None, data: GetVMInfoReply { vmid: msg.vmid, - config: Some(msg.config.clone()) - } - }); - - + config: Some(msg.config.clone()), + }, + }, + ); let reply = target_agent.ask(&msg).await; - // remove from caches if we fail to schedule if reply.is_err() { - self.agent_data_cache.alter(&target_agent.id(), |_,mut v| { + self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { v.extended_vm_set.remove(&msg.vmid); v @@ -620,7 +600,6 @@ impl Message for SchedulerActor { self.vm_data_cache.remove(&msg.vmid); } - Ok(reply?) } } @@ -635,7 +614,8 @@ impl Message for SchedulerActor { ) -> Self::Reply { let vm = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)).await?; tracing::trace!(?vm, "DeleteVM"); - if let Some(vm) = vm { // don't update cache, because we rely on link dying and updater task to remove from cache once the VM is fully down. + if let Some(vm) = vm { + // don't update cache, because we rely on link dying and updater task to remove from cache once the VM is fully down. vm.tell(&msg).send()?; Ok(DeleteVMReply) } else { @@ -654,7 +634,8 @@ impl Message for SchedulerActor { ) -> Self::Reply { let vm = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)).await?; tracing::trace!(?vm, "ShutdownVM"); - if let Some(vm) = vm { // don't update cache, because we rely on link dying and updater task to remove from cache once the VM is fully down. + if let Some(vm) = vm { + // don't update cache, because we rely on link dying and updater task to remove from cache once the VM is fully down. vm.tell(&msg).send()?; Ok(ShutdownVMReply) } else { diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index 577a103..a29db5b 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -30,7 +30,7 @@ pub struct VMActor { /// path to the Cloud Hypervisor socket, in /run/odorobo/vms//ch.sock pub vm_instance: VMInstance, pub migration_state: Option, - pub manifest: Option + pub manifest: Option, } impl Actor for VMActor { @@ -73,7 +73,7 @@ impl Actor for VMActor { vmid, vm_instance: vminstance, migration_state: None, - manifest: vm_manifest + manifest: vm_manifest, }) } @@ -166,7 +166,6 @@ impl Message for VMActor { } } - #[remote_message] impl Message for VMActor { type Reply = MigrateVMReceiveReply; diff --git a/odorobo/src/config.rs b/odorobo/src/config.rs index e13c9ae..5829928 100644 --- a/odorobo/src/config.rs +++ b/odorobo/src/config.rs @@ -143,7 +143,7 @@ pub struct Config { /// The number of VCPUs reserved for the agent. Defaults to 2. #[serde(default = "default_reserved_vcpus")] pub reserved_vcpus: u32, - + /// Arbitrary data that the infra team can set for notes for themselves. odorobo does not directly use these, but does include them in that can be used #[serde(default)] pub labels: AHashMap, diff --git a/odorobo/src/messages/agent.rs b/odorobo/src/messages/agent.rs index 67a1c91..44f4193 100644 --- a/odorobo/src/messages/agent.rs +++ b/odorobo/src/messages/agent.rs @@ -5,7 +5,6 @@ use ulid::Ulid; use crate::types::ObjectMetadata; - #[derive(Serialize, Deserialize)] pub struct GetAgentStatus; @@ -19,5 +18,5 @@ pub struct AgentStatus { pub used_vcpus: u32, pub used_ram: ByteSize, pub vms: Vec, - pub metadata: ObjectMetadata + pub metadata: ObjectMetadata, } diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index 0f80040..0852983 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -167,7 +167,7 @@ pub struct VirtualMachine { /// List of Affinity rules for scheduling. /// Affinity rules are ANDed or summed based on strictness - pub affinity: Option> + pub affinity: Option>, } #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] diff --git a/odoroboctl/src/cli.rs b/odoroboctl/src/cli.rs index 2e7c1a0..a92f0b6 100644 --- a/odoroboctl/src/cli.rs +++ b/odoroboctl/src/cli.rs @@ -1,17 +1,20 @@ use std::collections::BTreeMap; -use clap::{Parser, Subcommand}; -use reqwest::{Client, Response}; -use serde::{Deserialize}; -use stable_eyre::Result; -use odorobo::types::{AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, CreateVMRequest, MetadataTable, ObjectMetadata, Operator, VMData, VirtualMachine}; -use ulid::Ulid; use bytesize::ByteSize; use clap::{Parser, Subcommand}; +use clap::{Parser, Subcommand}; +use odorobo::types::{ + AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, CreateVMRequest, + MetadataTable, ObjectMetadata, Operator, VMData, VirtualMachine, +}; use odorobo::types::{CreateVMRequest, VMData, VirtualMachine}; use reqwest::{Client, Response}; +use reqwest::{Client, Response}; +use serde::Deserialize; use serde::Deserialize; use stable_eyre::Result; +use stable_eyre::Result; +use ulid::Ulid; use ulid::Ulid; #[derive(Parser)] @@ -118,25 +121,22 @@ pub async fn run_command(cli: Cli) -> Result<()> { }, metadata: Some(ObjectMetadata { labels: BTreeMap::new(), - annotations: BTreeMap::from([ - (String::from("distribution"), String::from("different")) - ]), + annotations: BTreeMap::from([( + String::from("distribution"), + String::from("different"), + )]), }), - affinity: Some(vec![ - AffinityRule { - strictness: AffinityStrictness::Required, - affinity_type: AffinityType::VirtualMachine, - inverse: false, - requirements: vec![ - AffinityRequirement { - key: String::from("distribution"), - table: MetadataTable::Annotation, - operator: Operator::NotIn, - values: vec![String::from("different")] - } - ] - } - ]), + affinity: Some(vec![AffinityRule { + strictness: AffinityStrictness::Required, + affinity_type: AffinityType::VirtualMachine, + inverse: false, + requirements: vec![AffinityRequirement { + key: String::from("distribution"), + table: MetadataTable::Annotation, + operator: Operator::NotIn, + values: vec![String::from("different")], + }], + }]), ..Default::default() }; From 06c78851c82ab77f10dfb042d070c1c13893b0c6 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 10 Jul 2026 02:22:41 -0500 Subject: [PATCH 12/35] fix bugs created by merge commit --- odorobo/src/actors/ip_management_actor.rs | 3 --- odorobo/src/actors/scheduler_actor.rs | 2 +- odorobo/src/ch_driver/actor.rs | 4 ++-- odoroboctl/src/cli.rs | 6 ------ 4 files changed, 3 insertions(+), 12 deletions(-) diff --git a/odorobo/src/actors/ip_management_actor.rs b/odorobo/src/actors/ip_management_actor.rs index fd21d1a..2f1fdd9 100644 --- a/odorobo/src/actors/ip_management_actor.rs +++ b/odorobo/src/actors/ip_management_actor.rs @@ -1,6 +1,3 @@ -use kameo::prelude::*; -use stable_eyre::{Report, Result}; - // idk if we ever agreed upon an OUI for fyra, but im reserving `FYR` for this // -cappy pub const FYRA_OUI: [u8; 3] = [0x46, 0x59, 0x52]; diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 8539352..93965db 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -301,7 +301,7 @@ impl SchedulerActor { /// - the cache likely needs to be updated automatically when a new vm is scheduled for info like used resources, because otherwise we have to deal with latency on that data we are using /// and then if someone tries to schedule lets say 10 VMs in a batch, we could end up scheduling them all to the same agent because the metadata hasn't updated. /// - there are a few solutions for this but they all kinda suck, mostly due to also making sure we deal with latency properly. I am ignoring the issue for now. - fn schedule_agent(&mut self, _msg: &CreateVM) -> Result, Report> { + fn schedule_agent(&mut self, msg: &CreateVM) -> Result, Report> { // todo: this could likely be better idiomatic rust. // I suspect there is a map-reduce operation that does the exact scoring thing I am trying to do. // I also assume there is a better function for the and_then diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index a29db5b..0427aa3 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -41,7 +41,7 @@ impl Actor for VMActor { #[tracing::instrument(skip_all)] async fn on_start((vmid, vm_config): Self::Args, actor_ref: ActorRef) -> Result { let mut vminstance = - VMInstance::spawn(&vmid.to_string(), vm_config.map(VmConfig::from), None).await?; + VMInstance::spawn(&vmid.to_string(), vm_config.clone().map(VmConfig::from), None).await?; // Take the child process out so we can watch for unexpected death. // destroy() handles a missing child_process gracefully. @@ -73,7 +73,7 @@ impl Actor for VMActor { vmid, vm_instance: vminstance, migration_state: None, - manifest: vm_manifest, + manifest: vm_config, }) } diff --git a/odoroboctl/src/cli.rs b/odoroboctl/src/cli.rs index a92f0b6..79ebc3b 100644 --- a/odoroboctl/src/cli.rs +++ b/odoroboctl/src/cli.rs @@ -2,19 +2,13 @@ use std::collections::BTreeMap; use bytesize::ByteSize; use clap::{Parser, Subcommand}; -use clap::{Parser, Subcommand}; use odorobo::types::{ AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, CreateVMRequest, MetadataTable, ObjectMetadata, Operator, VMData, VirtualMachine, }; -use odorobo::types::{CreateVMRequest, VMData, VirtualMachine}; -use reqwest::{Client, Response}; use reqwest::{Client, Response}; use serde::Deserialize; -use serde::Deserialize; use stable_eyre::Result; -use stable_eyre::Result; -use ulid::Ulid; use ulid::Ulid; #[derive(Parser)] From 2f1ad4e5037fab3c9772c055d3eb01913fa8931c Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Fri, 10 Jul 2026 02:23:33 -0500 Subject: [PATCH 13/35] cargo fmt --- odorobo/src/ch_driver/actor.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index 0427aa3..5b78cdf 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -40,8 +40,12 @@ impl Actor for VMActor { #[tracing::instrument(skip_all)] async fn on_start((vmid, vm_config): Self::Args, actor_ref: ActorRef) -> Result { - let mut vminstance = - VMInstance::spawn(&vmid.to_string(), vm_config.clone().map(VmConfig::from), None).await?; + let mut vminstance = VMInstance::spawn( + &vmid.to_string(), + vm_config.clone().map(VmConfig::from), + None, + ) + .await?; // Take the child process out so we can watch for unexpected death. // destroy() handles a missing child_process gracefully. From 156188488b69779af12a9b05019019f7f2e1db18 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Sat, 11 Jul 2026 03:13:43 -0500 Subject: [PATCH 14/35] random debugging stuff --- .zed/debug.json | 14 ++++++++++ odorobo/src/actors/scheduler_actor.rs | 37 +++++++++++++++------------ odoroboctl/src/cli.rs | 2 +- 3 files changed, 35 insertions(+), 18 deletions(-) create mode 100644 .zed/debug.json diff --git a/.zed/debug.json b/.zed/debug.json new file mode 100644 index 0000000..00271f1 --- /dev/null +++ b/.zed/debug.json @@ -0,0 +1,14 @@ +[ + { + "label": "Run manager agent", + "build": { + "command": "cargo", + "args": ["build"] + }, + "sourceLanguages": ["rust"], + "program": "$ZED_WORKTREE_ROOT/target/debug/odorobo", + "args": ["--manager-enabled"], + "request": "launch", + "adapter": "CodeLLDB" + } +] diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 93965db..c584c89 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -44,6 +44,8 @@ pub struct CachedVMActor { pub data: GetVMInfoReply, } +// todo: vms arent evicted from caches properly if the VM doesnt boot properly. need to determine why. + // todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. // especially when we haven't fully made decisions about some other things. // @@ -278,6 +280,9 @@ impl SchedulerActor { // intentionally ignoring results because we want to keep finding actors even if an attempt fails let _ = tokio::join!(vm_join_handle, agent_join_handle); + //info!(?vm_data_cache_arc_clone); + //info!(?agent_data_cache_arc_clone); + interval.tick().await; } })); @@ -302,21 +307,21 @@ impl SchedulerActor { /// and then if someone tries to schedule lets say 10 VMs in a batch, we could end up scheduling them all to the same agent because the metadata hasn't updated. /// - there are a few solutions for this but they all kinda suck, mostly due to also making sure we deal with latency properly. I am ignoring the issue for now. fn schedule_agent(&mut self, msg: &CreateVM) -> Result, Report> { - // todo: this could likely be better idiomatic rust. - // I suspect there is a map-reduce operation that does the exact scoring thing I am trying to do. - // I also assume there is a better function for the and_then - let best_agent = self - .agent_data_cache - .iter() - .map(|agent| (agent.actor_ref.clone(), self.score_agent(msg, &agent))) - .reduce(|best, new| if new.1 > best.1 { new } else { best }) - .and_then(|best| { - if best.1 == AgentScore::REJECTED { - None - } else { - Some(best.0) - } - }); + let mut best_agent = None; + let mut best_score = AgentScore::REJECTED; + + for agent in self.agent_data_cache.iter() { + let score = self.score_agent(msg, &agent); + + info!(?score, agent_id=?(agent.key()), "scored agent"); + + if score > best_score { + best_agent = Some(agent.actor_ref.clone()); + best_score = score; + } + } + + info!(?best_score, ?best_agent, "best agent"); best_agent.ok_or_eyre("No valid agents found.") } @@ -421,8 +426,6 @@ impl SchedulerActor { // agent_score = 1; // } - info!(?score, agent_id=?(agent.key()), "scored agent"); - score } } diff --git a/odoroboctl/src/cli.rs b/odoroboctl/src/cli.rs index 79ebc3b..4c6503a 100644 --- a/odoroboctl/src/cli.rs +++ b/odoroboctl/src/cli.rs @@ -110,7 +110,7 @@ pub async fn run_command(cli: Cli) -> Result<()> { vcpus: 4, max_vcpus: None, memory: ByteSize::gib(4), - image: "/var/lib/odorobo/f43.raw".to_string(), + image: "/var/lib/odorobo/f43-1.raw".to_string(), ..Default::default() }, metadata: Some(ObjectMetadata { From 188a2b9a1bf8a6875748a25c4bddc79bba28ea32 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Thu, 16 Jul 2026 04:47:55 -0500 Subject: [PATCH 15/35] fix: remove unused AffinityDirection enum --- odorobo/src/types.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index 0852983..e425bf1 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -195,12 +195,6 @@ pub enum AffinityType { /// if there are several metadata tables, their results will be ANDed together /// EX: if a requirement is checked against several VMs, it must pass all VMs for the requirement to be fulfilled. -#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] -pub enum AffinityDirection { - Normal, - Anti, -} - #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub struct AffinityRequirement { pub key: String, From 43b6f56f88c8aa64f5663999ac847127a7811db1 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Thu, 16 Jul 2026 04:53:27 -0500 Subject: [PATCH 16/35] feat: remove VMs from actors if unreachacble by updaters. this also fixes an issue where if a VM is never linked (ex: it didnt boot) it would never be removed from the caches. --- odorobo/src/actors/scheduler_actor.rs | 68 ++++++++++++++++++++++----- 1 file changed, 57 insertions(+), 11 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index c584c89..994bb09 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -41,11 +41,12 @@ pub struct CachedAgentActor { #[derive(Debug, Clone)] pub struct CachedVMActor { pub actor_ref: Option>, + /// The agent this VM was scheduled on. Used by the updater task to clean up + /// the agent's extended_vm_set if the VM actor dies before it's linked. + pub scheduled_agent_id: ActorId, pub data: GetVMInfoReply, } -// todo: vms arent evicted from caches properly if the VM doesnt boot properly. need to determine why. - // todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. // especially when we haven't fully made decisions about some other things. // @@ -120,6 +121,7 @@ impl SchedulerActor { vm_actorid_ulid_map: Arc>, data_cache: Arc>, keepalive_tasks: Arc>>, + agent_data_cache: Arc>, ) -> Result<(), Report> { trace!("running vm_actor_finder"); @@ -135,9 +137,15 @@ impl SchedulerActor { let vm_actorid_ulid_map_clone = Arc::clone(&vm_actorid_ulid_map); let data_cache_clone = Arc::clone(&data_cache); + let agent_data_cache_clone = Arc::clone(&agent_data_cache); let updater_task = tokio::spawn(async move { - Self::vm_updater_task(vm_actor, vm_actorid_ulid_map_clone, data_cache_clone) - .await; + Self::vm_updater_task( + vm_actor, + vm_actorid_ulid_map_clone, + data_cache_clone, + agent_data_cache_clone, + ) + .await; }); keepalive_tasks.insert(vm_actor_id, updater_task); @@ -151,6 +159,7 @@ impl SchedulerActor { actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, data_cache: Arc>, + agent_data_cache: Arc>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails = 0; @@ -164,6 +173,14 @@ impl SchedulerActor { vmid, CachedVMActor { actor_ref: Some(actor_ref.clone()), + // preserve the original scheduled_agent_id if an entry already exists + // (e.g. from CreateVM::handle inserting it before the actor was reachable), + // otherwise default to the vm actor's own id (won't match any agent entry, + // so no accidental cleanup of wrong agent). + scheduled_agent_id: data_cache + .get(&vmid) + .map(|e| e.scheduled_agent_id) + .unwrap_or(actor_ref.id()), data: data, }, ); @@ -174,8 +191,35 @@ impl SchedulerActor { } if fails > 5 { - // todo: possibly better error handling - warn!(?actor_ref, "can no longer reach agent actor.") + warn!(?actor_ref, "can no longer reach vm actor, cleaning up cache entries"); + + let vmid = vm_actorid_ulid_map.remove(&actor_ref.id()).map(|(_, vmid)| vmid); + + // if the actor was never reachable, there's no vm_actorid_ulid_map entry. + // try to recover the vmid from the data_cache by scanning for a stale entry + // that matches this actor_ref. + let vmid = vmid.or_else(|| { + data_cache.iter().find_map(|entry| { + if entry.actor_ref.as_ref().map(|r| r.id()) == Some(actor_ref.id()) { + Some(*entry.key()) + } else { + None + } + }) + }); + + if let Some(vmid) = vmid { + if let Some(cached_vm) = data_cache.get(&vmid) { + agent_data_cache.alter(&cached_vm.scheduled_agent_id, |_, mut v| { + v.extended_vm_set.remove(&vmid); + v + }); + } + + data_cache.remove(&vmid); + } + + return; } interval.tick().await; @@ -219,7 +263,6 @@ impl SchedulerActor { let mut fails = 0; loop { if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { - // todo: replace with kameo stream if data_cache.contains_key(&actor_ref.id()) { data_cache.alter(&actor_ref.id(), |_, mut v| { v.data = data; @@ -245,8 +288,9 @@ impl SchedulerActor { } if fails > 5 { - // todo: possibly better error handling - warn!(?actor_ref, "can no longer reach agent actor.") + warn!(?actor_ref, "can no longer reach agent actor, stopping updater"); + data_cache.remove(&actor_ref.id()); + return; } interval.tick().await; @@ -269,6 +313,7 @@ impl SchedulerActor { Arc::clone(&vm_actorid_ulid_map_arc_clone), Arc::clone(&vm_data_cache_arc_clone), Arc::clone(&vm_keepalive_tasks_arc_clone), + Arc::clone(&agent_data_cache_arc_clone), ); let agent_join_handle = Self::agent_actor_finder( @@ -355,7 +400,7 @@ impl SchedulerActor { return AgentScore::REJECTED; } - score.general += (agent_max_ram.as_u64() - agent.data.used_ram.as_u64()) as f32 + score.general += (agent_max_ram.as_u64() - agent_used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; // roughly based on: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ @@ -584,6 +629,7 @@ impl Message for SchedulerActor { msg.vmid, CachedVMActor { actor_ref: None, + scheduled_agent_id: target_agent.id(), data: GetVMInfoReply { vmid: msg.vmid, config: Some(msg.config.clone()), @@ -593,7 +639,7 @@ impl Message for SchedulerActor { let reply = target_agent.ask(&msg).await; - // remove from caches if we fail to schedule + // remove from caches if the agent rejected the request if reply.is_err() { self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { v.extended_vm_set.remove(&msg.vmid); From a784babdbf84e2ff2c630ccfb25d889f06e039a1 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Thu, 16 Jul 2026 05:37:31 -0500 Subject: [PATCH 17/35] chore: cargo fmt --- odorobo/src/actors/scheduler_actor.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 994bb09..29a088c 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -191,9 +191,14 @@ impl SchedulerActor { } if fails > 5 { - warn!(?actor_ref, "can no longer reach vm actor, cleaning up cache entries"); + warn!( + ?actor_ref, + "can no longer reach vm actor, cleaning up cache entries" + ); - let vmid = vm_actorid_ulid_map.remove(&actor_ref.id()).map(|(_, vmid)| vmid); + let vmid = vm_actorid_ulid_map + .remove(&actor_ref.id()) + .map(|(_, vmid)| vmid); // if the actor was never reachable, there's no vm_actorid_ulid_map entry. // try to recover the vmid from the data_cache by scanning for a stale entry @@ -288,7 +293,10 @@ impl SchedulerActor { } if fails > 5 { - warn!(?actor_ref, "can no longer reach agent actor, stopping updater"); + warn!( + ?actor_ref, + "can no longer reach agent actor, stopping updater" + ); data_cache.remove(&actor_ref.id()); return; } From 81d5b9c439f6ff1838cc326fc284db9666b87aa8 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Thu, 16 Jul 2026 15:13:00 -0500 Subject: [PATCH 18/35] fix: better extended_vm_set removal and setup --- odorobo/src/actors/scheduler_actor.rs | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 29a088c..113c887 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -281,8 +281,8 @@ impl SchedulerActor { actor_ref.id(), CachedAgentActor { actor_ref: actor_ref.clone(), - data: data, - extended_vm_set: AHashSet::new(), + data: data.clone(), + extended_vm_set: AHashSet::from_iter(data.vms.iter().copied()), }, ); } @@ -366,8 +366,6 @@ impl SchedulerActor { for agent in self.agent_data_cache.iter() { let score = self.score_agent(msg, &agent); - info!(?score, agent_id=?(agent.key()), "scored agent"); - if score > best_score { best_agent = Some(agent.actor_ref.clone()); best_score = score; @@ -596,10 +594,10 @@ impl Actor for SchedulerActor { }; if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&actor_id) { - // todo: this solution definitely isn't optimal. we likely should be storing which agent actor, this specific actor id is scheduled on and specifically removing it from that one. - // especially because things get iffy with migration, but im ignoring that for the moment. - for mut agent in self.agent_data_cache.iter_mut() { - agent.extended_vm_set.remove(&vmid); + if let Some(cached_vm) = self.vm_data_cache.get(&vmid) { + if let Some(mut agent) = self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) { + agent.extended_vm_set.remove(&vmid); + } } // todo: we likely should keep a copy of the VirtualMachine manifest in the cache. From ecbc2e928b4bf1dda249671e3d4c676d0bf9d318 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Thu, 16 Jul 2026 15:19:16 -0500 Subject: [PATCH 19/35] chore: clippy and fmt --- odorobo/src/actors/scheduler_actor.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 113c887..151c0e8 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -181,7 +181,7 @@ impl SchedulerActor { .get(&vmid) .map(|e| e.scheduled_agent_id) .unwrap_or(actor_ref.id()), - data: data, + data, }, ); @@ -487,8 +487,8 @@ fn evaluate_table_value(value_option: &Option<&String>, requirement: &AffinityRe }; match requirement.operator { - Operator::In => return requirement.values.iter().any(|e| *e == *value), - Operator::NotIn => return !requirement.values.iter().any(|e| *e == *value), + Operator::In => requirement.values.contains(value), + Operator::NotIn => !requirement.values.contains(value), Operator::Lt | Operator::Gt => { if requirement.values.len() != 1 { return false; @@ -503,12 +503,12 @@ fn evaluate_table_value(value_option: &Option<&String>, requirement: &AffinityRe }; if requirement.operator == Operator::Lt { - return value_number < requirement_value_number; + value_number < requirement_value_number } else { - return value_number > requirement_value_number; + value_number > requirement_value_number } } - }; + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -594,10 +594,11 @@ impl Actor for SchedulerActor { }; if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&actor_id) { - if let Some(cached_vm) = self.vm_data_cache.get(&vmid) { - if let Some(mut agent) = self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) { - agent.extended_vm_set.remove(&vmid); - } + if let Some(cached_vm) = self.vm_data_cache.get(&vmid) + && let Some(mut agent) = + self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) + { + agent.extended_vm_set.remove(&vmid); } // todo: we likely should keep a copy of the VirtualMachine manifest in the cache. From 7698f8f67936ab82d965cba96a1e51b891c8ebcd Mon Sep 17 00:00:00 2001 From: Willow C Reed Date: Fri, 31 Jul 2026 02:10:26 -0600 Subject: [PATCH 20/35] fix clippy warnings --- odorobo/src/actors/scheduler_actor.rs | 127 ++++++++++++++++---------- odorobo/src/types.rs | 4 +- 2 files changed, 79 insertions(+), 52 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index a734fae..74012dd 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -45,7 +45,7 @@ pub struct CachedAgentActor { pub struct CachedVMActor { pub actor_ref: Option>, /// The agent this VM was scheduled on. Used by the updater task to clean up - /// the agent's extended_vm_set if the VM actor dies before it's linked. + /// the agent's `extended_vm_set` if the VM actor dies before it's linked. pub scheduled_agent_id: ActorId, pub data: GetVMInfoReply, } @@ -99,19 +99,15 @@ static VCPU_OVERPROVISIONMENT_NUMERATOR: u32 = 2; static VCPU_OVERPROVISIONMENT_DENOMINATOR: u32 = 1; impl SchedulerActor { - async fn lookup_agent_by_actor_id( - &mut self, - actor_id: &ActorId, - ) -> Option> { + #[expect(dead_code, reason = "reserved for explicit placement by actor id")] + fn lookup_agent_by_actor_id(&self, actor_id: &ActorId) -> Option> { self.agent_data_cache .get(actor_id) .map(|data| data.actor_ref.clone()) } - async fn lookup_agent_by_hostname( - &mut self, - hostname: &str, - ) -> Option> { + #[expect(dead_code, reason = "reserved for explicit placement by hostname")] + fn lookup_agent_by_hostname(&self, hostname: &str) -> Option> { self.agent_data_cache .iter() .find(|data| data.data.hostname == hostname) @@ -120,7 +116,7 @@ impl SchedulerActor { // someone should likely give caleb a firm talking to about code duplication due to this section, but things are just different enough that trying to make them one function requires usage of a lot of generics which feels even worse. so i dont know what to do. cappy please fix. i hate this. async fn vm_actor_finder( - parent_actor_ref: RemoteActorRef, + parent_actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, data_cache: Arc>, keepalive_tasks: Arc>>, @@ -165,7 +161,7 @@ impl SchedulerActor { agent_data_cache: Arc>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); - let mut fails = 0; + let mut fails: u8 = 0; loop { if let Ok(data) = actor_ref.ask(&GetVMInfo { vmid: None }).await { let vmid = data.vmid; @@ -182,15 +178,14 @@ impl SchedulerActor { // so no accidental cleanup of wrong agent). scheduled_agent_id: data_cache .get(&vmid) - .map(|e| e.scheduled_agent_id) - .unwrap_or(actor_ref.id()), + .map_or_else(|| actor_ref.id(), |e| e.scheduled_agent_id), data, }, ); fails = 0; } else { - fails += 1; + fails = fails.saturating_add(1); } if fails > 5 { @@ -208,11 +203,8 @@ impl SchedulerActor { // that matches this actor_ref. let vmid = vmid.or_else(|| { data_cache.iter().find_map(|entry| { - if entry.actor_ref.as_ref().map(|r| r.id()) == Some(actor_ref.id()) { - Some(*entry.key()) - } else { - None - } + (entry.actor_ref.as_ref().map(RemoteActorRef::id) == Some(actor_ref.id())) + .then(|| *entry.key()) }) }); @@ -235,7 +227,7 @@ impl SchedulerActor { } async fn agent_actor_finder( - parent_actor_ref: RemoteActorRef, + parent_actor_ref: RemoteActorRef, data_cache: Arc>, keepalive_tasks: Arc>>, ) -> Result<(), Report> { @@ -268,7 +260,7 @@ impl SchedulerActor { data_cache: Arc>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); - let mut fails = 0; + let mut fails: u8 = 0; loop { if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { if data_cache.contains_key(&actor_ref.id()) { @@ -285,14 +277,14 @@ impl SchedulerActor { CachedAgentActor { actor_ref: actor_ref.clone(), data: data.clone(), - extended_vm_set: AHashSet::from_iter(data.vms.iter().copied()), + extended_vm_set: data.vms.iter().copied().collect(), }, ); } fails = 0; } else { - fails += 1; + fails = fails.saturating_add(1); } if fails > 5 { @@ -334,7 +326,13 @@ impl SchedulerActor { ); // intentionally ignoring results because we want to keep finding actors even if an attempt fails - let _ = tokio::join!(vm_join_handle, agent_join_handle); + let (vm_result, agent_result) = tokio::join!(vm_join_handle, agent_join_handle); + if let Err(error) = vm_result { + warn!(?error, "VM actor discovery failed"); + } + if let Err(error) = agent_result { + warn!(?error, "agent actor discovery failed"); + } //info!(?vm_data_cache_arc_clone); //info!(?agent_data_cache_arc_clone); @@ -356,13 +354,13 @@ impl SchedulerActor { /// The general score uses things like resource utilization to not over load any specific agent. /// /// - /// Affinity rules are roughly based on https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ + /// Affinity rules are roughly based on . /// /// todo: /// - the cache likely needs to be updated automatically when a new vm is scheduled for info like used resources, because otherwise we have to deal with latency on that data we are using /// and then if someone tries to schedule lets say 10 VMs in a batch, we could end up scheduling them all to the same agent because the metadata hasn't updated. /// - there are a few solutions for this but they all kinda suck, mostly due to also making sure we deal with latency properly. I am ignoring the issue for now. - fn schedule_agent(&mut self, msg: &CreateVM) -> Result, Report> { + fn schedule_agent(&self, msg: &CreateVM) -> Result, Report> { let mut best_agent = None; let mut best_score = AgentScore::REJECTED; @@ -390,29 +388,55 @@ impl SchedulerActor { ) -> AgentScore { let mut score = AgentScore::default(); - let agent_max_vcpus = agent.data.vcpus * VCPU_OVERPROVISIONMENT_NUMERATOR - / VCPU_OVERPROVISIONMENT_DENOMINATOR; + let agent_max_vcpus = agent + .data + .vcpus + .saturating_mul(VCPU_OVERPROVISIONMENT_NUMERATOR) + .checked_div(VCPU_OVERPROVISIONMENT_DENOMINATOR) + .unwrap_or(u32::MAX); // todo: do we care about VMData.max_vcpus? - let agent_used_vcpus = agent.data.used_vcpus + msg.config.data.vcpus; + let agent_used_vcpus = agent.data.used_vcpus.saturating_add(msg.config.data.vcpus); if agent_used_vcpus >= agent_max_vcpus { return AgentScore::REJECTED; } - score.general += (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; + #[expect( + clippy::cast_precision_loss, + reason = "the scheduler score intentionally uses f32 ratios" + )] + #[expect( + clippy::arithmetic_side_effects, + reason = "the preceding capacity check guarantees non-negative subtraction" + )] + let vcpu_headroom = (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; + score.general += vcpu_headroom; // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. let agent_max_ram = agent.data.ram; - let agent_used_ram = agent.data.used_ram + msg.config.data.memory; + let agent_used_ram = bytesize::ByteSize::b( + agent + .data + .used_ram + .as_u64() + .saturating_add(msg.config.data.memory.as_u64()), + ); if agent_used_ram >= agent_max_ram { return AgentScore::REJECTED; } - score.general += (agent_max_ram.as_u64() - agent_used_ram.as_u64()) as f32 + #[expect( + clippy::cast_precision_loss, + reason = "the scheduler score intentionally uses f32 ratios" + )] + let ram_headroom = agent_max_ram + .as_u64() + .saturating_sub(agent_used_ram.as_u64()) as f32 / agent_max_ram.as_u64() as f32; + score.general += ram_headroom; - // roughly based on: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/ + // Roughly based on . if let Some(affinity_rules) = &msg.config.affinity { for rule in affinity_rules { let mut metadata_tables: Vec = Vec::with_capacity(1); @@ -434,7 +458,7 @@ impl SchedulerActor { } } AffinityType::Agent => metadata_tables.push(agent.data.metadata.clone()), - }; + } let mut follows_rule = false; @@ -449,7 +473,7 @@ impl SchedulerActor { let value_option = table.get(&requirement.key); - if !evaluate_table_value(&value_option, requirement) { + if !evaluate_table_value(value_option, requirement) { requirement_outcome = false; break; } @@ -467,7 +491,10 @@ impl SchedulerActor { (AffinityStrictness::Required, false) => return AgentScore::REJECTED, (AffinityStrictness::Required, true) => {} // specifically do nothing (AffinityStrictness::Preferred { weight }, follows_rule) => { - score.affinity += follows_rule as i64 * weight; + let follows_rule = i64::from(follows_rule); + score.affinity = score + .affinity + .saturating_add(follows_rule.saturating_mul(weight)); } } } @@ -484,8 +511,8 @@ impl SchedulerActor { } } -fn evaluate_table_value(value_option: &Option<&String>, requirement: &AffinityRequirement) -> bool { - let Some(value) = *value_option else { +fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityRequirement) -> bool { + let Some(value) = value_option else { return false; }; @@ -557,7 +584,7 @@ impl Actor for SchedulerActor { info!(?peer_id, "Scheduler Actor started!"); - let mut scheduler_actor = SchedulerActor { + let mut scheduler_actor = Self { agent_data_cache: Arc::new(DashMap::new()), agent_keepalive_tasks: Arc::new(DashMap::new()), vm_actorid_ulid_map: Arc::new(DashMap::new()), @@ -574,29 +601,29 @@ impl Actor for SchedulerActor { async fn on_link_died( &mut self, actor_ref: WeakActorRef, - actor_id: ActorId, + id: ActorId, reason: ActorStopReason, ) -> Result, Self::Error> { - warn!(?actor_id, ?reason, "Linked actor died"); + warn!(?id, ?reason, "Linked actor died"); // check that scheduler actor is still alive. let Some(_) = actor_ref.upgrade() else { return Ok(ControlFlow::Break(ActorStopReason::Killed)); }; - if let Some((_, keepalive_task)) = self.agent_keepalive_tasks.remove(&actor_id) { - trace!(?actor_id, "Aborting agent keepalive task"); + if let Some((_, keepalive_task)) = self.agent_keepalive_tasks.remove(&id) { + trace!(?id, "Aborting agent keepalive task"); keepalive_task.abort(); - }; + } - self.agent_data_cache.remove(&actor_id); + self.agent_data_cache.remove(&id); - if let Some((_, keepalive_task)) = self.vm_keepalive_tasks.remove(&actor_id) { - trace!(?actor_id, "Aborting vm keepalive task"); + if let Some((_, keepalive_task)) = self.vm_keepalive_tasks.remove(&id) { + trace!(?id, "Aborting vm keepalive task"); keepalive_task.abort(); - }; + } - if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&actor_id) { + if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&id) { if let Some(cached_vm) = self.vm_data_cache.get(&vmid) && let Some(mut agent) = self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) @@ -608,7 +635,7 @@ impl Actor for SchedulerActor { // instead of removing the vm entirely, we should just modify the status to shutdown or crashed or something. // // another potential issue is that the link dying doesn't guarantee that the vm is dead. if there is a networking partition, things get iffy. - trace!(?actor_id, ?vmid, "Removing vm from vm_data_cache"); + trace!(?id, ?vmid, "Removing vm from vm_data_cache"); self.vm_data_cache.remove(&vmid); } diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index 07a7507..be0cf26 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -194,7 +194,7 @@ pub enum AffinityType { Agent, } -/// if there are several metadata tables, their results will be ANDed together +/// If there are several metadata tables, their results will be `ANDed` together /// EX: if a requirement is checked against several VMs, it must pass all VMs for the requirement to be fulfilled. #[derive(Serialize, Deserialize, Debug, JsonSchema, Clone)] pub struct AffinityRequirement { @@ -210,7 +210,7 @@ pub enum MetadataTable { Annotation, } -#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone, PartialEq)] +#[derive(Serialize, Deserialize, Debug, JsonSchema, Clone, PartialEq, Eq)] pub enum Operator { In, NotIn, From 929c4aa081373b584b175dd8a90c073c8b0d96a1 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 16:11:32 -0700 Subject: [PATCH 21/35] resolve review comments --- odorobo/src/actors/scheduler_actor.rs | 257 ++++++++++++++++++++------ odoroboctl/src/cli.rs | 13 +- 2 files changed, 205 insertions(+), 65 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 74012dd..09efbc6 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -88,7 +88,7 @@ pub struct SchedulerActor { // alternatively we could have the other manager nodes try to keep track of that data, but i think we are going to run into issues with keeping the state consistent between all nodes. // we may need to make some architecture designs about db consistency vs uptime vs speed in that situation, and im not doing that on my own. pub vm_actorid_ulid_map: Arc>, - pub vm_data_cache: Arc>, + pub vm_data_cache: Arc>>, pub vm_keepalive_tasks: Arc>>, pub cache_actor_finder: Option>, @@ -118,7 +118,7 @@ impl SchedulerActor { async fn vm_actor_finder( parent_actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, - data_cache: Arc>, + data_cache: Arc>>, keepalive_tasks: Arc>>, agent_data_cache: Arc>, ) -> Result<(), Report> { @@ -127,13 +127,16 @@ impl SchedulerActor { let mut vm_actor_stream = RemoteActorRef::::lookup_all(VM); while let Some(vm_actor) = vm_actor_stream.try_next().await? { - if !keepalive_tasks.contains_key(&vm_actor.id()) { + let vm_actor_id = vm_actor.id(); + let updater_is_running = keepalive_tasks + .get(&vm_actor_id) + .is_some_and(|task| !task.is_finished()); + if !updater_is_running { + keepalive_tasks.remove(&vm_actor_id); trace!(?vm_actor, "starting vm_updater_task"); parent_actor_ref.link_remote(&vm_actor).await?; - let vm_actor_id = vm_actor.id(); - let vm_actorid_ulid_map_clone = Arc::clone(&vm_actorid_ulid_map); let data_cache_clone = Arc::clone(&data_cache); let agent_data_cache_clone = Arc::clone(&agent_data_cache); @@ -154,10 +157,14 @@ impl SchedulerActor { Ok(()) } + #[expect( + clippy::too_many_lines, + reason = "the updater keeps polling and performs coordinated cache cleanup" + )] async fn vm_updater_task( actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, - data_cache: Arc>, + data_cache: Arc>>, agent_data_cache: Arc>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); @@ -168,20 +175,42 @@ impl SchedulerActor { vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that - data_cache.insert( - vmid, - CachedVMActor { - actor_ref: Some(actor_ref.clone()), - // preserve the original scheduled_agent_id if an entry already exists - // (e.g. from CreateVM::handle inserting it before the actor was reachable), - // otherwise default to the vm actor's own id (won't match any agent entry, - // so no accidental cleanup of wrong agent). - scheduled_agent_id: data_cache - .get(&vmid) - .map_or_else(|| actor_ref.id(), |e| e.scheduled_agent_id), - data, - }, - ); + let scheduled_agent_id = data_cache + .get(&vmid) + .and_then(|entries| { + entries.iter().find_map(|entry| { + entry + .actor_ref + .as_ref() + .is_none_or(|actor| actor.id() == actor_ref.id()) + .then_some(entry.scheduled_agent_id) + }) + }) + .unwrap_or_else(|| actor_ref.id()); + let cached_vm = CachedVMActor { + actor_ref: Some(actor_ref.clone()), + scheduled_agent_id, + data, + }; + data_cache + .entry(vmid) + .and_modify(|entries| { + if let Some(entry) = entries.iter_mut().find(|entry| { + entry + .actor_ref + .as_ref() + .is_some_and(|actor| actor.id() == actor_ref.id()) + }) { + *entry = cached_vm.clone(); + } else if let Some(entry) = + entries.iter_mut().find(|entry| entry.actor_ref.is_none()) + { + *entry = cached_vm.clone(); + } else { + entries.push(cached_vm.clone()); + } + }) + .or_insert_with(|| vec![cached_vm]); fails = 0; } else { @@ -203,20 +232,47 @@ impl SchedulerActor { // that matches this actor_ref. let vmid = vmid.or_else(|| { data_cache.iter().find_map(|entry| { - (entry.actor_ref.as_ref().map(RemoteActorRef::id) == Some(actor_ref.id())) + entry + .value() + .iter() + .any(|cached| { + cached + .actor_ref + .as_ref() + .is_some_and(|actor| actor.id() == actor_ref.id()) + }) .then(|| *entry.key()) }) }); if let Some(vmid) = vmid { - if let Some(cached_vm) = data_cache.get(&vmid) { - agent_data_cache.alter(&cached_vm.scheduled_agent_id, |_, mut v| { + let removed_cached_vm = data_cache.get_mut(&vmid).and_then(|mut entries| { + entries + .iter() + .position(|entry| { + entry + .actor_ref + .as_ref() + .is_some_and(|actor| actor.id() == actor_ref.id()) + }) + .map(|index| (index, entries.remove(index))) + }); + if let Some(cached_vm) = removed_cached_vm + && data_cache + .get(&vmid) + .is_none_or(|entries| entries.is_empty()) + { + agent_data_cache.alter(&cached_vm.1.scheduled_agent_id, |_, mut v| { v.extended_vm_set.remove(&vmid); v }); } - - data_cache.remove(&vmid); + if data_cache + .get(&vmid) + .is_some_and(|entries| entries.is_empty()) + { + data_cache.remove(&vmid); + } } return; @@ -236,13 +292,16 @@ impl SchedulerActor { let mut agent_actor_stream = RemoteActorRef::::lookup_all(AGENT); while let Some(agent_actor) = agent_actor_stream.try_next().await? { - if !keepalive_tasks.contains_key(&agent_actor.id()) { + let agent_actor_id = agent_actor.id(); + let updater_is_running = keepalive_tasks + .get(&agent_actor_id) + .is_some_and(|task| !task.is_finished()); + if !updater_is_running { + keepalive_tasks.remove(&agent_actor_id); trace!(?agent_actor, "starting agent_updater_task"); parent_actor_ref.link_remote(&agent_actor).await?; - let agent_actor_id = agent_actor.id(); - let data_cache_clone = Arc::clone(&data_cache); let updater_task = tokio::spawn(async move { Self::agent_updater_task(agent_actor, data_cache_clone).await; @@ -267,6 +326,7 @@ impl SchedulerActor { data_cache.alter(&actor_ref.id(), |_, mut v| { v.data = data; + v.extended_vm_set.retain(|vmid| v.data.vms.contains(vmid)); v.extended_vm_set.extend(v.data.vms.iter()); v @@ -444,16 +504,17 @@ impl SchedulerActor { match rule.affinity_type { AffinityType::VirtualMachine => { for vmid in &agent.extended_vm_set { - let Some(vm_data_cache_ref) = &self.vm_data_cache.get(vmid) else { - continue; - }; - - let Some(vm_manifest) = &vm_data_cache_ref.data.config else { + let Some(vm_data_cache_refs) = self.vm_data_cache.get(vmid) else { continue; }; - if let Some(metadata) = &vm_manifest.metadata { - metadata_tables.push(metadata.clone()); + for vm_data_cache_ref in vm_data_cache_refs.iter() { + let Some(vm_manifest) = &vm_data_cache_ref.data.config else { + continue; + }; + if let Some(metadata) = &vm_manifest.metadata { + metadata_tables.push(metadata.clone()); + } } } } @@ -541,6 +602,64 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq } } +#[cfg(test)] +mod tests { + use super::evaluate_table_value; + use crate::types::{AffinityRequirement, MetadataTable, Operator}; + use std::collections::BTreeMap; + + fn requirement(operator: Operator, values: &[&str]) -> AffinityRequirement { + AffinityRequirement { + key: "tier".to_owned(), + table: MetadataTable::Label, + operator, + values: values.iter().map(|value| (*value).to_owned()).collect(), + } + } + + #[test] + fn evaluates_membership_and_missing_keys() { + let metadata = BTreeMap::from([("tier".to_owned(), "frontend".to_owned())]); + assert!(evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::In, &["frontend", "api"]) + )); + assert!(!evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::In, &["backend"]) + )); + assert!(!evaluate_table_value( + None, + &requirement(Operator::In, &["frontend"]) + )); + } + + #[test] + fn evaluates_not_in_and_numeric_comparisons() { + let metadata = BTreeMap::from([("tier".to_owned(), "4".to_owned())]); + assert!(evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::NotIn, &["5"]) + )); + assert!(evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::Lt, &["5"]) + )); + assert!(evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::Gt, &["3"]) + )); + assert!(!evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::Lt, &["4", "5"]) + )); + assert!(!evaluate_table_value( + metadata.get("tier"), + &requirement(Operator::Gt, &["not-a-number"]) + )); + } +} + #[derive(Debug, Clone, Copy, PartialEq)] struct AgentScore { general: f32, @@ -623,20 +742,28 @@ impl Actor for SchedulerActor { keepalive_task.abort(); } - if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&id) { - if let Some(cached_vm) = self.vm_data_cache.get(&vmid) - && let Some(mut agent) = - self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) - { - agent.extended_vm_set.remove(&vmid); + if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&id) + && let Some(mut entries) = self.vm_data_cache.get_mut(&vmid) + { + let removed_cached_vm = entries + .iter() + .position(|entry| { + entry + .actor_ref + .as_ref() + .is_some_and(|actor| actor.id() == id) + }) + .map(|index| entries.remove(index)); + if entries.is_empty() { + drop(entries); + self.vm_data_cache.remove(&vmid); + if let Some(cached_vm) = removed_cached_vm + && let Some(mut agent) = + self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) + { + agent.extended_vm_set.remove(&vmid); + } } - - // todo: we likely should keep a copy of the VirtualMachine manifest in the cache. - // instead of removing the vm entirely, we should just modify the status to shutdown or crashed or something. - // - // another potential issue is that the link dying doesn't guarantee that the vm is dead. if there is a networking partition, things get iffy. - trace!(?id, ?vmid, "Removing vm from vm_data_cache"); - self.vm_data_cache.remove(&vmid); } // todo: attempt vm restarts if necessary. @@ -662,28 +789,38 @@ impl Message for SchedulerActor { return Err(eyre!("target agent is not in data cache")); } - self.vm_data_cache.insert( - msg.vmid, - CachedVMActor { + self.vm_data_cache + .entry(msg.vmid) + .or_default() + .push(CachedVMActor { actor_ref: None, scheduled_agent_id: target_agent.id(), data: GetVMInfoReply { vmid: msg.vmid, config: Some(msg.config.clone()), }, - }, - ); + }); let reply = target_agent.ask(&msg).await; - // remove from caches if the agent rejected the request if reply.is_err() { - self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { - v.extended_vm_set.remove(&msg.vmid); - - v - }); - self.vm_data_cache.remove(&msg.vmid); + // A lost reply is not a rejected create. Keep the cache if the VM actor exists. + let actor_exists = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) + .await? + .is_some(); + if !actor_exists { + self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { + v.extended_vm_set.remove(&msg.vmid); + v + }); + if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) { + entries.retain(|entry| entry.actor_ref.is_some()); + if entries.is_empty() { + drop(entries); + self.vm_data_cache.remove(&msg.vmid); + } + } + } } Ok(reply?) diff --git a/odoroboctl/src/cli.rs b/odoroboctl/src/cli.rs index 8450d37..a072d64 100644 --- a/odoroboctl/src/cli.rs +++ b/odoroboctl/src/cli.rs @@ -31,9 +31,12 @@ pub struct Cli { #[derive(Subcommand)] pub enum Command { - /// Create a VM via the scheduler endpoint, - /// optionally also booting it immediately after creation (if `--boot` is specified). - Create, + /// Create a VM via the scheduler endpoint. + Create { + /// Path or URI of the VM disk image. + #[arg(long, env = "ODOROBO_VM_IMAGE")] + image: String, + }, /// List VMs currently known by the manager/agent. List, @@ -101,7 +104,7 @@ pub async fn run_command(cli: Cli) -> Result<()> { let base_url = cli.manager_addr; match cli.command { - Command::Create => { + Command::Create { image } => { // TODO: setup actual cli args for these parameters. or just take in arbitrary json and serialize it into a VirtualMachine. let vm = VirtualMachine { data: VMData { @@ -110,7 +113,7 @@ pub async fn run_command(cli: Cli) -> Result<()> { vcpus: 4, max_vcpus: None, memory: ByteSize::gib(4), - image: "/var/lib/odorobo/f43-1.raw".to_owned(), + image, ..Default::default() }, metadata: Some(ObjectMetadata { From 8fe3e3bffff2e8d0a02f0198943cf50b82aefcb5 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 16:20:46 -0700 Subject: [PATCH 22/35] fix more of the review issues --- odorobo/src/actors/agent_actor.rs | 1 + odorobo/src/actors/scheduler_actor.rs | 248 ++++++++++++++++++-------- odorobo/src/messages/vm.rs | 2 + 3 files changed, 176 insertions(+), 75 deletions(-) diff --git a/odorobo/src/actors/agent_actor.rs b/odorobo/src/actors/agent_actor.rs index 5beade8..fd71734 100644 --- a/odorobo/src/actors/agent_actor.rs +++ b/odorobo/src/actors/agent_actor.rs @@ -128,6 +128,7 @@ impl Message for AgentActor { info!(?vmid, "VM Spawned successfully"); CreateVMReply { config: Some(msg.config), + actor_id: Some(actor_ref.id().to_bytes()), } } } diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 09efbc6..94f6d78 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -257,15 +257,19 @@ impl SchedulerActor { }) .map(|index| (index, entries.remove(index))) }); - if let Some(cached_vm) = removed_cached_vm - && data_cache - .get(&vmid) - .is_none_or(|entries| entries.is_empty()) - { - agent_data_cache.alter(&cached_vm.1.scheduled_agent_id, |_, mut v| { - v.extended_vm_set.remove(&vmid); - v + if let Some(cached_vm) = removed_cached_vm { + let scheduled_agent_id = cached_vm.1.scheduled_agent_id; + let agent_still_has_vm = data_cache.get(&vmid).is_some_and(|entries| { + entries + .iter() + .any(|entry| entry.scheduled_agent_id == scheduled_agent_id) }); + if !agent_still_has_vm { + agent_data_cache.alter(&scheduled_agent_id, |_, mut v| { + v.extended_vm_set.remove(&vmid); + v + }); + } } if data_cache .get(&vmid) @@ -457,7 +461,14 @@ impl SchedulerActor { // todo: do we care about VMData.max_vcpus? let agent_used_vcpus = agent.data.used_vcpus.saturating_add(msg.config.data.vcpus); - if agent_used_vcpus >= agent_max_vcpus { + if !has_capacity( + agent_max_vcpus, + agent.data.used_vcpus, + msg.config.data.vcpus, + agent.data.ram.as_u64(), + agent.data.used_ram.as_u64(), + msg.config.data.memory.as_u64(), + ) { return AgentScore::REJECTED; } @@ -472,7 +483,7 @@ impl SchedulerActor { let vcpu_headroom = (agent_max_vcpus - agent_used_vcpus) as f32 / agent_max_vcpus as f32; score.general += vcpu_headroom; - // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. + // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. let agent_max_ram = agent.data.ram; let agent_used_ram = bytesize::ByteSize::b( agent @@ -482,10 +493,6 @@ impl SchedulerActor { .saturating_add(msg.config.data.memory.as_u64()), ); - if agent_used_ram >= agent_max_ram { - return AgentScore::REJECTED; - } - #[expect( clippy::cast_precision_loss, reason = "the scheduler score intentionally uses f32 ratios" @@ -521,43 +528,12 @@ impl SchedulerActor { AffinityType::Agent => metadata_tables.push(agent.data.metadata.clone()), } - let mut follows_rule = false; - - for requirement in &rule.requirements { - let mut requirement_outcome = true; - - for object_metadata in &metadata_tables { - let table = match requirement.table { - MetadataTable::Label => &object_metadata.labels, - MetadataTable::Annotation => &object_metadata.annotations, - }; - - let value_option = table.get(&requirement.key); - - if !evaluate_table_value(value_option, requirement) { - requirement_outcome = false; - break; - } - } - - if requirement_outcome { - follows_rule = true; - break; - } - } - - follows_rule ^= rule.inverse; + let follows_rule = evaluate_affinity_rule(&metadata_tables, rule); - match (rule.strictness, follows_rule) { - (AffinityStrictness::Required, false) => return AgentScore::REJECTED, - (AffinityStrictness::Required, true) => {} // specifically do nothing - (AffinityStrictness::Preferred { weight }, follows_rule) => { - let follows_rule = i64::from(follows_rule); - score.affinity = score - .affinity - .saturating_add(follows_rule.saturating_mul(weight)); - } - } + let Some(affinity_delta) = affinity_delta(rule.strictness, follows_rule) else { + return AgentScore::REJECTED; + }; + score.affinity = score.affinity.saturating_add(affinity_delta); } } @@ -572,9 +548,61 @@ impl SchedulerActor { } } +const fn has_capacity( + max_vcpus: u32, + used_vcpus: u32, + requested_vcpus: u32, + max_ram: u64, + used_ram: u64, + requested_ram: u64, +) -> bool { + used_vcpus.saturating_add(requested_vcpus) < max_vcpus + && used_ram.saturating_add(requested_ram) < max_ram +} + +fn affinity_delta(strictness: AffinityStrictness, follows_rule: bool) -> Option { + match strictness { + AffinityStrictness::Required if !follows_rule => None, + AffinityStrictness::Required => Some(0), + AffinityStrictness::Preferred { weight } => { + Some(i64::from(follows_rule).saturating_mul(weight)) + } + } +} + +fn evaluate_affinity_rule( + metadata_tables: &[ObjectMetadata], + rule: &crate::types::AffinityRule, +) -> bool { + let mut follows_rule = false; + + for requirement in &rule.requirements { + let mut requirement_outcome = !metadata_tables.is_empty(); + + for object_metadata in metadata_tables { + let table = match requirement.table { + MetadataTable::Label => &object_metadata.labels, + MetadataTable::Annotation => &object_metadata.annotations, + }; + + if !evaluate_table_value(table.get(&requirement.key), requirement) { + requirement_outcome = false; + break; + } + } + + if requirement_outcome { + follows_rule = true; + break; + } + } + + follows_rule ^ rule.inverse +} + fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityRequirement) -> bool { let Some(value) = value_option else { - return false; + return matches!(requirement.operator, Operator::NotIn); }; match requirement.operator { @@ -604,8 +632,11 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq #[cfg(test)] mod tests { - use super::evaluate_table_value; - use crate::types::{AffinityRequirement, MetadataTable, Operator}; + use super::{affinity_delta, evaluate_affinity_rule, evaluate_table_value, has_capacity}; + use crate::types::{ + AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, MetadataTable, + ObjectMetadata, Operator, + }; use std::collections::BTreeMap; fn requirement(operator: Operator, values: &[&str]) -> AffinityRequirement { @@ -632,6 +663,10 @@ mod tests { None, &requirement(Operator::In, &["frontend"]) )); + assert!(evaluate_table_value( + None, + &requirement(Operator::NotIn, &["frontend"]) + )); } #[test] @@ -658,6 +693,46 @@ mod tests { &requirement(Operator::Gt, &["not-a-number"]) )); } + + #[test] + fn evaluates_inverse_and_empty_requirements() { + let metadata = ObjectMetadata { + labels: BTreeMap::from([("tier".to_owned(), "frontend".to_owned())]), + annotations: BTreeMap::new(), + }; + let rule = AffinityRule { + strictness: AffinityStrictness::Required, + affinity_type: AffinityType::Agent, + inverse: true, + requirements: vec![requirement(Operator::In, &["frontend"])], + }; + assert!(!evaluate_affinity_rule(&[metadata], &rule)); + + let empty_rule = AffinityRule { + strictness: AffinityStrictness::Required, + affinity_type: AffinityType::Agent, + inverse: false, + requirements: Vec::new(), + }; + assert!(!evaluate_affinity_rule(&[], &empty_rule)); + } + + #[test] + fn evaluates_required_preferred_and_capacity_rules() { + assert_eq!(affinity_delta(AffinityStrictness::Required, true), Some(0)); + assert_eq!(affinity_delta(AffinityStrictness::Required, false), None); + assert_eq!( + affinity_delta(AffinityStrictness::Preferred { weight: 7 }, true), + Some(7) + ); + assert_eq!( + affinity_delta(AffinityStrictness::Preferred { weight: 7 }, false), + Some(0) + ); + assert!(has_capacity(8, 2, 2, 16, 4, 4)); + assert!(!has_capacity(8, 6, 2, 16, 4, 4)); + assert!(!has_capacity(8, 2, 2, 16, 12, 4)); + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -754,15 +829,22 @@ impl Actor for SchedulerActor { .is_some_and(|actor| actor.id() == id) }) .map(|index| entries.remove(index)); + if let Some(cached_vm) = removed_cached_vm { + let scheduled_agent_id = cached_vm.scheduled_agent_id; + let agent_still_has_vm = entries + .iter() + .any(|entry| entry.scheduled_agent_id == scheduled_agent_id); + if !agent_still_has_vm { + self.agent_data_cache + .alter(&scheduled_agent_id, |_, mut v| { + v.extended_vm_set.remove(&vmid); + v + }); + } + } if entries.is_empty() { drop(entries); self.vm_data_cache.remove(&vmid); - if let Some(cached_vm) = removed_cached_vm - && let Some(mut agent) = - self.agent_data_cache.get_mut(&cached_vm.scheduled_agent_id) - { - agent.extended_vm_set.remove(&vmid); - } } } @@ -803,22 +885,38 @@ impl Message for SchedulerActor { let reply = target_agent.ask(&msg).await; - if reply.is_err() { - // A lost reply is not a rejected create. Keep the cache if the VM actor exists. - let actor_exists = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) - .await? - .is_some(); - if !actor_exists { - self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { - v.extended_vm_set.remove(&msg.vmid); - v - }); - if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) { - entries.retain(|entry| entry.actor_ref.is_some()); - if entries.is_empty() { - drop(entries); - self.vm_data_cache.remove(&msg.vmid); - } + if let Ok(reply) = &reply + && let Some(actor_id_bytes) = &reply.actor_id + && let Ok(actor_id) = ActorId::from_bytes(actor_id_bytes) + { + self.vm_actorid_ulid_map.insert(actor_id, msg.vmid); + let actor_ref = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) + .await + .ok() + .flatten(); + if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) + && let Some(entry) = entries.iter_mut().find(|entry| entry.actor_ref.is_none()) + { + entry.actor_ref = actor_ref; + } + } + + if reply.is_err() + && RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) + .await + .ok() + .flatten() + .is_none() + { + self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { + v.extended_vm_set.remove(&msg.vmid); + v + }); + if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) { + entries.retain(|entry| entry.actor_ref.is_some()); + if entries.is_empty() { + drop(entries); + self.vm_data_cache.remove(&msg.vmid); } } } diff --git a/odorobo/src/messages/vm.rs b/odorobo/src/messages/vm.rs index c48ce5c..c5d8dd8 100644 --- a/odorobo/src/messages/vm.rs +++ b/odorobo/src/messages/vm.rs @@ -29,6 +29,8 @@ pub struct CreateVM { #[derive(Serialize, Deserialize, Reply, Debug, JsonSchema)] pub struct CreateVMReply { pub config: Option, + /// Serialized ID of the VM actor created by the agent. + pub actor_id: Option>, } /// Message to delete a VM's config from the agent, shutting it down From 4a972aa7d73b448b4b77b799b70ed5f54483fa39 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 17:19:42 -0700 Subject: [PATCH 23/35] fix review comments and remove unnecessary exception --- odorobo/src/actors/scheduler_actor.rs | 43 ++++++++++++++------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 94f6d78..3d5b26f 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -99,6 +99,25 @@ static VCPU_OVERPROVISIONMENT_NUMERATOR: u32 = 2; static VCPU_OVERPROVISIONMENT_DENOMINATOR: u32 = 1; impl SchedulerActor { + fn update_cached_vm_entry( + entries: &mut Vec, + actor_id: ActorId, + cached_vm: CachedVMActor, + ) { + if let Some(entry) = entries.iter_mut().find(|entry| { + entry + .actor_ref + .as_ref() + .is_some_and(|actor| actor.id() == actor_id) + }) { + *entry = cached_vm; + } else if let Some(entry) = entries.iter_mut().find(|entry| entry.actor_ref.is_none()) { + *entry = cached_vm; + } else { + entries.push(cached_vm); + } + } + #[expect(dead_code, reason = "reserved for explicit placement by actor id")] fn lookup_agent_by_actor_id(&self, actor_id: &ActorId) -> Option> { self.agent_data_cache @@ -157,10 +176,6 @@ impl SchedulerActor { Ok(()) } - #[expect( - clippy::too_many_lines, - reason = "the updater keeps polling and performs coordinated cache cleanup" - )] async fn vm_updater_task( actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, @@ -195,20 +210,7 @@ impl SchedulerActor { data_cache .entry(vmid) .and_modify(|entries| { - if let Some(entry) = entries.iter_mut().find(|entry| { - entry - .actor_ref - .as_ref() - .is_some_and(|actor| actor.id() == actor_ref.id()) - }) { - *entry = cached_vm.clone(); - } else if let Some(entry) = - entries.iter_mut().find(|entry| entry.actor_ref.is_none()) - { - *entry = cached_vm.clone(); - } else { - entries.push(cached_vm.clone()); - } + Self::update_cached_vm_entry(entries, actor_ref.id(), cached_vm.clone()); }) .or_insert_with(|| vec![cached_vm]); @@ -330,8 +332,9 @@ impl SchedulerActor { data_cache.alter(&actor_ref.id(), |_, mut v| { v.data = data; - v.extended_vm_set.retain(|vmid| v.data.vms.contains(vmid)); - v.extended_vm_set.extend(v.data.vms.iter()); + // Keep optimistic reservations until their VM lifecycle reports that they + // no longer exist. Status updates can race with CreateVM and arrive first. + v.extended_vm_set.extend(v.data.vms.iter().copied()); v }); From a415ef29d8c5e7c1b8bb9c9fd77bb05135fa3a59 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 17:31:09 -0700 Subject: [PATCH 24/35] make sure stuff in the cache actually gets run --- odorobo/src/actors/scheduler_actor.rs | 84 ++++++++++++++++++++++++++- 1 file changed, 82 insertions(+), 2 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 3d5b26f..048215b 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,7 +1,7 @@ use std::cmp::Ordering; use std::ops::ControlFlow; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; @@ -48,6 +48,7 @@ pub struct CachedVMActor { /// the agent's `extended_vm_set` if the VM actor dies before it's linked. pub scheduled_agent_id: ActorId, pub data: GetVMInfoReply, + pub cached_at: Instant, } // todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. @@ -97,6 +98,7 @@ pub struct SchedulerActor { // todo: this might need to be a runtime thing but this makes it easy to write for now and could easily be switched out later. static VCPU_OVERPROVISIONMENT_NUMERATOR: u32 = 2; static VCPU_OVERPROVISIONMENT_DENOMINATOR: u32 = 1; +const UNRESOLVED_VM_CACHE_TIMEOUT: Duration = Duration::from_secs(30); impl SchedulerActor { fn update_cached_vm_entry( @@ -118,6 +120,47 @@ impl SchedulerActor { } } + fn cleanup_unresolved_vm_cache( + data_cache: &DashMap>, + agent_data_cache: &DashMap, + ) { + let now = Instant::now(); + let vmids: Vec<_> = data_cache.iter().map(|entry| *entry.key()).collect(); + + for vmid in vmids { + let mut removed_agent_ids = Vec::new(); + let cache_is_empty = data_cache.get_mut(&vmid).map_or(false, |mut entries| { + entries.retain(|entry| { + let expired = entry.actor_ref.is_none() + && now.duration_since(entry.cached_at) >= UNRESOLVED_VM_CACHE_TIMEOUT; + if expired { + removed_agent_ids.push(entry.scheduled_agent_id); + } + !expired + }); + entries.is_empty() + }); + + if cache_is_empty { + data_cache.remove(&vmid); + } + + for agent_id in removed_agent_ids { + let agent_still_has_vm = data_cache.get(&vmid).is_some_and(|entries| { + entries + .iter() + .any(|entry| entry.scheduled_agent_id == agent_id) + }); + if !agent_still_has_vm { + agent_data_cache.alter(&agent_id, |_, mut agent| { + agent.extended_vm_set.remove(&vmid); + agent + }); + } + } + } + } + #[expect(dead_code, reason = "reserved for explicit placement by actor id")] fn lookup_agent_by_actor_id(&self, actor_id: &ActorId) -> Option> { self.agent_data_cache @@ -206,6 +249,7 @@ impl SchedulerActor { actor_ref: Some(actor_ref.clone()), scheduled_agent_id, data, + cached_at: Instant::now(), }; data_cache .entry(vmid) @@ -401,6 +445,11 @@ impl SchedulerActor { warn!(?error, "agent actor discovery failed"); } + Self::cleanup_unresolved_vm_cache( + &vm_data_cache_arc_clone, + &agent_data_cache_arc_clone, + ); + //info!(?vm_data_cache_arc_clone); //info!(?agent_data_cache_arc_clone); @@ -635,12 +684,19 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq #[cfg(test)] mod tests { - use super::{affinity_delta, evaluate_affinity_rule, evaluate_table_value, has_capacity}; + use super::{ + CachedVMActor, SchedulerActor, affinity_delta, evaluate_affinity_rule, + evaluate_table_value, has_capacity, + }; + use crate::messages::vm::GetVMInfoReply; use crate::types::{ AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, MetadataTable, ObjectMetadata, Operator, }; + use dashmap::DashMap; use std::collections::BTreeMap; + use std::time::{Duration, Instant}; + use ulid::Ulid; fn requirement(operator: Operator, values: &[&str]) -> AffinityRequirement { AffinityRequirement { @@ -651,6 +707,29 @@ mod tests { } } + #[test] + fn removes_expired_unresolved_vm_placeholders() { + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let agent_id = super::ActorId::new(1); + let data_cache = DashMap::new(); + data_cache.insert( + vmid, + vec![CachedVMActor { + actor_ref: None, + scheduled_agent_id: agent_id, + data: GetVMInfoReply { vmid, config: None }, + cached_at: Instant::now() + .checked_sub(Duration::from_secs(31)) + .expect("test timestamp should be representable"), + }], + ); + let agent_data_cache = DashMap::new(); + + SchedulerActor::cleanup_unresolved_vm_cache(&data_cache, &agent_data_cache); + + assert!(!data_cache.contains_key(&vmid)); + } + #[test] fn evaluates_membership_and_missing_keys() { let metadata = BTreeMap::from([("tier".to_owned(), "frontend".to_owned())]); @@ -884,6 +963,7 @@ impl Message for SchedulerActor { vmid: msg.vmid, config: Some(msg.config.clone()), }, + cached_at: Instant::now(), }); let reply = target_agent.ask(&msg).await; From 3b3f4211702d512d4f93ba4d9288b398a28ef107 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 17:52:08 -0700 Subject: [PATCH 25/35] more issues fixed - Pending VM CPU/RAM is now reserved immediately, preventing back-to-back creates from oversubscribing an agent. - Exact-fit capacity is accepted by using `<=`. - VM cache ownership is optional for restart-discovered VMs instead of incorrectly using the VM actor ID as an agent ID. - Unconfirmed VM cache entries expire even if an actor reference was briefly found, fixing the create/stop discovery leak. - Agent affinity state is rebuilt from current status plus active scheduler reservations, preventing stale VM IDs while preserving create-race safety. - Removed the hard-coded required affinity rule from `odoroboctl create`, allowing creation in an empty cluster. - Fixed the clippy warning and added reservation/capacity regression tests. --- odorobo/src/actors/scheduler_actor.rs | 189 ++++++++++++++++++-------- odoroboctl/src/cli.rs | 26 +--- 2 files changed, 132 insertions(+), 83 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 048215b..99d2992 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -46,9 +46,11 @@ pub struct CachedVMActor { pub actor_ref: Option>, /// The agent this VM was scheduled on. Used by the updater task to clean up /// the agent's `extended_vm_set` if the VM actor dies before it's linked. - pub scheduled_agent_id: ActorId, + pub scheduled_agent_id: Option, pub data: GetVMInfoReply, pub cached_at: Instant, + /// False until a VM updater has successfully polled this actor. + pub updater_confirmed: bool, } // todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. @@ -129,12 +131,12 @@ impl SchedulerActor { for vmid in vmids { let mut removed_agent_ids = Vec::new(); - let cache_is_empty = data_cache.get_mut(&vmid).map_or(false, |mut entries| { + let cache_is_empty = data_cache.get_mut(&vmid).is_some_and(|mut entries| { entries.retain(|entry| { - let expired = entry.actor_ref.is_none() + let expired = !entry.updater_confirmed && now.duration_since(entry.cached_at) >= UNRESOLVED_VM_CACHE_TIMEOUT; - if expired { - removed_agent_ids.push(entry.scheduled_agent_id); + if expired && let Some(agent_id) = entry.scheduled_agent_id { + removed_agent_ids.push(agent_id); } !expired }); @@ -149,7 +151,7 @@ impl SchedulerActor { let agent_still_has_vm = data_cache.get(&vmid).is_some_and(|entries| { entries .iter() - .any(|entry| entry.scheduled_agent_id == agent_id) + .any(|entry| entry.scheduled_agent_id == Some(agent_id)) }); if !agent_still_has_vm { agent_data_cache.alter(&agent_id, |_, mut agent| { @@ -233,23 +235,22 @@ impl SchedulerActor { vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that - let scheduled_agent_id = data_cache - .get(&vmid) - .and_then(|entries| { - entries.iter().find_map(|entry| { - entry - .actor_ref - .as_ref() - .is_none_or(|actor| actor.id() == actor_ref.id()) - .then_some(entry.scheduled_agent_id) - }) + let scheduled_agent_id = data_cache.get(&vmid).and_then(|entries| { + entries.iter().find_map(|entry| { + entry + .actor_ref + .as_ref() + .is_none_or(|actor| actor.id() == actor_ref.id()) + .then_some(entry.scheduled_agent_id) + .flatten() }) - .unwrap_or_else(|| actor_ref.id()); + }); let cached_vm = CachedVMActor { actor_ref: Some(actor_ref.clone()), scheduled_agent_id, data, cached_at: Instant::now(), + updater_confirmed: true, }; data_cache .entry(vmid) @@ -303,12 +304,13 @@ impl SchedulerActor { }) .map(|index| (index, entries.remove(index))) }); - if let Some(cached_vm) = removed_cached_vm { - let scheduled_agent_id = cached_vm.1.scheduled_agent_id; + if let Some(cached_vm) = removed_cached_vm + && let Some(scheduled_agent_id) = cached_vm.1.scheduled_agent_id + { let agent_still_has_vm = data_cache.get(&vmid).is_some_and(|entries| { entries .iter() - .any(|entry| entry.scheduled_agent_id == scheduled_agent_id) + .any(|entry| entry.scheduled_agent_id == Some(scheduled_agent_id)) }); if !agent_still_has_vm { agent_data_cache.alter(&scheduled_agent_id, |_, mut v| { @@ -335,6 +337,7 @@ impl SchedulerActor { async fn agent_actor_finder( parent_actor_ref: RemoteActorRef, data_cache: Arc>, + vm_data_cache: Arc>>, keepalive_tasks: Arc>>, ) -> Result<(), Report> { trace!("running agent_actor_finder"); @@ -353,8 +356,10 @@ impl SchedulerActor { parent_actor_ref.link_remote(&agent_actor).await?; let data_cache_clone = Arc::clone(&data_cache); + let vm_data_cache_clone = Arc::clone(&vm_data_cache); let updater_task = tokio::spawn(async move { - Self::agent_updater_task(agent_actor, data_cache_clone).await; + Self::agent_updater_task(agent_actor, data_cache_clone, vm_data_cache_clone) + .await; }); keepalive_tasks.insert(agent_actor_id, updater_task); @@ -367,28 +372,28 @@ impl SchedulerActor { async fn agent_updater_task( actor_ref: RemoteActorRef, data_cache: Arc>, + vm_data_cache: Arc>>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails: u8 = 0; loop { if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { + let reserved_vms = reserved_vms_for_agent(&vm_data_cache, actor_ref.id()); if data_cache.contains_key(&actor_ref.id()) { data_cache.alter(&actor_ref.id(), |_, mut v| { v.data = data; - - // Keep optimistic reservations until their VM lifecycle reports that they - // no longer exist. Status updates can race with CreateVM and arrive first. - v.extended_vm_set.extend(v.data.vms.iter().copied()); - + v.extended_vm_set = + v.data.vms.iter().copied().chain(reserved_vms).collect(); v }); } else { + let extended_vm_set = data.vms.iter().copied().chain(reserved_vms).collect(); data_cache.insert( actor_ref.id(), CachedAgentActor { actor_ref: actor_ref.clone(), - data: data.clone(), - extended_vm_set: data.vms.iter().copied().collect(), + data, + extended_vm_set, }, ); } @@ -433,6 +438,7 @@ impl SchedulerActor { let agent_join_handle = Self::agent_actor_finder( actor_ref.clone(), Arc::clone(&agent_data_cache_arc_clone), + Arc::clone(&vm_data_cache_arc_clone), Arc::clone(&agent_keepalive_tasks_arc_clone), ); @@ -511,14 +517,18 @@ impl SchedulerActor { .checked_div(VCPU_OVERPROVISIONMENT_DENOMINATOR) .unwrap_or(u32::MAX); // todo: do we care about VMData.max_vcpus? - let agent_used_vcpus = agent.data.used_vcpus.saturating_add(msg.config.data.vcpus); + let (pending_vcpus, pending_ram) = + pending_resources_for_agent(&self.vm_data_cache, agent.actor_ref.id(), &agent.data.vms); + let used_vcpus = agent.data.used_vcpus.saturating_add(pending_vcpus); + let used_ram = agent.data.used_ram.as_u64().saturating_add(pending_ram); + let agent_used_vcpus = used_vcpus.saturating_add(msg.config.data.vcpus); if !has_capacity( agent_max_vcpus, - agent.data.used_vcpus, + used_vcpus, msg.config.data.vcpus, agent.data.ram.as_u64(), - agent.data.used_ram.as_u64(), + used_ram, msg.config.data.memory.as_u64(), ) { return AgentScore::REJECTED; @@ -537,13 +547,8 @@ impl SchedulerActor { // todo: add ram overprovisionment. not adding this to scheduler until it works on the hypervisor side. let agent_max_ram = agent.data.ram; - let agent_used_ram = bytesize::ByteSize::b( - agent - .data - .used_ram - .as_u64() - .saturating_add(msg.config.data.memory.as_u64()), - ); + let agent_used_ram = + bytesize::ByteSize::b(used_ram.saturating_add(msg.config.data.memory.as_u64())); #[expect( clippy::cast_precision_loss, @@ -608,8 +613,44 @@ const fn has_capacity( used_ram: u64, requested_ram: u64, ) -> bool { - used_vcpus.saturating_add(requested_vcpus) < max_vcpus - && used_ram.saturating_add(requested_ram) < max_ram + used_vcpus.saturating_add(requested_vcpus) <= max_vcpus + && used_ram.saturating_add(requested_ram) <= max_ram +} + +fn reserved_vms_for_agent( + vm_data_cache: &DashMap>, + agent_id: ActorId, +) -> Vec { + vm_data_cache + .iter() + .filter(|entry| { + entry + .iter() + .any(|vm| vm.scheduled_agent_id == Some(agent_id)) + }) + .map(|entry| *entry.key()) + .collect() +} + +fn pending_resources_for_agent( + vm_data_cache: &DashMap>, + agent_id: ActorId, + confirmed_vms: &[Ulid], +) -> (u32, u64) { + vm_data_cache + .iter() + .filter(|entry| !confirmed_vms.contains(entry.key())) + .flat_map(|entry| { + entry + .iter() + .filter(|vm| vm.scheduled_agent_id == Some(agent_id)) + .filter_map(|vm| vm.data.config.as_ref()) + .map(|config| (config.data.vcpus, config.data.memory.as_u64())) + .collect::>() + }) + .fold((0u32, 0u64), |(vcpus, ram), (vm_vcpus, vm_ram)| { + (vcpus.saturating_add(vm_vcpus), ram.saturating_add(vm_ram)) + }) } fn affinity_delta(strictness: AffinityStrictness, follows_rule: bool) -> Option { @@ -686,13 +727,14 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq mod tests { use super::{ CachedVMActor, SchedulerActor, affinity_delta, evaluate_affinity_rule, - evaluate_table_value, has_capacity, + evaluate_table_value, has_capacity, pending_resources_for_agent, reserved_vms_for_agent, }; use crate::messages::vm::GetVMInfoReply; use crate::types::{ AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, MetadataTable, - ObjectMetadata, Operator, + ObjectMetadata, Operator, VirtualMachine, }; + use bytesize::ByteSize; use dashmap::DashMap; use std::collections::BTreeMap; use std::time::{Duration, Instant}; @@ -716,11 +758,12 @@ mod tests { vmid, vec![CachedVMActor { actor_ref: None, - scheduled_agent_id: agent_id, + scheduled_agent_id: Some(agent_id), data: GetVMInfoReply { vmid, config: None }, cached_at: Instant::now() .checked_sub(Duration::from_secs(31)) .expect("test timestamp should be representable"), + updater_confirmed: false, }], ); let agent_data_cache = DashMap::new(); @@ -730,6 +773,39 @@ mod tests { assert!(!data_cache.contains_key(&vmid)); } + #[test] + fn reserves_pending_vm_resources_until_agent_status_confirms_them() { + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let agent_id = super::ActorId::new(1); + let mut config = VirtualMachine::default(); + config.data.vcpus = 4; + config.data.memory = ByteSize::gib(8); + let data_cache = DashMap::new(); + data_cache.insert( + vmid, + vec![CachedVMActor { + actor_ref: None, + scheduled_agent_id: Some(agent_id), + data: GetVMInfoReply { + vmid, + config: Some(config), + }, + cached_at: Instant::now(), + updater_confirmed: false, + }], + ); + + assert_eq!(reserved_vms_for_agent(&data_cache, agent_id), vec![vmid]); + assert_eq!( + pending_resources_for_agent(&data_cache, agent_id, &[]), + (4, ByteSize::gib(8).as_u64()) + ); + assert_eq!( + pending_resources_for_agent(&data_cache, agent_id, &[vmid]), + (0, 0) + ); + } + #[test] fn evaluates_membership_and_missing_keys() { let metadata = BTreeMap::from([("tier".to_owned(), "frontend".to_owned())]); @@ -812,8 +888,10 @@ mod tests { Some(0) ); assert!(has_capacity(8, 2, 2, 16, 4, 4)); - assert!(!has_capacity(8, 6, 2, 16, 4, 4)); - assert!(!has_capacity(8, 2, 2, 16, 12, 4)); + assert!(has_capacity(8, 6, 2, 16, 4, 4)); + assert!(has_capacity(8, 2, 2, 16, 12, 4)); + assert!(!has_capacity(8, 7, 2, 16, 4, 4)); + assert!(!has_capacity(8, 2, 2, 16, 13, 4)); } } @@ -911,11 +989,12 @@ impl Actor for SchedulerActor { .is_some_and(|actor| actor.id() == id) }) .map(|index| entries.remove(index)); - if let Some(cached_vm) = removed_cached_vm { - let scheduled_agent_id = cached_vm.scheduled_agent_id; + if let Some(cached_vm) = removed_cached_vm + && let Some(scheduled_agent_id) = cached_vm.scheduled_agent_id + { let agent_still_has_vm = entries .iter() - .any(|entry| entry.scheduled_agent_id == scheduled_agent_id); + .any(|entry| entry.scheduled_agent_id == Some(scheduled_agent_id)); if !agent_still_has_vm { self.agent_data_cache .alter(&scheduled_agent_id, |_, mut v| { @@ -958,12 +1037,13 @@ impl Message for SchedulerActor { .or_default() .push(CachedVMActor { actor_ref: None, - scheduled_agent_id: target_agent.id(), + scheduled_agent_id: Some(target_agent.id()), data: GetVMInfoReply { vmid: msg.vmid, config: Some(msg.config.clone()), }, cached_at: Instant::now(), + updater_confirmed: false, }); let reply = target_agent.ask(&msg).await; @@ -973,15 +1053,6 @@ impl Message for SchedulerActor { && let Ok(actor_id) = ActorId::from_bytes(actor_id_bytes) { self.vm_actorid_ulid_map.insert(actor_id, msg.vmid); - let actor_ref = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) - .await - .ok() - .flatten(); - if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) - && let Some(entry) = entries.iter_mut().find(|entry| entry.actor_ref.is_none()) - { - entry.actor_ref = actor_ref; - } } if reply.is_err() @@ -996,7 +1067,7 @@ impl Message for SchedulerActor { v }); if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) { - entries.retain(|entry| entry.actor_ref.is_some()); + entries.retain(|entry| entry.updater_confirmed); if entries.is_empty() { drop(entries); self.vm_data_cache.remove(&msg.vmid); diff --git a/odoroboctl/src/cli.rs b/odoroboctl/src/cli.rs index a072d64..902f7ce 100644 --- a/odoroboctl/src/cli.rs +++ b/odoroboctl/src/cli.rs @@ -1,11 +1,6 @@ -use std::collections::BTreeMap; - use bytesize::ByteSize; use clap::{Parser, Subcommand}; -use odorobo::types::{ - AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, CreateVMRequest, - MetadataTable, ObjectMetadata, Operator, VMData, VirtualMachine, -}; +use odorobo::types::{CreateVMRequest, VMData, VirtualMachine}; use reqwest::{Client, Response}; use serde::Deserialize; use stable_eyre::Result; @@ -116,24 +111,7 @@ pub async fn run_command(cli: Cli) -> Result<()> { image, ..Default::default() }, - metadata: Some(ObjectMetadata { - labels: BTreeMap::new(), - annotations: BTreeMap::from([( - String::from("distribution"), - String::from("different"), - )]), - }), - affinity: Some(vec![AffinityRule { - strictness: AffinityStrictness::Required, - affinity_type: AffinityType::VirtualMachine, - inverse: false, - requirements: vec![AffinityRequirement { - key: String::from("distribution"), - table: MetadataTable::Annotation, - operator: Operator::NotIn, - values: vec![String::from("different")], - }], - }]), + ..Default::default() }; From 46bce53b0928f598c7801ea6803231895877a521 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Mon, 17 Aug 2026 20:21:19 -0500 Subject: [PATCH 26/35] remove extended vm set and replace with scheduler authoritative vm placement map --- odorobo/src/actors/scheduler_actor.rs | 359 +++++++++----------------- 1 file changed, 120 insertions(+), 239 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 99d2992..be769b3 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -17,6 +17,7 @@ use crate::types::AffinityType; use crate::types::MetadataTable; use crate::types::ObjectMetadata; use crate::types::Operator; +use crate::types::VirtualMachine; use crate::utils::actor_names::AGENT; use crate::utils::actor_names::VM; use crate::utils::actor_names::vm_actor_id; @@ -35,27 +36,32 @@ use ulid::Ulid; #[derive(Debug, Clone)] pub struct CachedAgentActor { pub actor_ref: RemoteActorRef, + /// Latest observation reported by the agent. This is not scheduler placement state. pub data: AgentStatus, - /// this is a set of all VMs that may be on an agent. it is used for rules such as affinity to make sure we don't schedule things in ways that arent allowed - /// We don't know for a fact these VMs are scheduled due to latency and boot up delay, but they may be scheduled. - pub extended_vm_set: AHashSet, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VmLifecycle { + Pending, + Running, +} + +#[derive(Debug, Clone)] +pub struct VmPlacement { + pub agent_id: ActorId, + pub config: VirtualMachine, + pub lifecycle: VmLifecycle, + pub created_at: Instant, + pub last_confirmed_at: Option, } #[derive(Debug, Clone)] pub struct CachedVMActor { pub actor_ref: Option>, - /// The agent this VM was scheduled on. Used by the updater task to clean up - /// the agent's `extended_vm_set` if the VM actor dies before it's linked. - pub scheduled_agent_id: Option, pub data: GetVMInfoReply, pub cached_at: Instant, - /// False until a VM updater has successfully polled this actor. - pub updater_confirmed: bool, } -// todo: i dont like the way this cache is setup. I think we may need to change it later, but it is hard to figure out what the optimal solution is without doing it at least once. -// especially when we haven't fully made decisions about some other things. -// // todo: we should improve the cache to not have agents and vms send the full data on every update. // I looked at kameo streams to make this better, but they aren't really intended for this kind of long term update use case. // They use rust futures::stream which seems to be more intended for you have an iterator for example that will create data, but not like full on sending messages. @@ -68,29 +74,9 @@ pub struct SchedulerActor { pub agent_data_cache: Arc>, pub agent_keepalive_tasks: Arc>>, - // todo: we might need a better way to store this. - // we 100% need a way to store vms even if we don't know their actorid (ex: actor hasn't been started or is shutdown) - // we also might want to be able to store them without a ulid, possibly - // so we might need a vec of vms and then to just store maps/indexes of actorid and ulid to vector index - // and then like a freelist or something. - // i dont really love that option either though cause it feels overkill. - // maybe we sure just be using a proper database entirely? - // idk. will figure it out later. - // - // new related problem: i just realized vmid, actorid pairs dont have to be unique. - // if a vm is migrating from one actor to another, there might be two actors with the same vmid. - // - // additional context (05/05/2026): we almost may need a way to store them without a ulid, due to how CH migration works. - // the question becomes if we want to abstract CH migration away entirely from the scheduler. - // we could also possibly ignore it for the non-HA scheduler. - // I (caleb) want to ask cappy (and possibly Lea) about these problems. - // - // the best solution for at least some of this is almost certainly having an external reliable DB (such as etcd) to store some of these things permanently. - // we will need that specifically for what VMs are supposed to be running, because if a large percentage of the cluster goes down, including the manager, we need a way to recover. - // and i dont think leaving that on dashboard which could have high latency is a good idea. - // alternatively we could have the other manager nodes try to keep track of that data, but i think we are going to run into issues with keeping the state consistent between all nodes. - // we may need to make some architecture designs about db consistency vs uptime vs speed in that situation, and im not doing that on my own. pub vm_actorid_ulid_map: Arc>, + pub vm_placements: Arc>, + /// this is a vec because a vmid/ulid can be scheduled on multiple boxes simultaneously during migration pub vm_data_cache: Arc>>, pub vm_keepalive_tasks: Arc>>, @@ -123,44 +109,48 @@ impl SchedulerActor { } fn cleanup_unresolved_vm_cache( + placements: &DashMap, data_cache: &DashMap>, - agent_data_cache: &DashMap, ) { let now = Instant::now(); - let vmids: Vec<_> = data_cache.iter().map(|entry| *entry.key()).collect(); - - for vmid in vmids { - let mut removed_agent_ids = Vec::new(); - let cache_is_empty = data_cache.get_mut(&vmid).is_some_and(|mut entries| { - entries.retain(|entry| { - let expired = !entry.updater_confirmed - && now.duration_since(entry.cached_at) >= UNRESOLVED_VM_CACHE_TIMEOUT; - if expired && let Some(agent_id) = entry.scheduled_agent_id { - removed_agent_ids.push(agent_id); - } - !expired - }); - entries.is_empty() - }); + let expired: Vec<_> = placements + .iter() + .filter(|entry| { + entry.lifecycle == VmLifecycle::Pending + && now.duration_since(entry.created_at) >= UNRESOLVED_VM_CACHE_TIMEOUT + }) + .map(|entry| *entry.key()) + .collect(); + + for vmid in expired { + Self::remove_vm_state(vmid, placements, data_cache); + } + } - if cache_is_empty { - data_cache.remove(&vmid); - } + /// takes a list of observed vmids and gives you a set of every vmid that could be on an agent. + fn placement_vm_ids( + placements: &DashMap, + agent_id: ActorId, + observed: &[Ulid], + ) -> AHashSet { + observed + .iter() + .copied() + .chain( + placements + .iter() + .filter_map(|entry| (entry.agent_id == agent_id).then_some(*entry.key())), + ) + .collect() + } - for agent_id in removed_agent_ids { - let agent_still_has_vm = data_cache.get(&vmid).is_some_and(|entries| { - entries - .iter() - .any(|entry| entry.scheduled_agent_id == Some(agent_id)) - }); - if !agent_still_has_vm { - agent_data_cache.alter(&agent_id, |_, mut agent| { - agent.extended_vm_set.remove(&vmid); - agent - }); - } - } - } + fn remove_vm_state( + vmid: Ulid, + placements: &DashMap, + data_cache: &DashMap>, + ) { + placements.remove(&vmid); + data_cache.remove(&vmid); } #[expect(dead_code, reason = "reserved for explicit placement by actor id")] @@ -184,7 +174,7 @@ impl SchedulerActor { vm_actorid_ulid_map: Arc>, data_cache: Arc>>, keepalive_tasks: Arc>>, - agent_data_cache: Arc>, + placements: Arc>, ) -> Result<(), Report> { trace!("running vm_actor_finder"); @@ -203,13 +193,13 @@ impl SchedulerActor { let vm_actorid_ulid_map_clone = Arc::clone(&vm_actorid_ulid_map); let data_cache_clone = Arc::clone(&data_cache); - let agent_data_cache_clone = Arc::clone(&agent_data_cache); + let placements_clone = Arc::clone(&placements); let updater_task = tokio::spawn(async move { Self::vm_updater_task( vm_actor, vm_actorid_ulid_map_clone, data_cache_clone, - agent_data_cache_clone, + placements_clone, ) .await; }); @@ -225,7 +215,7 @@ impl SchedulerActor { actor_ref: RemoteActorRef, vm_actorid_ulid_map: Arc>, data_cache: Arc>>, - agent_data_cache: Arc>, + placements: Arc>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails: u8 = 0; @@ -235,22 +225,14 @@ impl SchedulerActor { vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that - let scheduled_agent_id = data_cache.get(&vmid).and_then(|entries| { - entries.iter().find_map(|entry| { - entry - .actor_ref - .as_ref() - .is_none_or(|actor| actor.id() == actor_ref.id()) - .then_some(entry.scheduled_agent_id) - .flatten() - }) - }); + if let Some(mut placement) = placements.get_mut(&vmid) { + placement.lifecycle = VmLifecycle::Running; + placement.last_confirmed_at = Some(Instant::now()); + } let cached_vm = CachedVMActor { actor_ref: Some(actor_ref.clone()), - scheduled_agent_id, data, cached_at: Instant::now(), - updater_confirmed: true, }; data_cache .entry(vmid) @@ -293,38 +275,7 @@ impl SchedulerActor { }); if let Some(vmid) = vmid { - let removed_cached_vm = data_cache.get_mut(&vmid).and_then(|mut entries| { - entries - .iter() - .position(|entry| { - entry - .actor_ref - .as_ref() - .is_some_and(|actor| actor.id() == actor_ref.id()) - }) - .map(|index| (index, entries.remove(index))) - }); - if let Some(cached_vm) = removed_cached_vm - && let Some(scheduled_agent_id) = cached_vm.1.scheduled_agent_id - { - let agent_still_has_vm = data_cache.get(&vmid).is_some_and(|entries| { - entries - .iter() - .any(|entry| entry.scheduled_agent_id == Some(scheduled_agent_id)) - }); - if !agent_still_has_vm { - agent_data_cache.alter(&scheduled_agent_id, |_, mut v| { - v.extended_vm_set.remove(&vmid); - v - }); - } - } - if data_cache - .get(&vmid) - .is_some_and(|entries| entries.is_empty()) - { - data_cache.remove(&vmid); - } + Self::remove_vm_state(vmid, &placements, &data_cache); } return; @@ -337,7 +288,6 @@ impl SchedulerActor { async fn agent_actor_finder( parent_actor_ref: RemoteActorRef, data_cache: Arc>, - vm_data_cache: Arc>>, keepalive_tasks: Arc>>, ) -> Result<(), Report> { trace!("running agent_actor_finder"); @@ -356,10 +306,8 @@ impl SchedulerActor { parent_actor_ref.link_remote(&agent_actor).await?; let data_cache_clone = Arc::clone(&data_cache); - let vm_data_cache_clone = Arc::clone(&vm_data_cache); let updater_task = tokio::spawn(async move { - Self::agent_updater_task(agent_actor, data_cache_clone, vm_data_cache_clone) - .await; + Self::agent_updater_task(agent_actor, data_cache_clone).await; }); keepalive_tasks.insert(agent_actor_id, updater_task); @@ -372,28 +320,22 @@ impl SchedulerActor { async fn agent_updater_task( actor_ref: RemoteActorRef, data_cache: Arc>, - vm_data_cache: Arc>>, ) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails: u8 = 0; loop { if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { - let reserved_vms = reserved_vms_for_agent(&vm_data_cache, actor_ref.id()); if data_cache.contains_key(&actor_ref.id()) { data_cache.alter(&actor_ref.id(), |_, mut v| { v.data = data; - v.extended_vm_set = - v.data.vms.iter().copied().chain(reserved_vms).collect(); v }); } else { - let extended_vm_set = data.vms.iter().copied().chain(reserved_vms).collect(); data_cache.insert( actor_ref.id(), CachedAgentActor { actor_ref: actor_ref.clone(), data, - extended_vm_set, }, ); } @@ -422,6 +364,7 @@ impl SchedulerActor { let vm_actorid_ulid_map_arc_clone = Arc::clone(&self.vm_actorid_ulid_map); let vm_data_cache_arc_clone = Arc::clone(&self.vm_data_cache); + let vm_placements_arc_clone = Arc::clone(&self.vm_placements); let vm_keepalive_tasks_arc_clone = Arc::clone(&self.vm_keepalive_tasks); self.cache_actor_finder = Some(tokio::spawn(async move { @@ -432,13 +375,12 @@ impl SchedulerActor { Arc::clone(&vm_actorid_ulid_map_arc_clone), Arc::clone(&vm_data_cache_arc_clone), Arc::clone(&vm_keepalive_tasks_arc_clone), - Arc::clone(&agent_data_cache_arc_clone), + Arc::clone(&vm_placements_arc_clone), ); let agent_join_handle = Self::agent_actor_finder( actor_ref.clone(), Arc::clone(&agent_data_cache_arc_clone), - Arc::clone(&vm_data_cache_arc_clone), Arc::clone(&agent_keepalive_tasks_arc_clone), ); @@ -452,8 +394,8 @@ impl SchedulerActor { } Self::cleanup_unresolved_vm_cache( + &vm_placements_arc_clone, &vm_data_cache_arc_clone, - &agent_data_cache_arc_clone, ); //info!(?vm_data_cache_arc_clone); @@ -518,7 +460,7 @@ impl SchedulerActor { .unwrap_or(u32::MAX); // todo: do we care about VMData.max_vcpus? let (pending_vcpus, pending_ram) = - pending_resources_for_agent(&self.vm_data_cache, agent.actor_ref.id(), &agent.data.vms); + pending_resources_for_agent(&self.vm_placements, agent.actor_ref.id()); let used_vcpus = agent.data.used_vcpus.saturating_add(pending_vcpus); let used_ram = agent.data.used_ram.as_u64().saturating_add(pending_ram); let agent_used_vcpus = used_vcpus.saturating_add(msg.config.data.vcpus); @@ -567,8 +509,13 @@ impl SchedulerActor { match rule.affinity_type { AffinityType::VirtualMachine => { - for vmid in &agent.extended_vm_set { - let Some(vm_data_cache_refs) = self.vm_data_cache.get(vmid) else { + let vmids = Self::placement_vm_ids( + &self.vm_placements, + agent.actor_ref.id(), + &agent.data.vms, + ); + for vmid in vmids { + let Some(vm_data_cache_refs) = self.vm_data_cache.get(&vmid) else { continue; }; @@ -617,37 +564,14 @@ const fn has_capacity( && used_ram.saturating_add(requested_ram) <= max_ram } -fn reserved_vms_for_agent( - vm_data_cache: &DashMap>, - agent_id: ActorId, -) -> Vec { - vm_data_cache - .iter() - .filter(|entry| { - entry - .iter() - .any(|vm| vm.scheduled_agent_id == Some(agent_id)) - }) - .map(|entry| *entry.key()) - .collect() -} - fn pending_resources_for_agent( - vm_data_cache: &DashMap>, + placements: &DashMap, agent_id: ActorId, - confirmed_vms: &[Ulid], ) -> (u32, u64) { - vm_data_cache + placements .iter() - .filter(|entry| !confirmed_vms.contains(entry.key())) - .flat_map(|entry| { - entry - .iter() - .filter(|vm| vm.scheduled_agent_id == Some(agent_id)) - .filter_map(|vm| vm.data.config.as_ref()) - .map(|config| (config.data.vcpus, config.data.memory.as_u64())) - .collect::>() - }) + .filter(|entry| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) + .map(|entry| (entry.config.data.vcpus, entry.config.data.memory.as_u64())) .fold((0u32, 0u64), |(vcpus, ram), (vm_vcpus, vm_ram)| { (vcpus.saturating_add(vm_vcpus), ram.saturating_add(vm_ram)) }) @@ -726,10 +650,10 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq #[cfg(test)] mod tests { use super::{ - CachedVMActor, SchedulerActor, affinity_delta, evaluate_affinity_rule, - evaluate_table_value, has_capacity, pending_resources_for_agent, reserved_vms_for_agent, + CachedVMActor, SchedulerActor, VmLifecycle, VmPlacement, affinity_delta, + evaluate_affinity_rule, evaluate_table_value, has_capacity, pending_resources_for_agent, }; - use crate::messages::vm::GetVMInfoReply; + use crate::types::{ AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, MetadataTable, ObjectMetadata, Operator, VirtualMachine, @@ -753,23 +677,24 @@ mod tests { fn removes_expired_unresolved_vm_placeholders() { let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); let agent_id = super::ActorId::new(1); - let data_cache = DashMap::new(); - data_cache.insert( + let placements: DashMap = DashMap::new(); + placements.insert( vmid, - vec![CachedVMActor { - actor_ref: None, - scheduled_agent_id: Some(agent_id), - data: GetVMInfoReply { vmid, config: None }, - cached_at: Instant::now() + VmPlacement { + agent_id, + config: VirtualMachine::default(), + lifecycle: VmLifecycle::Pending, + created_at: Instant::now() .checked_sub(Duration::from_secs(31)) .expect("test timestamp should be representable"), - updater_confirmed: false, - }], + last_confirmed_at: None, + }, ); - let agent_data_cache = DashMap::new(); + let data_cache: DashMap> = DashMap::new(); - SchedulerActor::cleanup_unresolved_vm_cache(&data_cache, &agent_data_cache); + SchedulerActor::cleanup_unresolved_vm_cache(&placements, &data_cache); + assert!(!placements.contains_key(&vmid)); assert!(!data_cache.contains_key(&vmid)); } @@ -780,30 +705,22 @@ mod tests { let mut config = VirtualMachine::default(); config.data.vcpus = 4; config.data.memory = ByteSize::gib(8); - let data_cache = DashMap::new(); - data_cache.insert( + let placements = DashMap::new(); + placements.insert( vmid, - vec![CachedVMActor { - actor_ref: None, - scheduled_agent_id: Some(agent_id), - data: GetVMInfoReply { - vmid, - config: Some(config), - }, - cached_at: Instant::now(), - updater_confirmed: false, - }], + VmPlacement { + agent_id, + config, + lifecycle: VmLifecycle::Pending, + created_at: Instant::now(), + last_confirmed_at: None, + }, ); - assert_eq!(reserved_vms_for_agent(&data_cache, agent_id), vec![vmid]); assert_eq!( - pending_resources_for_agent(&data_cache, agent_id, &[]), + pending_resources_for_agent(&placements, agent_id), (4, ByteSize::gib(8).as_u64()) ); - assert_eq!( - pending_resources_for_agent(&data_cache, agent_id, &[vmid]), - (0, 0) - ); } #[test] @@ -942,6 +859,7 @@ impl Actor for SchedulerActor { agent_data_cache: Arc::new(DashMap::new()), agent_keepalive_tasks: Arc::new(DashMap::new()), vm_actorid_ulid_map: Arc::new(DashMap::new()), + vm_placements: Arc::new(DashMap::new()), vm_data_cache: Arc::new(DashMap::new()), vm_keepalive_tasks: Arc::new(DashMap::new()), cache_actor_finder: None, @@ -977,36 +895,8 @@ impl Actor for SchedulerActor { keepalive_task.abort(); } - if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&id) - && let Some(mut entries) = self.vm_data_cache.get_mut(&vmid) - { - let removed_cached_vm = entries - .iter() - .position(|entry| { - entry - .actor_ref - .as_ref() - .is_some_and(|actor| actor.id() == id) - }) - .map(|index| entries.remove(index)); - if let Some(cached_vm) = removed_cached_vm - && let Some(scheduled_agent_id) = cached_vm.scheduled_agent_id - { - let agent_still_has_vm = entries - .iter() - .any(|entry| entry.scheduled_agent_id == Some(scheduled_agent_id)); - if !agent_still_has_vm { - self.agent_data_cache - .alter(&scheduled_agent_id, |_, mut v| { - v.extended_vm_set.remove(&vmid); - v - }); - } - } - if entries.is_empty() { - drop(entries); - self.vm_data_cache.remove(&vmid); - } + if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&id) { + Self::remove_vm_state(vmid, &self.vm_placements, &self.vm_data_cache); } // todo: attempt vm restarts if necessary. @@ -1025,25 +915,26 @@ impl Message for SchedulerActor { ) -> Self::Reply { let target_agent = self.schedule_agent(&msg)?; - // we add to cache first, because we want to make sure future requests assume this vm exists. if the message fails, we clean it up afterward. - if let Some(mut cached_data) = self.agent_data_cache.get_mut(&target_agent.id()) { - cached_data.extended_vm_set.insert(msg.vmid); - } else { - return Err(eyre!("target agent is not in data cache")); - } - + self.vm_placements.insert( + msg.vmid, + VmPlacement { + agent_id: target_agent.id(), + config: msg.config.clone(), + lifecycle: VmLifecycle::Pending, + created_at: Instant::now(), + last_confirmed_at: None, + }, + ); self.vm_data_cache .entry(msg.vmid) .or_default() .push(CachedVMActor { actor_ref: None, - scheduled_agent_id: Some(target_agent.id()), data: GetVMInfoReply { vmid: msg.vmid, config: Some(msg.config.clone()), }, cached_at: Instant::now(), - updater_confirmed: false, }); let reply = target_agent.ask(&msg).await; @@ -1062,17 +953,7 @@ impl Message for SchedulerActor { .flatten() .is_none() { - self.agent_data_cache.alter(&target_agent.id(), |_, mut v| { - v.extended_vm_set.remove(&msg.vmid); - v - }); - if let Some(mut entries) = self.vm_data_cache.get_mut(&msg.vmid) { - entries.retain(|entry| entry.updater_confirmed); - if entries.is_empty() { - drop(entries); - self.vm_data_cache.remove(&msg.vmid); - } - } + Self::remove_vm_state(msg.vmid, &self.vm_placements, &self.vm_data_cache); } Ok(reply?) From 72c22576ac480a66bc7aff469dac230301ce14b9 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Mon, 17 Aug 2026 21:13:42 -0500 Subject: [PATCH 27/35] move dashmap accesses to messages due to race condition issues --- odorobo/src/actors/scheduler_actor.rs | 479 +++++++++++++++----------- 1 file changed, 275 insertions(+), 204 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index be769b3..333d537 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -1,6 +1,7 @@ use std::cmp::Ordering; + use std::ops::ControlFlow; -use std::sync::Arc; + use std::time::{Duration, Instant}; use crate::actors::agent_actor::AgentActor; @@ -21,9 +22,7 @@ use crate::types::VirtualMachine; use crate::utils::actor_names::AGENT; use crate::utils::actor_names::VM; use crate::utils::actor_names::vm_actor_id; -use ahash::AHashSet; -use dashmap::DashMap; -use dashmap::mapref::multiple::RefMulti; +use ahash::{AHashMap, AHashSet}; use kameo::prelude::*; use libp2p::futures::TryStreamExt; use stable_eyre::eyre::OptionExt; @@ -33,10 +32,45 @@ use tracing::trace; use tracing::{info, warn}; use ulid::Ulid; +#[derive(Debug)] +struct VmActorDiscovered { + actor_ref: RemoteActorRef, +} + +#[derive(Debug)] +struct AgentActorDiscovered { + actor_ref: RemoteActorRef, +} + +#[derive(Debug)] +struct VmUpdated { + actor_id: ActorId, + data: GetVMInfoReply, +} + +#[derive(Debug)] +struct VmUpdaterStopped { + actor_id: ActorId, +} + +#[derive(Debug)] +struct AgentUpdated { + actor_id: ActorId, + actor_ref: RemoteActorRef, + data: AgentStatus, +} + +#[derive(Debug)] +struct AgentUpdaterStopped { + actor_id: ActorId, +} + +#[derive(Debug)] +struct ReconcileVmPlacements; + #[derive(Debug, Clone)] pub struct CachedAgentActor { pub actor_ref: RemoteActorRef, - /// Latest observation reported by the agent. This is not scheduler placement state. pub data: AgentStatus, } @@ -71,14 +105,14 @@ pub struct CachedVMActor { // Option 2 is easier to write and uses less compute, but uses more network bandwidth. #[derive(RemoteActor)] pub struct SchedulerActor { - pub agent_data_cache: Arc>, - pub agent_keepalive_tasks: Arc>>, + pub agent_data_cache: AHashMap, + pub agent_keepalive_tasks: AHashMap>, - pub vm_actorid_ulid_map: Arc>, - pub vm_placements: Arc>, + pub vm_actorid_ulid_map: AHashMap, + pub vm_placements: AHashMap, /// this is a vec because a vmid/ulid can be scheduled on multiple boxes simultaneously during migration - pub vm_data_cache: Arc>>, - pub vm_keepalive_tasks: Arc>>, + pub vm_data_cache: AHashMap>, + pub vm_keepalive_tasks: AHashMap>, pub cache_actor_finder: Option>, } @@ -109,17 +143,17 @@ impl SchedulerActor { } fn cleanup_unresolved_vm_cache( - placements: &DashMap, - data_cache: &DashMap>, + placements: &mut AHashMap, + data_cache: &mut AHashMap>, ) { let now = Instant::now(); let expired: Vec<_> = placements .iter() - .filter(|entry| { + .filter(|(_, entry)| { entry.lifecycle == VmLifecycle::Pending && now.duration_since(entry.created_at) >= UNRESOLVED_VM_CACHE_TIMEOUT }) - .map(|entry| *entry.key()) + .map(|(vmid, _)| *vmid) .collect(); for vmid in expired { @@ -129,7 +163,7 @@ impl SchedulerActor { /// takes a list of observed vmids and gives you a set of every vmid that could be on an agent. fn placement_vm_ids( - placements: &DashMap, + placements: &AHashMap, agent_id: ActorId, observed: &[Ulid], ) -> AHashSet { @@ -139,15 +173,15 @@ impl SchedulerActor { .chain( placements .iter() - .filter_map(|entry| (entry.agent_id == agent_id).then_some(*entry.key())), + .filter_map(|(vmid, entry)| (entry.agent_id == agent_id).then_some(*vmid)), ) .collect() } fn remove_vm_state( vmid: Ulid, - placements: &DashMap, - data_cache: &DashMap>, + placements: &mut AHashMap, + data_cache: &mut AHashMap>, ) { placements.remove(&vmid); data_cache.remove(&vmid); @@ -163,83 +197,45 @@ impl SchedulerActor { #[expect(dead_code, reason = "reserved for explicit placement by hostname")] fn lookup_agent_by_hostname(&self, hostname: &str) -> Option> { self.agent_data_cache - .iter() + .values() .find(|data| data.data.hostname == hostname) .map(|data| data.actor_ref.clone()) } - // someone should likely give caleb a firm talking to about code duplication due to this section, but things are just different enough that trying to make them one function requires usage of a lot of generics which feels even worse. so i dont know what to do. cappy please fix. i hate this. - async fn vm_actor_finder( - parent_actor_ref: RemoteActorRef, - vm_actorid_ulid_map: Arc>, - data_cache: Arc>>, - keepalive_tasks: Arc>>, - placements: Arc>, - ) -> Result<(), Report> { + async fn vm_actor_finder(parent_actor_ref: ActorRef) -> Result<(), Report> { trace!("running vm_actor_finder"); let mut vm_actor_stream = RemoteActorRef::::lookup_all(VM); while let Some(vm_actor) = vm_actor_stream.try_next().await? { - let vm_actor_id = vm_actor.id(); - let updater_is_running = keepalive_tasks - .get(&vm_actor_id) - .is_some_and(|task| !task.is_finished()); - if !updater_is_running { - keepalive_tasks.remove(&vm_actor_id); - trace!(?vm_actor, "starting vm_updater_task"); - - parent_actor_ref.link_remote(&vm_actor).await?; - - let vm_actorid_ulid_map_clone = Arc::clone(&vm_actorid_ulid_map); - let data_cache_clone = Arc::clone(&data_cache); - let placements_clone = Arc::clone(&placements); - let updater_task = tokio::spawn(async move { - Self::vm_updater_task( - vm_actor, - vm_actorid_ulid_map_clone, - data_cache_clone, - placements_clone, - ) - .await; - }); - - keepalive_tasks.insert(vm_actor_id, updater_task); - } + parent_actor_ref + .tell(VmActorDiscovered { + actor_ref: vm_actor, + }) + .send() + .await?; } Ok(()) } - async fn vm_updater_task( - actor_ref: RemoteActorRef, - vm_actorid_ulid_map: Arc>, - data_cache: Arc>>, - placements: Arc>, - ) { + async fn vm_updater_task(scheduler: ActorRef, actor_ref: RemoteActorRef) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails: u8 = 0; loop { if let Ok(data) = actor_ref.ask(&GetVMInfo { vmid: None }).await { - let vmid = data.vmid; - - vm_actorid_ulid_map.insert(actor_ref.id(), vmid); // should we be doing this on every loop? idk. but we at least need to do it on the first iteration given we don't know the mapping before that - - if let Some(mut placement) = placements.get_mut(&vmid) { - placement.lifecycle = VmLifecycle::Running; - placement.last_confirmed_at = Some(Instant::now()); - } - let cached_vm = CachedVMActor { - actor_ref: Some(actor_ref.clone()), - data, - cached_at: Instant::now(), - }; - data_cache - .entry(vmid) - .and_modify(|entries| { - Self::update_cached_vm_entry(entries, actor_ref.id(), cached_vm.clone()); + let send_result = scheduler + .tell(VmUpdated { + actor_id: actor_ref.id(), + data, }) - .or_insert_with(|| vec![cached_vm]); + .send() + .await + .map_err(|error| eyre!("failed to send VM update: {error}")); + if let Err(error) = send_result { + warn!(?error, "VM updater could not notify scheduler"); + return; + } fails = 0; } else { @@ -252,32 +248,16 @@ impl SchedulerActor { "can no longer reach vm actor, cleaning up cache entries" ); - let vmid = vm_actorid_ulid_map - .remove(&actor_ref.id()) - .map(|(_, vmid)| vmid); - - // if the actor was never reachable, there's no vm_actorid_ulid_map entry. - // try to recover the vmid from the data_cache by scanning for a stale entry - // that matches this actor_ref. - let vmid = vmid.or_else(|| { - data_cache.iter().find_map(|entry| { - entry - .value() - .iter() - .any(|cached| { - cached - .actor_ref - .as_ref() - .is_some_and(|actor| actor.id() == actor_ref.id()) - }) - .then(|| *entry.key()) + let send_result = scheduler + .tell(VmUpdaterStopped { + actor_id: actor_ref.id(), }) - }); - - if let Some(vmid) = vmid { - Self::remove_vm_state(vmid, &placements, &data_cache); + .send() + .await + .map_err(|error| eyre!("failed to send VM stop: {error}")); + if let Err(error) = send_result { + warn!(?error, "VM updater could not notify scheduler"); } - return; } @@ -285,61 +265,41 @@ impl SchedulerActor { } } - async fn agent_actor_finder( - parent_actor_ref: RemoteActorRef, - data_cache: Arc>, - keepalive_tasks: Arc>>, - ) -> Result<(), Report> { + async fn agent_actor_finder(parent_actor_ref: ActorRef) -> Result<(), Report> { trace!("running agent_actor_finder"); let mut agent_actor_stream = RemoteActorRef::::lookup_all(AGENT); while let Some(agent_actor) = agent_actor_stream.try_next().await? { - let agent_actor_id = agent_actor.id(); - let updater_is_running = keepalive_tasks - .get(&agent_actor_id) - .is_some_and(|task| !task.is_finished()); - if !updater_is_running { - keepalive_tasks.remove(&agent_actor_id); - trace!(?agent_actor, "starting agent_updater_task"); - - parent_actor_ref.link_remote(&agent_actor).await?; - - let data_cache_clone = Arc::clone(&data_cache); - let updater_task = tokio::spawn(async move { - Self::agent_updater_task(agent_actor, data_cache_clone).await; - }); - - keepalive_tasks.insert(agent_actor_id, updater_task); - } + parent_actor_ref + .tell(AgentActorDiscovered { + actor_ref: agent_actor, + }) + .send() + .await?; } Ok(()) } - async fn agent_updater_task( - actor_ref: RemoteActorRef, - data_cache: Arc>, - ) { + async fn agent_updater_task(scheduler: ActorRef, actor_ref: RemoteActorRef) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails: u8 = 0; loop { if let Ok(data) = actor_ref.ask(&GetAgentStatus).await { - if data_cache.contains_key(&actor_ref.id()) { - data_cache.alter(&actor_ref.id(), |_, mut v| { - v.data = data; - v - }); - } else { - data_cache.insert( - actor_ref.id(), - CachedAgentActor { - actor_ref: actor_ref.clone(), - data, - }, - ); + let send_result = scheduler + .tell(AgentUpdated { + actor_id: actor_ref.id(), + actor_ref: actor_ref.clone(), + data, + }) + .send() + .await + .map_err(|error| eyre!("failed to send agent update: {error}")); + if let Err(error) = send_result { + warn!(?error, "agent updater could not notify scheduler"); + return; } - fails = 0; } else { fails = fails.saturating_add(1); @@ -350,7 +310,16 @@ impl SchedulerActor { ?actor_ref, "can no longer reach agent actor, stopping updater" ); - data_cache.remove(&actor_ref.id()); + let send_result = scheduler + .tell(AgentUpdaterStopped { + actor_id: actor_ref.id(), + }) + .send() + .await + .map_err(|error| eyre!("failed to send agent stop: {error}")); + if let Err(error) = send_result { + warn!(?error, "agent updater could not notify scheduler"); + } return; } @@ -358,49 +327,19 @@ impl SchedulerActor { } } - fn start_actor_finder(&mut self, actor_ref: RemoteActorRef) { - let agent_data_cache_arc_clone = Arc::clone(&self.agent_data_cache); - let agent_keepalive_tasks_arc_clone = Arc::clone(&self.agent_keepalive_tasks); - - let vm_actorid_ulid_map_arc_clone = Arc::clone(&self.vm_actorid_ulid_map); - let vm_data_cache_arc_clone = Arc::clone(&self.vm_data_cache); - let vm_placements_arc_clone = Arc::clone(&self.vm_placements); - let vm_keepalive_tasks_arc_clone = Arc::clone(&self.vm_keepalive_tasks); - + fn start_actor_finder(&mut self, actor_ref: ActorRef) { self.cache_actor_finder = Some(tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_secs(1)); loop { - let vm_join_handle = Self::vm_actor_finder( - actor_ref.clone(), - Arc::clone(&vm_actorid_ulid_map_arc_clone), - Arc::clone(&vm_data_cache_arc_clone), - Arc::clone(&vm_keepalive_tasks_arc_clone), - Arc::clone(&vm_placements_arc_clone), - ); - - let agent_join_handle = Self::agent_actor_finder( - actor_ref.clone(), - Arc::clone(&agent_data_cache_arc_clone), - Arc::clone(&agent_keepalive_tasks_arc_clone), - ); - - // intentionally ignoring results because we want to keep finding actors even if an attempt fails - let (vm_result, agent_result) = tokio::join!(vm_join_handle, agent_join_handle); + let vm_result = Self::vm_actor_finder(actor_ref.clone()).await; + let agent_result = Self::agent_actor_finder(actor_ref.clone()).await; if let Err(error) = vm_result { warn!(?error, "VM actor discovery failed"); } if let Err(error) = agent_result { warn!(?error, "agent actor discovery failed"); } - - Self::cleanup_unresolved_vm_cache( - &vm_placements_arc_clone, - &vm_data_cache_arc_clone, - ); - - //info!(?vm_data_cache_arc_clone); - //info!(?agent_data_cache_arc_clone); - + actor_ref.tell(ReconcileVmPlacements).send().await.ok(); interval.tick().await; } })); @@ -428,8 +367,8 @@ impl SchedulerActor { let mut best_agent = None; let mut best_score = AgentScore::REJECTED; - for agent in self.agent_data_cache.iter() { - let score = self.score_agent(msg, &agent); + for agent in self.agent_data_cache.values() { + let score = self.score_agent(msg, agent); if score > best_score { best_agent = Some(agent.actor_ref.clone()); @@ -445,11 +384,7 @@ impl SchedulerActor { // this function intentionally only checks against the cache. this has some positives and negatives: // positive: it will never trigger any network requests so its very fast, and having to do network requests for scoring whenever we want to schedule a vm is likely a bad idea // negative: it technically has a delayed view of the cluster, meaning that some things that happened in the future, may not exist yet. so we need to be careful about how this is done so affinity rules are not accidentally broken. mostly this means, if we do anything that could affect the outcome of an affinity rule (ex: network request to an agent), we need to update the cache, before we do the action. - fn score_agent( - &self, - msg: &CreateVM, - agent: &RefMulti<'_, ActorId, CachedAgentActor>, - ) -> AgentScore { + fn score_agent(&self, msg: &CreateVM, agent: &CachedAgentActor) -> AgentScore { let mut score = AgentScore::default(); let agent_max_vcpus = agent @@ -565,13 +500,13 @@ const fn has_capacity( } fn pending_resources_for_agent( - placements: &DashMap, + placements: &AHashMap, agent_id: ActorId, ) -> (u32, u64) { placements .iter() - .filter(|entry| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) - .map(|entry| (entry.config.data.vcpus, entry.config.data.memory.as_u64())) + .filter(|(_, entry)| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) + .map(|(_, entry)| (entry.config.data.vcpus, entry.config.data.memory.as_u64())) .fold((0u32, 0u64), |(vcpus, ram), (vm_vcpus, vm_ram)| { (vcpus.saturating_add(vm_vcpus), ram.saturating_add(vm_ram)) }) @@ -658,8 +593,8 @@ mod tests { AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, MetadataTable, ObjectMetadata, Operator, VirtualMachine, }; + use ahash::AHashMap; use bytesize::ByteSize; - use dashmap::DashMap; use std::collections::BTreeMap; use std::time::{Duration, Instant}; use ulid::Ulid; @@ -677,7 +612,7 @@ mod tests { fn removes_expired_unresolved_vm_placeholders() { let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); let agent_id = super::ActorId::new(1); - let placements: DashMap = DashMap::new(); + let mut placements: AHashMap = AHashMap::new(); placements.insert( vmid, VmPlacement { @@ -690,9 +625,9 @@ mod tests { last_confirmed_at: None, }, ); - let data_cache: DashMap> = DashMap::new(); + let mut data_cache: AHashMap> = AHashMap::new(); - SchedulerActor::cleanup_unresolved_vm_cache(&placements, &data_cache); + SchedulerActor::cleanup_unresolved_vm_cache(&mut placements, &mut data_cache); assert!(!placements.contains_key(&vmid)); assert!(!data_cache.contains_key(&vmid)); @@ -705,7 +640,7 @@ mod tests { let mut config = VirtualMachine::default(); config.data.vcpus = 4; config.data.memory = ByteSize::gib(8); - let placements = DashMap::new(); + let mut placements: AHashMap = AHashMap::new(); placements.insert( vmid, VmPlacement { @@ -856,16 +791,16 @@ impl Actor for SchedulerActor { info!(?peer_id, "Scheduler Actor started!"); let mut scheduler_actor = Self { - agent_data_cache: Arc::new(DashMap::new()), - agent_keepalive_tasks: Arc::new(DashMap::new()), - vm_actorid_ulid_map: Arc::new(DashMap::new()), - vm_placements: Arc::new(DashMap::new()), - vm_data_cache: Arc::new(DashMap::new()), - vm_keepalive_tasks: Arc::new(DashMap::new()), + agent_data_cache: AHashMap::new(), + agent_keepalive_tasks: AHashMap::new(), + vm_actorid_ulid_map: AHashMap::new(), + vm_placements: AHashMap::new(), + vm_data_cache: AHashMap::new(), + vm_keepalive_tasks: AHashMap::new(), cache_actor_finder: None, }; - scheduler_actor.start_actor_finder(actor_ref.into_remote_ref().await); + scheduler_actor.start_actor_finder(actor_ref.clone()); Ok(scheduler_actor) } @@ -883,20 +818,20 @@ impl Actor for SchedulerActor { return Ok(ControlFlow::Break(ActorStopReason::Killed)); }; - if let Some((_, keepalive_task)) = self.agent_keepalive_tasks.remove(&id) { + if let Some(keepalive_task) = self.agent_keepalive_tasks.remove(&id) { trace!(?id, "Aborting agent keepalive task"); keepalive_task.abort(); } self.agent_data_cache.remove(&id); - if let Some((_, keepalive_task)) = self.vm_keepalive_tasks.remove(&id) { + if let Some(keepalive_task) = self.vm_keepalive_tasks.remove(&id) { trace!(?id, "Aborting vm keepalive task"); keepalive_task.abort(); } - if let Some((_, vmid)) = self.vm_actorid_ulid_map.remove(&id) { - Self::remove_vm_state(vmid, &self.vm_placements, &self.vm_data_cache); + if let Some(vmid) = self.vm_actorid_ulid_map.remove(&id) { + Self::remove_vm_state(vmid, &mut self.vm_placements, &mut self.vm_data_cache); } // todo: attempt vm restarts if necessary. @@ -905,6 +840,142 @@ impl Actor for SchedulerActor { } } +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, msg: VmActorDiscovered, ctx: &mut Context) { + let actor_id = msg.actor_ref.id(); + let updater_is_running = self + .vm_keepalive_tasks + .get(&actor_id) + .is_some_and(|task| !task.is_finished()); + if updater_is_running { + return; + } + self.vm_keepalive_tasks.remove(&actor_id); + if let Err(error) = ctx.actor_ref().link_remote(&msg.actor_ref).await { + warn!(?error, ?actor_id, "failed to link VM actor"); + return; + } + let scheduler = ctx.actor_ref().clone(); + let actor_ref = msg.actor_ref; + let task = tokio::spawn(async move { + Self::vm_updater_task(scheduler, actor_ref).await; + }); + self.vm_keepalive_tasks.insert(actor_id, task); + } +} + +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, msg: AgentActorDiscovered, ctx: &mut Context) { + let actor_id = msg.actor_ref.id(); + let updater_is_running = self + .agent_keepalive_tasks + .get(&actor_id) + .is_some_and(|task| !task.is_finished()); + if updater_is_running { + return; + } + self.agent_keepalive_tasks.remove(&actor_id); + if let Err(error) = ctx.actor_ref().link_remote(&msg.actor_ref).await { + warn!(?error, ?actor_id, "failed to link agent actor"); + return; + } + let scheduler = ctx.actor_ref().clone(); + let actor_ref = msg.actor_ref; + let task = tokio::spawn(async move { + Self::agent_updater_task(scheduler, actor_ref).await; + }); + self.agent_keepalive_tasks.insert(actor_id, task); + } +} + +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, msg: VmUpdated, _ctx: &mut Context) { + let vmid = msg.data.vmid; + self.vm_actorid_ulid_map.insert(msg.actor_id, vmid); + if let Some(placement) = self.vm_placements.get_mut(&vmid) { + placement.lifecycle = VmLifecycle::Running; + placement.last_confirmed_at = Some(Instant::now()); + } + let cached_vm = CachedVMActor { + actor_ref: RemoteActorRef::lookup(vm_actor_id(vmid)) + .await + .ok() + .flatten(), + data: msg.data, + cached_at: Instant::now(), + }; + let entries = self.vm_data_cache.entry(vmid).or_default(); + Self::update_cached_vm_entry(entries, msg.actor_id, cached_vm); + } +} + +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, msg: VmUpdaterStopped, _ctx: &mut Context) { + self.vm_keepalive_tasks.remove(&msg.actor_id); + if let Some(vmid) = self.vm_actorid_ulid_map.remove(&msg.actor_id) { + Self::remove_vm_state(vmid, &mut self.vm_placements, &mut self.vm_data_cache); + } else { + let vmids: Vec<_> = self + .vm_data_cache + .iter() + .filter_map(|(vmid, entries)| { + entries + .iter() + .any(|entry| { + entry + .actor_ref + .as_ref() + .is_some_and(|actor| actor.id() == msg.actor_id) + }) + .then_some(*vmid) + }) + .collect(); + for vmid in vmids { + Self::remove_vm_state(vmid, &mut self.vm_placements, &mut self.vm_data_cache); + } + } + } +} + +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, msg: AgentUpdated, _ctx: &mut Context) { + self.agent_data_cache.insert( + msg.actor_id, + CachedAgentActor { + actor_ref: msg.actor_ref, + data: msg.data, + }, + ); + } +} + +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, msg: AgentUpdaterStopped, _ctx: &mut Context) { + self.agent_keepalive_tasks.remove(&msg.actor_id); + self.agent_data_cache.remove(&msg.actor_id); + } +} + +impl Message for SchedulerActor { + type Reply = (); + + async fn handle(&mut self, _msg: ReconcileVmPlacements, _ctx: &mut Context) { + Self::cleanup_unresolved_vm_cache(&mut self.vm_placements, &mut self.vm_data_cache); + } +} + impl Message for SchedulerActor { type Reply = Result; @@ -953,7 +1024,7 @@ impl Message for SchedulerActor { .flatten() .is_none() { - Self::remove_vm_state(msg.vmid, &self.vm_placements, &self.vm_data_cache); + Self::remove_vm_state(msg.vmid, &mut self.vm_placements, &mut self.vm_data_cache); } Ok(reply?) @@ -1013,7 +1084,7 @@ impl Message for SchedulerActor { ) -> Self::Reply { let mut vms = Vec::new(); - for agent in self.agent_data_cache.iter() { + for agent in self.agent_data_cache.values() { vms.extend_from_slice(agent.data.vms.as_slice()); } From bdf19d1e27435d3f2fc1b4747365a5ed02f13b91 Mon Sep 17 00:00:00 2001 From: Caleb Jones Date: Mon, 17 Aug 2026 21:18:10 -0500 Subject: [PATCH 28/35] clippy --- odorobo/src/actors/scheduler_actor.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 333d537..2f27411 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -454,7 +454,7 @@ impl SchedulerActor { continue; }; - for vm_data_cache_ref in vm_data_cache_refs.iter() { + for vm_data_cache_ref in vm_data_cache_refs { let Some(vm_manifest) = &vm_data_cache_ref.data.config else { continue; }; @@ -800,7 +800,7 @@ impl Actor for SchedulerActor { cache_actor_finder: None, }; - scheduler_actor.start_actor_finder(actor_ref.clone()); + scheduler_actor.start_actor_finder(actor_ref); Ok(scheduler_actor) } From 8d85c9f2bfc249faa36c63cba91c9b94fb5041bb Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 20:14:18 -0700 Subject: [PATCH 29/35] apply mado's second suggestion --- odorobo/src/actors/scheduler_actor.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 2f27411..7c8e71e 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -561,15 +561,15 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq Operator::In => requirement.values.contains(value), Operator::NotIn => !requirement.values.contains(value), Operator::Lt | Operator::Gt => { - if requirement.values.len() != 1 { + let [requirement_value] = &requirement.values[..] else { return false; - } + }; let Ok(value_number): Result = value.parse() else { return false; }; - let Ok(requirement_value_number): Result = requirement.values[0].parse() else { + let Ok(requirement_value_number): Result = requirement_value.parse() else { return false; }; From 4b42ca34c76600ba6c84cc430f0a80c5e50e6da4 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 20:26:55 -0700 Subject: [PATCH 30/35] make vm_placements safer for migration, remove zed launch config --- .gitignore | 5 +- .zed/debug.json | 14 -- odorobo/src/actors/scheduler_actor.rs | 234 ++++++++++++++++++-------- odorobo/src/types.rs | 1 + 4 files changed, 171 insertions(+), 83 deletions(-) delete mode 100644 .zed/debug.json diff --git a/.gitignore b/.gitignore index 69ed440..0e257a6 100644 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,9 @@ dev # option (not recommended) you can uncomment the following to ignore the entire idea folder. .idea/ +# Zed launch configurations are machine-specific. +.zed/debug.json + # AI Slop # Anyone who uses LLMs for slop gets to keep their own slop context. # sharing context is technically good, but bad for optics for the anti-AI crowd, so @@ -34,4 +37,4 @@ CLAUDE.md # mac bullshit -.DS_Store \ No newline at end of file +.DS_Store diff --git a/.zed/debug.json b/.zed/debug.json deleted file mode 100644 index 00271f1..0000000 --- a/.zed/debug.json +++ /dev/null @@ -1,14 +0,0 @@ -[ - { - "label": "Run manager agent", - "build": { - "command": "cargo", - "args": ["build"] - }, - "sourceLanguages": ["rust"], - "program": "$ZED_WORKTREE_ROOT/target/debug/odorobo", - "args": ["--manager-enabled"], - "request": "launch", - "adapter": "CodeLLDB" - } -] diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 7c8e71e..04976fe 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -44,7 +44,7 @@ struct AgentActorDiscovered { #[derive(Debug)] struct VmUpdated { - actor_id: ActorId, + actor_ref: RemoteActorRef, data: GetVMInfoReply, } @@ -109,7 +109,8 @@ pub struct SchedulerActor { pub agent_keepalive_tasks: AHashMap>, pub vm_actorid_ulid_map: AHashMap, - pub vm_placements: AHashMap, + /// A VM may be placed on multiple agents while it is migrating. + pub vm_placements: AHashMap>, /// this is a vec because a vmid/ulid can be scheduled on multiple boxes simultaneously during migration pub vm_data_cache: AHashMap>, pub vm_keepalive_tasks: AHashMap>, @@ -143,50 +144,121 @@ impl SchedulerActor { } fn cleanup_unresolved_vm_cache( - placements: &mut AHashMap, + placements: &mut AHashMap>, data_cache: &mut AHashMap>, ) { let now = Instant::now(); - let expired: Vec<_> = placements - .iter() - .filter(|(_, entry)| { - entry.lifecycle == VmLifecycle::Pending - && now.duration_since(entry.created_at) >= UNRESOLVED_VM_CACHE_TIMEOUT + let empty_vmids: Vec<_> = placements + .iter_mut() + .filter_map(|(vmid, entries)| { + entries.retain(|entry| { + entry.lifecycle != VmLifecycle::Pending + || now.duration_since(entry.created_at) < UNRESOLVED_VM_CACHE_TIMEOUT + }); + entries.is_empty().then_some(*vmid) }) - .map(|(vmid, _)| *vmid) .collect(); - for vmid in expired { + for vmid in empty_vmids { Self::remove_vm_state(vmid, placements, data_cache); } } /// takes a list of observed vmids and gives you a set of every vmid that could be on an agent. fn placement_vm_ids( - placements: &AHashMap, + placements: &AHashMap>, agent_id: ActorId, observed: &[Ulid], ) -> AHashSet { observed .iter() .copied() - .chain( - placements + .chain(placements.iter().filter_map(|(vmid, entries)| { + entries .iter() - .filter_map(|(vmid, entry)| (entry.agent_id == agent_id).then_some(*vmid)), - ) + .any(|entry| entry.agent_id == agent_id) + .then_some(*vmid) + })) .collect() } fn remove_vm_state( vmid: Ulid, - placements: &mut AHashMap, + placements: &mut AHashMap>, data_cache: &mut AHashMap>, ) { placements.remove(&vmid); data_cache.remove(&vmid); } + fn remove_vm_actor(actor_id: ActorId, data_cache: &mut AHashMap>) { + let empty_vmids: Vec<_> = data_cache + .iter_mut() + .filter_map(|(vmid, entries)| { + entries.retain(|entry| { + entry + .actor_ref + .as_ref() + .is_none_or(|actor| actor.id() != actor_id) + }); + entries.is_empty().then_some(*vmid) + }) + .collect(); + for vmid in empty_vmids { + data_cache.remove(&vmid); + } + } + + fn reconcile_agent_placements( + agent_id: ActorId, + status: &AgentStatus, + placements: &mut AHashMap>, + ) { + let now = Instant::now(); + let observed: AHashSet<_> = status.vms.iter().copied().collect(); + let missing_agent_placements: Vec<_> = observed + .iter() + .filter_map(|vmid| { + let entries = placements.get(vmid)?; + (!entries.iter().any(|entry| entry.agent_id == agent_id)) + .then(|| (*vmid, entries[0].config.clone())) + }) + .collect(); + for (vmid, config) in missing_agent_placements { + placements.entry(vmid).or_default().push(VmPlacement { + agent_id, + config, + lifecycle: VmLifecycle::Running, + created_at: now, + last_confirmed_at: Some(now), + }); + } + + let empty_vmids: Vec<_> = placements + .iter_mut() + .filter_map(|(vmid, entries)| { + entries.retain(|entry| { + entry.agent_id != agent_id + || entry.lifecycle == VmLifecycle::Pending + || observed.contains(vmid) + }); + for entry in entries + .iter_mut() + .filter(|entry| entry.agent_id == agent_id) + { + if observed.contains(vmid) { + entry.lifecycle = VmLifecycle::Running; + entry.last_confirmed_at = Some(now); + } + } + entries.is_empty().then_some(*vmid) + }) + .collect(); + for vmid in empty_vmids { + placements.remove(&vmid); + } + } + #[expect(dead_code, reason = "reserved for explicit placement by actor id")] fn lookup_agent_by_actor_id(&self, actor_id: &ActorId) -> Option> { self.agent_data_cache @@ -226,7 +298,7 @@ impl SchedulerActor { if let Ok(data) = actor_ref.ask(&GetVMInfo { vmid: None }).await { let send_result = scheduler .tell(VmUpdated { - actor_id: actor_ref.id(), + actor_ref: actor_ref.clone(), data, }) .send() @@ -500,13 +572,14 @@ const fn has_capacity( } fn pending_resources_for_agent( - placements: &AHashMap, + placements: &AHashMap>, agent_id: ActorId, ) -> (u32, u64) { placements - .iter() - .filter(|(_, entry)| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) - .map(|(_, entry)| (entry.config.data.vcpus, entry.config.data.memory.as_u64())) + .values() + .flat_map(|entries| entries.iter()) + .filter(|entry| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) + .map(|entry| (entry.config.data.vcpus, entry.config.data.memory.as_u64())) .fold((0u32, 0u64), |(vcpus, ram), (vm_vcpus, vm_ram)| { (vcpus.saturating_add(vm_vcpus), ram.saturating_add(vm_ram)) }) @@ -589,6 +662,7 @@ mod tests { evaluate_affinity_rule, evaluate_table_value, has_capacity, pending_resources_for_agent, }; + use crate::messages::agent::AgentStatus; use crate::types::{ AffinityRequirement, AffinityRule, AffinityStrictness, AffinityType, MetadataTable, ObjectMetadata, Operator, VirtualMachine, @@ -612,10 +686,10 @@ mod tests { fn removes_expired_unresolved_vm_placeholders() { let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); let agent_id = super::ActorId::new(1); - let mut placements: AHashMap = AHashMap::new(); + let mut placements: AHashMap> = AHashMap::new(); placements.insert( vmid, - VmPlacement { + vec![VmPlacement { agent_id, config: VirtualMachine::default(), lifecycle: VmLifecycle::Pending, @@ -623,7 +697,7 @@ mod tests { .checked_sub(Duration::from_secs(31)) .expect("test timestamp should be representable"), last_confirmed_at: None, - }, + }], ); let mut data_cache: AHashMap> = AHashMap::new(); @@ -640,16 +714,16 @@ mod tests { let mut config = VirtualMachine::default(); config.data.vcpus = 4; config.data.memory = ByteSize::gib(8); - let mut placements: AHashMap = AHashMap::new(); + let mut placements: AHashMap> = AHashMap::new(); placements.insert( vmid, - VmPlacement { + vec![VmPlacement { agent_id, config, lifecycle: VmLifecycle::Pending, created_at: Instant::now(), last_confirmed_at: None, - }, + }], ); assert_eq!( @@ -658,6 +732,53 @@ mod tests { ); } + #[test] + fn reconciling_source_agent_preserves_destination_migration_placement() { + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let source_agent = super::ActorId::new(1); + let destination_agent = super::ActorId::new(2); + let running_placement = |agent_id| VmPlacement { + agent_id, + config: VirtualMachine::default(), + lifecycle: VmLifecycle::Running, + created_at: Instant::now(), + last_confirmed_at: Some(Instant::now()), + }; + let mut placements = AHashMap::from([(vmid, vec![running_placement(source_agent)])]); + let destination_status = AgentStatus { + hostname: "destination".to_owned(), + vcpus: 1, + ram: ByteSize::b(1), + used_vcpus: 0, + used_ram: ByteSize::b(0), + vms: vec![vmid], + metadata: ObjectMetadata::default(), + }; + SchedulerActor::reconcile_agent_placements( + destination_agent, + &destination_status, + &mut placements, + ); + + let source_status = AgentStatus { + hostname: "source".to_owned(), + vcpus: 1, + ram: ByteSize::b(1), + used_vcpus: 0, + used_ram: ByteSize::b(0), + vms: Vec::new(), + metadata: ObjectMetadata::default(), + }; + + SchedulerActor::reconcile_agent_placements(source_agent, &source_status, &mut placements); + + let remaining = placements + .get(&vmid) + .expect("destination placement remains"); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].agent_id, destination_agent); + } + #[test] fn evaluates_membership_and_missing_keys() { let metadata = BTreeMap::from([("tier".to_owned(), "frontend".to_owned())]); @@ -716,7 +837,10 @@ mod tests { inverse: true, requirements: vec![requirement(Operator::In, &["frontend"])], }; - assert!(!evaluate_affinity_rule(&[metadata], &rule)); + assert!(!evaluate_affinity_rule( + std::slice::from_ref(&metadata), + &rule + )); let empty_rule = AffinityRule { strictness: AffinityStrictness::Required, @@ -830,9 +954,8 @@ impl Actor for SchedulerActor { keepalive_task.abort(); } - if let Some(vmid) = self.vm_actorid_ulid_map.remove(&id) { - Self::remove_vm_state(vmid, &mut self.vm_placements, &mut self.vm_data_cache); - } + self.vm_actorid_ulid_map.remove(&id); + Self::remove_vm_actor(id, &mut self.vm_data_cache); // todo: attempt vm restarts if necessary. @@ -897,21 +1020,15 @@ impl Message for SchedulerActor { async fn handle(&mut self, msg: VmUpdated, _ctx: &mut Context) { let vmid = msg.data.vmid; - self.vm_actorid_ulid_map.insert(msg.actor_id, vmid); - if let Some(placement) = self.vm_placements.get_mut(&vmid) { - placement.lifecycle = VmLifecycle::Running; - placement.last_confirmed_at = Some(Instant::now()); - } + let actor_id = msg.actor_ref.id(); + self.vm_actorid_ulid_map.insert(actor_id, vmid); let cached_vm = CachedVMActor { - actor_ref: RemoteActorRef::lookup(vm_actor_id(vmid)) - .await - .ok() - .flatten(), + actor_ref: Some(msg.actor_ref), data: msg.data, cached_at: Instant::now(), }; let entries = self.vm_data_cache.entry(vmid).or_default(); - Self::update_cached_vm_entry(entries, msg.actor_id, cached_vm); + Self::update_cached_vm_entry(entries, actor_id, cached_vm); } } @@ -920,28 +1037,8 @@ impl Message for SchedulerActor { async fn handle(&mut self, msg: VmUpdaterStopped, _ctx: &mut Context) { self.vm_keepalive_tasks.remove(&msg.actor_id); - if let Some(vmid) = self.vm_actorid_ulid_map.remove(&msg.actor_id) { - Self::remove_vm_state(vmid, &mut self.vm_placements, &mut self.vm_data_cache); - } else { - let vmids: Vec<_> = self - .vm_data_cache - .iter() - .filter_map(|(vmid, entries)| { - entries - .iter() - .any(|entry| { - entry - .actor_ref - .as_ref() - .is_some_and(|actor| actor.id() == msg.actor_id) - }) - .then_some(*vmid) - }) - .collect(); - for vmid in vmids { - Self::remove_vm_state(vmid, &mut self.vm_placements, &mut self.vm_data_cache); - } - } + self.vm_actorid_ulid_map.remove(&msg.actor_id); + Self::remove_vm_actor(msg.actor_id, &mut self.vm_data_cache); } } @@ -949,6 +1046,7 @@ impl Message for SchedulerActor { type Reply = (); async fn handle(&mut self, msg: AgentUpdated, _ctx: &mut Context) { + Self::reconcile_agent_placements(msg.actor_id, &msg.data, &mut self.vm_placements); self.agent_data_cache.insert( msg.actor_id, CachedAgentActor { @@ -986,16 +1084,16 @@ impl Message for SchedulerActor { ) -> Self::Reply { let target_agent = self.schedule_agent(&msg)?; - self.vm_placements.insert( - msg.vmid, - VmPlacement { + self.vm_placements + .entry(msg.vmid) + .or_default() + .push(VmPlacement { agent_id: target_agent.id(), config: msg.config.clone(), lifecycle: VmLifecycle::Pending, created_at: Instant::now(), last_confirmed_at: None, - }, - ); + }); self.vm_data_cache .entry(msg.vmid) .or_default() diff --git a/odorobo/src/types.rs b/odorobo/src/types.rs index be0cf26..8a7ede9 100644 --- a/odorobo/src/types.rs +++ b/odorobo/src/types.rs @@ -178,6 +178,7 @@ pub struct AffinityRule { /// If true, the outcome of the requirements is inverted. #[serde(default)] pub inverse: bool, + /// `ORed` together pub requirements: Vec, } From eab2da4dab603838fb012a72b7c25c76919dcde3 Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 20:36:38 -0700 Subject: [PATCH 31/35] add shrink_to_fit for the vectors to reduce memory usage post-migration --- odorobo/src/actors/scheduler_actor.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 04976fe..b3fba6d 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -124,6 +124,12 @@ static VCPU_OVERPROVISIONMENT_DENOMINATOR: u32 = 1; const UNRESOLVED_VM_CACHE_TIMEOUT: Duration = Duration::from_secs(30); impl SchedulerActor { + fn shrink_non_migrating_entries(entries: &mut Vec) { + if entries.len() <= 1 && entries.capacity() > entries.len() { + entries.shrink_to_fit(); + } + } + fn update_cached_vm_entry( entries: &mut Vec, actor_id: ActorId, @@ -141,6 +147,7 @@ impl SchedulerActor { } else { entries.push(cached_vm); } + Self::shrink_non_migrating_entries(entries); } fn cleanup_unresolved_vm_cache( @@ -155,6 +162,7 @@ impl SchedulerActor { entry.lifecycle != VmLifecycle::Pending || now.duration_since(entry.created_at) < UNRESOLVED_VM_CACHE_TIMEOUT }); + Self::shrink_non_migrating_entries(entries); entries.is_empty().then_some(*vmid) }) .collect(); @@ -201,6 +209,7 @@ impl SchedulerActor { .as_ref() .is_none_or(|actor| actor.id() != actor_id) }); + Self::shrink_non_migrating_entries(entries); entries.is_empty().then_some(*vmid) }) .collect(); @@ -251,6 +260,7 @@ impl SchedulerActor { entry.last_confirmed_at = Some(now); } } + Self::shrink_non_migrating_entries(entries); entries.is_empty().then_some(*vmid) }) .collect(); From 21cda5d2f5be601ced5dd28ed8bbefa21c48bc9d Mon Sep 17 00:00:00 2001 From: Cypress Reed Date: Mon, 17 Aug 2026 20:46:35 -0700 Subject: [PATCH 32/35] memory optimization --- odorobo/src/actors/scheduler_actor.rs | 126 ++++++++++++++++---------- 1 file changed, 78 insertions(+), 48 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index b3fba6d..51adda9 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -83,7 +83,6 @@ pub enum VmLifecycle { #[derive(Debug, Clone)] pub struct VmPlacement { pub agent_id: ActorId, - pub config: VirtualMachine, pub lifecycle: VmLifecycle, pub created_at: Instant, pub last_confirmed_at: Option, @@ -92,7 +91,6 @@ pub struct VmPlacement { #[derive(Debug, Clone)] pub struct CachedVMActor { pub actor_ref: Option>, - pub data: GetVMInfoReply, pub cached_at: Instant, } @@ -109,6 +107,8 @@ pub struct SchedulerActor { pub agent_keepalive_tasks: AHashMap>, pub vm_actorid_ulid_map: AHashMap, + /// Canonical manifest storage; placements and actor entries refer to the VM ID. + pub vm_manifests: AHashMap, /// A VM may be placed on multiple agents while it is migrating. pub vm_placements: AHashMap>, /// this is a vec because a vmid/ulid can be scheduled on multiple boxes simultaneously during migration @@ -151,6 +151,7 @@ impl SchedulerActor { } fn cleanup_unresolved_vm_cache( + manifests: &mut AHashMap, placements: &mut AHashMap>, data_cache: &mut AHashMap>, ) { @@ -168,7 +169,7 @@ impl SchedulerActor { .collect(); for vmid in empty_vmids { - Self::remove_vm_state(vmid, placements, data_cache); + Self::remove_vm_state(vmid, manifests, placements, data_cache); } } @@ -192,9 +193,11 @@ impl SchedulerActor { fn remove_vm_state( vmid: Ulid, + manifests: &mut AHashMap, placements: &mut AHashMap>, data_cache: &mut AHashMap>, ) { + manifests.remove(&vmid); placements.remove(&vmid); data_cache.remove(&vmid); } @@ -221,22 +224,24 @@ impl SchedulerActor { fn reconcile_agent_placements( agent_id: ActorId, status: &AgentStatus, + manifests: &AHashMap, placements: &mut AHashMap>, ) { let now = Instant::now(); let observed: AHashSet<_> = status.vms.iter().copied().collect(); let missing_agent_placements: Vec<_> = observed .iter() - .filter_map(|vmid| { - let entries = placements.get(vmid)?; - (!entries.iter().any(|entry| entry.agent_id == agent_id)) - .then(|| (*vmid, entries[0].config.clone())) + .filter(|vmid| manifests.contains_key(vmid)) + .filter(|vmid| { + placements + .get(vmid) + .is_none_or(|entries| !entries.iter().any(|entry| entry.agent_id == agent_id)) }) + .copied() .collect(); - for (vmid, config) in missing_agent_placements { + for vmid in missing_agent_placements { placements.entry(vmid).or_default().push(VmPlacement { agent_id, - config, lifecycle: VmLifecycle::Running, created_at: now, last_confirmed_at: Some(now), @@ -476,8 +481,11 @@ impl SchedulerActor { .checked_div(VCPU_OVERPROVISIONMENT_DENOMINATOR) .unwrap_or(u32::MAX); // todo: do we care about VMData.max_vcpus? - let (pending_vcpus, pending_ram) = - pending_resources_for_agent(&self.vm_placements, agent.actor_ref.id()); + let (pending_vcpus, pending_ram) = pending_resources_for_agent( + &self.vm_manifests, + &self.vm_placements, + agent.actor_ref.id(), + ); let used_vcpus = agent.data.used_vcpus.saturating_add(pending_vcpus); let used_ram = agent.data.used_ram.as_u64().saturating_add(pending_ram); let agent_used_vcpus = used_vcpus.saturating_add(msg.config.data.vcpus); @@ -522,7 +530,7 @@ impl SchedulerActor { // Roughly based on . if let Some(affinity_rules) = &msg.config.affinity { for rule in affinity_rules { - let mut metadata_tables: Vec = Vec::with_capacity(1); + let mut metadata_tables: Vec<&ObjectMetadata> = Vec::with_capacity(1); match rule.affinity_type { AffinityType::VirtualMachine => { @@ -532,21 +540,16 @@ impl SchedulerActor { &agent.data.vms, ); for vmid in vmids { - let Some(vm_data_cache_refs) = self.vm_data_cache.get(&vmid) else { - continue; - }; - - for vm_data_cache_ref in vm_data_cache_refs { - let Some(vm_manifest) = &vm_data_cache_ref.data.config else { - continue; - }; - if let Some(metadata) = &vm_manifest.metadata { - metadata_tables.push(metadata.clone()); - } + if let Some(metadata) = self + .vm_manifests + .get(&vmid) + .and_then(|manifest| manifest.metadata.as_ref()) + { + metadata_tables.push(metadata); } } } - AffinityType::Agent => metadata_tables.push(agent.data.metadata.clone()), + AffinityType::Agent => metadata_tables.push(&agent.data.metadata), } let follows_rule = evaluate_affinity_rule(&metadata_tables, rule); @@ -582,14 +585,19 @@ const fn has_capacity( } fn pending_resources_for_agent( + manifests: &AHashMap, placements: &AHashMap>, agent_id: ActorId, ) -> (u32, u64) { placements - .values() - .flat_map(|entries| entries.iter()) - .filter(|entry| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) - .map(|entry| (entry.config.data.vcpus, entry.config.data.memory.as_u64())) + .iter() + .flat_map(|(vmid, entries)| entries.iter().map(move |entry| (vmid, entry))) + .filter(|(_, entry)| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) + .filter_map(|(vmid, _)| { + manifests + .get(vmid) + .map(|manifest| (manifest.data.vcpus, manifest.data.memory.as_u64())) + }) .fold((0u32, 0u64), |(vcpus, ram), (vm_vcpus, vm_ram)| { (vcpus.saturating_add(vm_vcpus), ram.saturating_add(vm_ram)) }) @@ -606,7 +614,7 @@ fn affinity_delta(strictness: AffinityStrictness, follows_rule: bool) -> Option< } fn evaluate_affinity_rule( - metadata_tables: &[ObjectMetadata], + metadata_tables: &[&ObjectMetadata], rule: &crate::types::AffinityRule, ) -> bool { let mut follows_rule = false; @@ -701,7 +709,6 @@ mod tests { vmid, vec![VmPlacement { agent_id, - config: VirtualMachine::default(), lifecycle: VmLifecycle::Pending, created_at: Instant::now() .checked_sub(Duration::from_secs(31)) @@ -709,10 +716,16 @@ mod tests { last_confirmed_at: None, }], ); + let mut manifests = AHashMap::from([(vmid, VirtualMachine::default())]); let mut data_cache: AHashMap> = AHashMap::new(); - SchedulerActor::cleanup_unresolved_vm_cache(&mut placements, &mut data_cache); + SchedulerActor::cleanup_unresolved_vm_cache( + &mut manifests, + &mut placements, + &mut data_cache, + ); + assert!(!manifests.contains_key(&vmid)); assert!(!placements.contains_key(&vmid)); assert!(!data_cache.contains_key(&vmid)); } @@ -729,15 +742,15 @@ mod tests { vmid, vec![VmPlacement { agent_id, - config, lifecycle: VmLifecycle::Pending, created_at: Instant::now(), last_confirmed_at: None, }], ); + let manifests = AHashMap::from([(vmid, config)]); assert_eq!( - pending_resources_for_agent(&placements, agent_id), + pending_resources_for_agent(&manifests, &placements, agent_id), (4, ByteSize::gib(8).as_u64()) ); } @@ -749,11 +762,11 @@ mod tests { let destination_agent = super::ActorId::new(2); let running_placement = |agent_id| VmPlacement { agent_id, - config: VirtualMachine::default(), lifecycle: VmLifecycle::Running, created_at: Instant::now(), last_confirmed_at: Some(Instant::now()), }; + let manifests = AHashMap::from([(vmid, VirtualMachine::default())]); let mut placements = AHashMap::from([(vmid, vec![running_placement(source_agent)])]); let destination_status = AgentStatus { hostname: "destination".to_owned(), @@ -767,8 +780,10 @@ mod tests { SchedulerActor::reconcile_agent_placements( destination_agent, &destination_status, + &manifests, &mut placements, ); + assert_eq!(manifests.len(), 1); let source_status = AgentStatus { hostname: "source".to_owned(), @@ -780,7 +795,12 @@ mod tests { metadata: ObjectMetadata::default(), }; - SchedulerActor::reconcile_agent_placements(source_agent, &source_status, &mut placements); + SchedulerActor::reconcile_agent_placements( + source_agent, + &source_status, + &manifests, + &mut placements, + ); let remaining = placements .get(&vmid) @@ -847,10 +867,7 @@ mod tests { inverse: true, requirements: vec![requirement(Operator::In, &["frontend"])], }; - assert!(!evaluate_affinity_rule( - std::slice::from_ref(&metadata), - &rule - )); + assert!(!evaluate_affinity_rule(&[&metadata], &rule)); let empty_rule = AffinityRule { strictness: AffinityStrictness::Required, @@ -928,6 +945,7 @@ impl Actor for SchedulerActor { agent_data_cache: AHashMap::new(), agent_keepalive_tasks: AHashMap::new(), vm_actorid_ulid_map: AHashMap::new(), + vm_manifests: AHashMap::new(), vm_placements: AHashMap::new(), vm_data_cache: AHashMap::new(), vm_keepalive_tasks: AHashMap::new(), @@ -1032,9 +1050,11 @@ impl Message for SchedulerActor { let vmid = msg.data.vmid; let actor_id = msg.actor_ref.id(); self.vm_actorid_ulid_map.insert(actor_id, vmid); + if let Some(manifest) = msg.data.config { + self.vm_manifests.insert(vmid, manifest); + } let cached_vm = CachedVMActor { actor_ref: Some(msg.actor_ref), - data: msg.data, cached_at: Instant::now(), }; let entries = self.vm_data_cache.entry(vmid).or_default(); @@ -1056,7 +1076,12 @@ impl Message for SchedulerActor { type Reply = (); async fn handle(&mut self, msg: AgentUpdated, _ctx: &mut Context) { - Self::reconcile_agent_placements(msg.actor_id, &msg.data, &mut self.vm_placements); + Self::reconcile_agent_placements( + msg.actor_id, + &msg.data, + &self.vm_manifests, + &mut self.vm_placements, + ); self.agent_data_cache.insert( msg.actor_id, CachedAgentActor { @@ -1080,7 +1105,11 @@ impl Message for SchedulerActor { type Reply = (); async fn handle(&mut self, _msg: ReconcileVmPlacements, _ctx: &mut Context) { - Self::cleanup_unresolved_vm_cache(&mut self.vm_placements, &mut self.vm_data_cache); + Self::cleanup_unresolved_vm_cache( + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); } } @@ -1094,12 +1123,12 @@ impl Message for SchedulerActor { ) -> Self::Reply { let target_agent = self.schedule_agent(&msg)?; + self.vm_manifests.insert(msg.vmid, msg.config.clone()); self.vm_placements .entry(msg.vmid) .or_default() .push(VmPlacement { agent_id: target_agent.id(), - config: msg.config.clone(), lifecycle: VmLifecycle::Pending, created_at: Instant::now(), last_confirmed_at: None, @@ -1109,10 +1138,6 @@ impl Message for SchedulerActor { .or_default() .push(CachedVMActor { actor_ref: None, - data: GetVMInfoReply { - vmid: msg.vmid, - config: Some(msg.config.clone()), - }, cached_at: Instant::now(), }); @@ -1132,7 +1157,12 @@ impl Message for SchedulerActor { .flatten() .is_none() { - Self::remove_vm_state(msg.vmid, &mut self.vm_placements, &mut self.vm_data_cache); + Self::remove_vm_state( + msg.vmid, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); } Ok(reply?) From f90cce3e2f5a647036624597820d24a5b856abee Mon Sep 17 00:00:00 2001 From: Willow C Reed Date: Tue, 18 Aug 2026 20:29:01 -0600 Subject: [PATCH 33/35] fix VM placeholder stuff, preserve VM state during migration, handle placement removal, fix test that fails erroneously sometimes --- odorobo/src/actors/scheduler_actor.rs | 63 ++++++++++++++++++++++++++- odorobo/src/config.rs | 32 +++++++++----- 2 files changed, 82 insertions(+), 13 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 51adda9..2efff3c 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -221,6 +221,26 @@ impl SchedulerActor { } } + fn remove_agent_placements( + agent_id: ActorId, + manifests: &mut AHashMap, + placements: &mut AHashMap>, + data_cache: &mut AHashMap>, + ) { + let empty_vmids: Vec<_> = placements + .iter_mut() + .filter_map(|(vmid, entries)| { + entries.retain(|entry| entry.agent_id != agent_id); + Self::shrink_non_migrating_entries(entries); + entries.is_empty().then_some(*vmid) + }) + .collect(); + + for vmid in empty_vmids { + Self::remove_vm_state(vmid, manifests, placements, data_cache); + } + } + fn reconcile_agent_placements( agent_id: ActorId, status: &AgentStatus, @@ -976,14 +996,33 @@ impl Actor for SchedulerActor { } self.agent_data_cache.remove(&id); + Self::remove_agent_placements( + id, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); if let Some(keepalive_task) = self.vm_keepalive_tasks.remove(&id) { trace!(?id, "Aborting vm keepalive task"); keepalive_task.abort(); } - self.vm_actorid_ulid_map.remove(&id); + let vmid = self.vm_actorid_ulid_map.remove(&id); Self::remove_vm_actor(id, &mut self.vm_data_cache); + if let Some(vmid) = vmid + && self + .vm_data_cache + .get(&vmid) + .is_none_or(|entries| entries.iter().all(|entry| entry.actor_ref.is_none())) + { + Self::remove_vm_state( + vmid, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); + } // todo: attempt vm restarts if necessary. @@ -1067,8 +1106,22 @@ impl Message for SchedulerActor { async fn handle(&mut self, msg: VmUpdaterStopped, _ctx: &mut Context) { self.vm_keepalive_tasks.remove(&msg.actor_id); - self.vm_actorid_ulid_map.remove(&msg.actor_id); + let vmid = self.vm_actorid_ulid_map.remove(&msg.actor_id); Self::remove_vm_actor(msg.actor_id, &mut self.vm_data_cache); + + if let Some(vmid) = vmid + && self + .vm_data_cache + .get(&vmid) + .is_none_or(|entries| entries.iter().all(|entry| entry.actor_ref.is_none())) + { + Self::remove_vm_state( + vmid, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); + } } } @@ -1098,6 +1151,12 @@ impl Message for SchedulerActor { async fn handle(&mut self, msg: AgentUpdaterStopped, _ctx: &mut Context) { self.agent_keepalive_tasks.remove(&msg.actor_id); self.agent_data_cache.remove(&msg.actor_id); + Self::remove_agent_placements( + msg.actor_id, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); } } diff --git a/odorobo/src/config.rs b/odorobo/src/config.rs index 73126a7..050563d 100644 --- a/odorobo/src/config.rs +++ b/odorobo/src/config.rs @@ -21,17 +21,27 @@ fn default_gateway() -> Ipv4Addr { } /// Infers the default upstream interface from the system's default route fn default_upstream_iface() -> String { - // ip route - let out = std::process::Command::new("ip") - .arg("route") - .output() - .unwrap(); - let output = String::from_utf8(out.stdout).unwrap(); - - let default_route = output.lines().find(|l| l.starts_with("default")).unwrap(); - let iface = default_route.split_whitespace().nth(4).unwrap(); + let Ok(output) = std::process::Command::new("ip").arg("route").output() else { + warn!("cannot inspect routes; defaulting upstream interface to eth0"); + return "eth0".to_owned(); + }; + + let Ok(output) = String::from_utf8(output.stdout) else { + warn!("route output is not valid UTF-8; defaulting upstream interface to eth0"); + return "eth0".to_owned(); + }; + + let Some(iface) = output + .lines() + .find(|line| line.starts_with("default")) + .and_then(|line| line.split_whitespace().nth(4)) + else { + warn!("no default route found; defaulting upstream interface to eth0"); + return "eth0".to_owned(); + }; + info!("inferring default upstream interface: {}", iface); - iface.into() + iface.to_owned() } /// DHCP server config @@ -203,7 +213,7 @@ mod tests { bridge: "vmbr0".to_owned(), gateway: Ipv4Addr::new(10, 10, 100, 1), subnet: Ipv4Net::new(Ipv4Addr::new(10, 10, 100, 0), 24).unwrap(), - upstream_iface: default_upstream_iface(), + upstream_iface: "test0".to_owned(), }, }, ..Default::default() From fd898d400b9972e01c998f0f3163368d1786dab2 Mon Sep 17 00:00:00 2001 From: Willow C Reed Date: Tue, 18 Aug 2026 20:39:36 -0600 Subject: [PATCH 34/35] performance optimizations --- odorobo/src/actors/scheduler_actor.rs | 105 +++++++++++++++----------- odorobo/src/ch_driver/actor.rs | 17 ++++- odorobo/src/messages/vm.rs | 11 ++- 3 files changed, 88 insertions(+), 45 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index 2efff3c..d473641 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -8,8 +8,8 @@ use crate::actors::agent_actor::AgentActor; use crate::ch_driver::actor::VMActor; use crate::messages::agent::{AgentStatus, GetAgentStatus}; use crate::messages::vm::{ - AgentListVMs, AgentListVMsReply, CreateVM, CreateVMReply, DeleteVM, DeleteVMReply, GetVMInfo, - GetVMInfoReply, ShutdownVM, ShutdownVMReply, + AgentListVMs, AgentListVMsReply, CreateVM, CreateVMReply, DeleteVM, DeleteVMReply, + GetVMHeartbeat, GetVMInfo, GetVMInfoReply, ShutdownVM, ShutdownVMReply, }; use crate::messages::{Ping, Pong}; use crate::types::AffinityRequirement; @@ -91,7 +91,6 @@ pub struct VmPlacement { #[derive(Debug, Clone)] pub struct CachedVMActor { pub actor_ref: Option>, - pub cached_at: Instant, } // todo: we should improve the cache to not have agents and vms send the full data on every update. @@ -125,7 +124,7 @@ const UNRESOLVED_VM_CACHE_TIMEOUT: Duration = Duration::from_secs(30); impl SchedulerActor { fn shrink_non_migrating_entries(entries: &mut Vec) { - if entries.len() <= 1 && entries.capacity() > entries.len() { + if entries.len() <= 1 && entries.capacity() >= 4 { entries.shrink_to_fit(); } } @@ -329,21 +328,29 @@ impl SchedulerActor { async fn vm_updater_task(scheduler: ActorRef, actor_ref: RemoteActorRef) { let mut interval = tokio::time::interval(Duration::from_secs(1)); let mut fails: u8 = 0; + let mut initialized = false; + loop { - if let Ok(data) = actor_ref.ask(&GetVMInfo { vmid: None }).await { - let send_result = scheduler - .tell(VmUpdated { - actor_ref: actor_ref.clone(), - data, - }) - .send() - .await - .map_err(|error| eyre!("failed to send VM update: {error}")); - if let Err(error) = send_result { - warn!(?error, "VM updater could not notify scheduler"); - return; + if !initialized { + if let Ok(data) = actor_ref.ask(&GetVMInfo { vmid: None }).await { + let send_result = scheduler + .tell(VmUpdated { + actor_ref: actor_ref.clone(), + data, + }) + .send() + .await + .map_err(|error| eyre!("failed to send VM update: {error}")); + if let Err(error) = send_result { + warn!(?error, "VM updater could not notify scheduler"); + return; + } + initialized = true; + fails = 0; + } else { + fails = fails.saturating_add(1); } - + } else if actor_ref.ask(&GetVMHeartbeat).await.is_ok() { fails = 0; } else { fails = fails.saturating_add(1); @@ -436,7 +443,7 @@ impl SchedulerActor { fn start_actor_finder(&mut self, actor_ref: ActorRef) { self.cache_actor_finder = Some(tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_secs(1)); + let mut interval = tokio::time::interval(Duration::from_secs(5)); loop { let vm_result = Self::vm_actor_finder(actor_ref.clone()).await; let agent_result = Self::agent_actor_finder(actor_ref.clone()).await; @@ -471,11 +478,12 @@ impl SchedulerActor { /// and then if someone tries to schedule lets say 10 VMs in a batch, we could end up scheduling them all to the same agent because the metadata hasn't updated. /// - there are a few solutions for this but they all kinda suck, mostly due to also making sure we deal with latency properly. I am ignoring the issue for now. fn schedule_agent(&self, msg: &CreateVM) -> Result, Report> { + let pending_resources = pending_resources_by_agent(&self.vm_manifests, &self.vm_placements); let mut best_agent = None; let mut best_score = AgentScore::REJECTED; for agent in self.agent_data_cache.values() { - let score = self.score_agent(msg, agent); + let score = self.score_agent(msg, agent, &pending_resources); if score > best_score { best_agent = Some(agent.actor_ref.clone()); @@ -491,7 +499,12 @@ impl SchedulerActor { // this function intentionally only checks against the cache. this has some positives and negatives: // positive: it will never trigger any network requests so its very fast, and having to do network requests for scoring whenever we want to schedule a vm is likely a bad idea // negative: it technically has a delayed view of the cluster, meaning that some things that happened in the future, may not exist yet. so we need to be careful about how this is done so affinity rules are not accidentally broken. mostly this means, if we do anything that could affect the outcome of an affinity rule (ex: network request to an agent), we need to update the cache, before we do the action. - fn score_agent(&self, msg: &CreateVM, agent: &CachedAgentActor) -> AgentScore { + fn score_agent( + &self, + msg: &CreateVM, + agent: &CachedAgentActor, + pending_resources: &AHashMap, + ) -> AgentScore { let mut score = AgentScore::default(); let agent_max_vcpus = agent @@ -501,11 +514,10 @@ impl SchedulerActor { .checked_div(VCPU_OVERPROVISIONMENT_DENOMINATOR) .unwrap_or(u32::MAX); // todo: do we care about VMData.max_vcpus? - let (pending_vcpus, pending_ram) = pending_resources_for_agent( - &self.vm_manifests, - &self.vm_placements, - agent.actor_ref.id(), - ); + let (pending_vcpus, pending_ram) = pending_resources + .get(&agent.actor_ref.id()) + .copied() + .unwrap_or_default(); let used_vcpus = agent.data.used_vcpus.saturating_add(pending_vcpus); let used_ram = agent.data.used_ram.as_u64().saturating_add(pending_ram); let agent_used_vcpus = used_vcpus.saturating_add(msg.config.data.vcpus); @@ -604,23 +616,36 @@ const fn has_capacity( && used_ram.saturating_add(requested_ram) <= max_ram } +fn pending_resources_by_agent( + manifests: &AHashMap, + placements: &AHashMap>, +) -> AHashMap { + let mut resources = AHashMap::new(); + for (vmid, entries) in placements { + let Some(manifest) = manifests.get(vmid) else { + continue; + }; + for entry in entries { + if entry.lifecycle == VmLifecycle::Pending { + let totals = resources.entry(entry.agent_id).or_insert((0u32, 0u64)); + totals.0 = totals.0.saturating_add(manifest.data.vcpus); + totals.1 = totals.1.saturating_add(manifest.data.memory.as_u64()); + } + } + } + resources +} + +#[cfg(test)] fn pending_resources_for_agent( manifests: &AHashMap, placements: &AHashMap>, agent_id: ActorId, ) -> (u32, u64) { - placements - .iter() - .flat_map(|(vmid, entries)| entries.iter().map(move |entry| (vmid, entry))) - .filter(|(_, entry)| entry.agent_id == agent_id && entry.lifecycle == VmLifecycle::Pending) - .filter_map(|(vmid, _)| { - manifests - .get(vmid) - .map(|manifest| (manifest.data.vcpus, manifest.data.memory.as_u64())) - }) - .fold((0u32, 0u64), |(vcpus, ram), (vm_vcpus, vm_ram)| { - (vcpus.saturating_add(vm_vcpus), ram.saturating_add(vm_ram)) - }) + pending_resources_by_agent(manifests, placements) + .get(&agent_id) + .copied() + .unwrap_or_default() } fn affinity_delta(strictness: AffinityStrictness, follows_rule: bool) -> Option { @@ -1094,7 +1119,6 @@ impl Message for SchedulerActor { } let cached_vm = CachedVMActor { actor_ref: Some(msg.actor_ref), - cached_at: Instant::now(), }; let entries = self.vm_data_cache.entry(vmid).or_default(); Self::update_cached_vm_entry(entries, actor_id, cached_vm); @@ -1195,10 +1219,7 @@ impl Message for SchedulerActor { self.vm_data_cache .entry(msg.vmid) .or_default() - .push(CachedVMActor { - actor_ref: None, - cached_at: Instant::now(), - }); + .push(CachedVMActor { actor_ref: None }); let reply = target_agent.ask(&msg).await; diff --git a/odorobo/src/ch_driver/actor.rs b/odorobo/src/ch_driver/actor.rs index 122eed3..4d6251f 100644 --- a/odorobo/src/ch_driver/actor.rs +++ b/odorobo/src/ch_driver/actor.rs @@ -1,6 +1,6 @@ use crate::messages::vm::{ - DeleteVM, GetVMInfo, GetVMInfoReply, MigrateVMReceive, MigrateVMReceiveReply, PrepMigration, - ShutdownVM, + DeleteVM, GetVMHeartbeat, GetVMHeartbeatReply, GetVMInfo, GetVMInfoReply, MigrateVMReceive, + MigrateVMReceiveReply, PrepMigration, ShutdownVM, }; use crate::{ch_driver::VMInstance, types::VirtualMachine}; use cloud_hypervisor_client::models::{ @@ -170,6 +170,19 @@ impl Message for VMActor { } } +#[remote_message] +impl Message for VMActor { + type Reply = GetVMHeartbeatReply; + + async fn handle( + &mut self, + _msg: GetVMHeartbeat, + _ctx: &mut Context, + ) -> Self::Reply { + GetVMHeartbeatReply { vmid: self.vmid } + } +} + #[remote_message] impl Message for VMActor { type Reply = MigrateVMReceiveReply; diff --git a/odorobo/src/messages/vm.rs b/odorobo/src/messages/vm.rs index c5d8dd8..54c275c 100644 --- a/odorobo/src/messages/vm.rs +++ b/odorobo/src/messages/vm.rs @@ -92,7 +92,7 @@ pub struct AgentListVMsReply { pub vms: Vec, } -/// Get VM info +/// Get VM info, including the full manifest. #[derive(Serialize, Deserialize, Debug)] pub struct GetVMInfo { pub vmid: Option, @@ -103,3 +103,12 @@ pub struct GetVMInfoReply { pub vmid: Ulid, pub config: Option, } + +/// Lightweight VM liveness check used by the scheduler heartbeat. +#[derive(Serialize, Deserialize, Debug)] +pub struct GetVMHeartbeat; + +#[derive(Serialize, Deserialize, Reply, Debug, Clone, Copy)] +pub struct GetVMHeartbeatReply { + pub vmid: Ulid, +} From f00e86de3a6a390f7f704b417284bf6606bb959d Mon Sep 17 00:00:00 2001 From: Willow C Reed Date: Tue, 18 Aug 2026 20:57:38 -0600 Subject: [PATCH 35/35] cleanup some stuff, add tests --- odorobo/src/actors/scheduler_actor.rs | 212 +++++++++++++++++++++----- 1 file changed, 174 insertions(+), 38 deletions(-) diff --git a/odorobo/src/actors/scheduler_actor.rs b/odorobo/src/actors/scheduler_actor.rs index d473641..33fcbbc 100644 --- a/odorobo/src/actors/scheduler_actor.rs +++ b/odorobo/src/actors/scheduler_actor.rs @@ -42,6 +42,12 @@ struct AgentActorDiscovered { actor_ref: RemoteActorRef, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CachedActorKind { + Agent, + Vm, +} + #[derive(Debug)] struct VmUpdated { actor_ref: RemoteActorRef, @@ -113,6 +119,7 @@ pub struct SchedulerActor { /// this is a vec because a vmid/ulid can be scheduled on multiple boxes simultaneously during migration pub vm_data_cache: AHashMap>, pub vm_keepalive_tasks: AHashMap>, + actor_kinds: AHashMap, pub cache_actor_finder: Option>, } @@ -240,6 +247,54 @@ impl SchedulerActor { } } + fn rollback_failed_create( + vmid: Ulid, + actor_exists: bool, + manifests: &mut AHashMap, + placements: &mut AHashMap>, + data_cache: &mut AHashMap>, + ) { + if !actor_exists { + Self::remove_vm_state(vmid, manifests, placements, data_cache); + } + } + + fn cleanup_agent_actor(&mut self, actor_id: ActorId) { + if let Some(keepalive_task) = self.agent_keepalive_tasks.remove(&actor_id) { + trace!(?actor_id, "Aborting agent keepalive task"); + keepalive_task.abort(); + } + self.agent_data_cache.remove(&actor_id); + Self::remove_agent_placements( + actor_id, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); + } + + fn cleanup_vm_actor(&mut self, actor_id: ActorId) { + if let Some(keepalive_task) = self.vm_keepalive_tasks.remove(&actor_id) { + trace!(?actor_id, "Aborting VM keepalive task"); + keepalive_task.abort(); + } + let vmid = self.vm_actorid_ulid_map.remove(&actor_id); + Self::remove_vm_actor(actor_id, &mut self.vm_data_cache); + if let Some(vmid) = vmid + && self + .vm_data_cache + .get(&vmid) + .is_none_or(|entries| entries.iter().all(|entry| entry.actor_ref.is_none())) + { + Self::remove_vm_state( + vmid, + &mut self.vm_manifests, + &mut self.vm_placements, + &mut self.vm_data_cache, + ); + } + } + fn reconcile_agent_placements( agent_id: ActorId, status: &AgentStatus, @@ -721,7 +776,7 @@ fn evaluate_table_value(value_option: Option<&String>, requirement: &AffinityReq #[cfg(test)] mod tests { use super::{ - CachedVMActor, SchedulerActor, VmLifecycle, VmPlacement, affinity_delta, + CachedActorKind, CachedVMActor, SchedulerActor, VmLifecycle, VmPlacement, affinity_delta, evaluate_affinity_rule, evaluate_table_value, has_capacity, pending_resources_for_agent, }; @@ -854,6 +909,110 @@ mod tests { assert_eq!(remaining[0].agent_id, destination_agent); } + #[test] + fn failed_create_rolls_back_state_without_an_actor() { + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let mut manifests = AHashMap::from([(vmid, VirtualMachine::default())]); + let mut placements = AHashMap::from([( + vmid, + vec![VmPlacement { + agent_id: super::ActorId::new(1), + lifecycle: VmLifecycle::Pending, + created_at: Instant::now(), + last_confirmed_at: None, + }], + )]); + let mut data_cache = AHashMap::from([(vmid, vec![CachedVMActor { actor_ref: None }])]); + + SchedulerActor::rollback_failed_create( + vmid, + false, + &mut manifests, + &mut placements, + &mut data_cache, + ); + + assert!(!manifests.contains_key(&vmid)); + assert!(!placements.contains_key(&vmid)); + assert!(!data_cache.contains_key(&vmid)); + } + + #[test] + fn failed_create_keeps_state_if_actor_exists() { + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let mut manifests = AHashMap::from([(vmid, VirtualMachine::default())]); + let mut placements = AHashMap::new(); + let mut data_cache = AHashMap::new(); + + SchedulerActor::rollback_failed_create( + vmid, + true, + &mut manifests, + &mut placements, + &mut data_cache, + ); + + assert!(manifests.contains_key(&vmid)); + } + + #[test] + fn agent_cleanup_does_not_remove_unrelated_vm_state() { + let agent_id = super::ActorId::new(1); + let vm_actor_id = super::ActorId::new(2); + let placement_agent_id = super::ActorId::new(3); + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let mut scheduler = SchedulerActor { + agent_data_cache: AHashMap::new(), + agent_keepalive_tasks: AHashMap::new(), + vm_actorid_ulid_map: AHashMap::from([(vm_actor_id, vmid)]), + vm_manifests: AHashMap::from([(vmid, VirtualMachine::default())]), + vm_placements: AHashMap::from([( + vmid, + vec![VmPlacement { + agent_id: placement_agent_id, + lifecycle: VmLifecycle::Running, + created_at: Instant::now(), + last_confirmed_at: Some(Instant::now()), + }], + )]), + vm_data_cache: AHashMap::from([(vmid, vec![CachedVMActor { actor_ref: None }])]), + vm_keepalive_tasks: AHashMap::new(), + actor_kinds: AHashMap::from([(agent_id, CachedActorKind::Agent)]), + cache_actor_finder: None, + }; + + scheduler.cleanup_agent_actor(agent_id); + + assert!(scheduler.vm_placements.contains_key(&vmid)); + assert!(scheduler.vm_manifests.contains_key(&vmid)); + assert!(scheduler.vm_actorid_ulid_map.contains_key(&vm_actor_id)); + assert!(scheduler.vm_data_cache.contains_key(&vmid)); + } + + #[test] + fn vm_cleanup_does_not_remove_unrelated_agent_state() { + let agent_id = super::ActorId::new(1); + let vm_actor_id = super::ActorId::new(2); + let vmid = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ulid"); + let mut scheduler = SchedulerActor { + agent_data_cache: AHashMap::new(), + agent_keepalive_tasks: AHashMap::new(), + vm_actorid_ulid_map: AHashMap::from([(vm_actor_id, vmid)]), + vm_manifests: AHashMap::from([(vmid, VirtualMachine::default())]), + vm_placements: AHashMap::new(), + vm_data_cache: AHashMap::from([(vmid, vec![CachedVMActor { actor_ref: None }])]), + vm_keepalive_tasks: AHashMap::new(), + actor_kinds: AHashMap::from([(agent_id, CachedActorKind::Agent)]), + cache_actor_finder: None, + }; + + scheduler.cleanup_vm_actor(vm_actor_id); + + assert!(scheduler.actor_kinds.contains_key(&agent_id)); + assert!(!scheduler.vm_manifests.contains_key(&vmid)); + assert!(!scheduler.vm_data_cache.contains_key(&vmid)); + } + #[test] fn evaluates_membership_and_missing_keys() { let metadata = BTreeMap::from([("tier".to_owned(), "frontend".to_owned())]); @@ -994,6 +1153,7 @@ impl Actor for SchedulerActor { vm_placements: AHashMap::new(), vm_data_cache: AHashMap::new(), vm_keepalive_tasks: AHashMap::new(), + actor_kinds: AHashMap::new(), cache_actor_finder: None, }; @@ -1015,38 +1175,10 @@ impl Actor for SchedulerActor { return Ok(ControlFlow::Break(ActorStopReason::Killed)); }; - if let Some(keepalive_task) = self.agent_keepalive_tasks.remove(&id) { - trace!(?id, "Aborting agent keepalive task"); - keepalive_task.abort(); - } - - self.agent_data_cache.remove(&id); - Self::remove_agent_placements( - id, - &mut self.vm_manifests, - &mut self.vm_placements, - &mut self.vm_data_cache, - ); - - if let Some(keepalive_task) = self.vm_keepalive_tasks.remove(&id) { - trace!(?id, "Aborting vm keepalive task"); - keepalive_task.abort(); - } - - let vmid = self.vm_actorid_ulid_map.remove(&id); - Self::remove_vm_actor(id, &mut self.vm_data_cache); - if let Some(vmid) = vmid - && self - .vm_data_cache - .get(&vmid) - .is_none_or(|entries| entries.iter().all(|entry| entry.actor_ref.is_none())) - { - Self::remove_vm_state( - vmid, - &mut self.vm_manifests, - &mut self.vm_placements, - &mut self.vm_data_cache, - ); + match self.actor_kinds.remove(&id) { + Some(CachedActorKind::Agent) => self.cleanup_agent_actor(id), + Some(CachedActorKind::Vm) => self.cleanup_vm_actor(id), + None => {} } // todo: attempt vm restarts if necessary. @@ -1072,6 +1204,7 @@ impl Message for SchedulerActor { warn!(?error, ?actor_id, "failed to link VM actor"); return; } + self.actor_kinds.insert(actor_id, CachedActorKind::Vm); let scheduler = ctx.actor_ref().clone(); let actor_ref = msg.actor_ref; let task = tokio::spawn(async move { @@ -1098,6 +1231,7 @@ impl Message for SchedulerActor { warn!(?error, ?actor_id, "failed to link agent actor"); return; } + self.actor_kinds.insert(actor_id, CachedActorKind::Agent); let scheduler = ctx.actor_ref().clone(); let actor_ref = msg.actor_ref; let task = tokio::spawn(async move { @@ -1130,6 +1264,7 @@ impl Message for SchedulerActor { async fn handle(&mut self, msg: VmUpdaterStopped, _ctx: &mut Context) { self.vm_keepalive_tasks.remove(&msg.actor_id); + self.actor_kinds.remove(&msg.actor_id); let vmid = self.vm_actorid_ulid_map.remove(&msg.actor_id); Self::remove_vm_actor(msg.actor_id, &mut self.vm_data_cache); @@ -1174,6 +1309,7 @@ impl Message for SchedulerActor { async fn handle(&mut self, msg: AgentUpdaterStopped, _ctx: &mut Context) { self.agent_keepalive_tasks.remove(&msg.actor_id); + self.actor_kinds.remove(&msg.actor_id); self.agent_data_cache.remove(&msg.actor_id); Self::remove_agent_placements( msg.actor_id, @@ -1230,15 +1366,15 @@ impl Message for SchedulerActor { self.vm_actorid_ulid_map.insert(actor_id, msg.vmid); } - if reply.is_err() - && RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) + if reply.is_err() { + let actor_exists = RemoteActorRef::::lookup(vm_actor_id(msg.vmid)) .await .ok() .flatten() - .is_none() - { - Self::remove_vm_state( + .is_some(); + Self::rollback_failed_create( msg.vmid, + actor_exists, &mut self.vm_manifests, &mut self.vm_placements, &mut self.vm_data_cache,