Skip to main content

khora_control/
budget_channel.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Unidirectional budget channel from the DCC cold thread to the Scheduler hot thread.
16//!
17//! The DCC sends updated budgets at ~20 Hz. The Scheduler consumes them once per frame
18//! at the start of `run_frame()`. If multiple budgets arrive between frames, only the
19//! last one is kept ("last wins" semantics).
20
21use crossbeam_channel::{Receiver, Sender};
22use khora_core::control::gorna::{AgentId, ResourceBudget};
23use std::collections::HashMap;
24use std::sync::{Arc, RwLock};
25
26/// Shared inner state for the budget channel.
27struct BudgetChannelInner {
28    senders: HashMap<AgentId, Sender<ResourceBudget>>,
29    receivers: HashMap<AgentId, Receiver<ResourceBudget>>,
30    current: RwLock<HashMap<AgentId, ResourceBudget>>,
31}
32
33/// Budget channel for cold → hot thread communication.
34///
35/// Cloneable — both the DCC (sender) and Scheduler (receiver) hold clones.
36#[derive(Clone)]
37pub struct BudgetChannel {
38    inner: Arc<BudgetChannelInner>,
39}
40
41impl BudgetChannel {
42    /// Creates a new budget channel with a sender/receiver pair per agent.
43    pub fn new(agent_ids: &[AgentId]) -> Self {
44        let mut senders = HashMap::new();
45        let mut receivers = HashMap::new();
46
47        for &id in agent_ids {
48            let (tx, rx) = crossbeam_channel::unbounded();
49            senders.insert(id, tx);
50            receivers.insert(id, rx);
51        }
52
53        Self {
54            inner: Arc::new(BudgetChannelInner {
55                senders,
56                receivers,
57                current: RwLock::new(HashMap::new()),
58            }),
59        }
60    }
61
62    /// Cold thread: sends a budget for an agent. Non-blocking, last wins.
63    pub fn send(&self, agent_id: AgentId, budget: ResourceBudget) {
64        if let Some(tx) = self.inner.senders.get(&agent_id) {
65            let _ = tx.try_send(budget);
66        }
67    }
68
69    /// Hot thread: syncs all pending budgets at the start of each frame.
70    /// Drains all channels and keeps only the last budget per agent.
71    pub fn sync(&self) {
72        let mut current = self
73            .inner
74            .current
75            .write()
76            .unwrap_or_else(|e| e.into_inner());
77        for (&agent_id, rx) in &self.inner.receivers {
78            while let Ok(budget) = rx.try_recv() {
79                current.insert(agent_id, budget);
80            }
81        }
82    }
83
84    /// Hot thread: reads the current budget for an agent.
85    pub fn get(&self, agent_id: AgentId) -> Option<ResourceBudget> {
86        self.inner
87            .current
88            .read()
89            .unwrap_or_else(|e| e.into_inner())
90            .get(&agent_id)
91            .cloned()
92    }
93}