hataori_runtime/
domain.rs

1use crate::atomic::AtomicTryUpdate;
2use crate::{
3    action::{ActionValue, RegisteredAction},
4    error::{ActionError, ResourceKind, RuntimeError},
5};
6use hataori_runtime_foundation::protocol::{
7    ActionId, DomainId, LocalityId, ObjectId, RequestId, TraceId,
8};
9use rayon::{ThreadPool, ThreadPoolBuilder};
10use std::{
11    collections::BTreeMap,
12    sync::{
13        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
14        mpsc::{self, Receiver, SyncSender},
15        Arc,
16    },
17};
18
19#[derive(Clone, Copy, Debug, Eq, PartialEq)]
20pub struct DomainConfig {
21    pub id: DomainId,
22    pub workers: usize,
23    pub queue_capacity: usize,
24}
25
26impl Default for DomainConfig {
27    fn default() -> Self {
28        Self {
29            id: DomainId::DEFAULT,
30            workers: 1,
31            queue_capacity: 64,
32        }
33    }
34}
35
36#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
37pub struct DomainStats {
38    pub queued: usize,
39    pub running: usize,
40    pub pending_completions: usize,
41    pub completed: u64,
42    pub cancelled_before_start: u64,
43    pub total_queue_latency_ns: u64,
44    pub max_queue_latency_ns: u64,
45    pub total_execution_ns: u64,
46    pub max_execution_ns: u64,
47}
48
49#[derive(Clone, Copy, Debug)]
50pub(crate) struct ObjectJob {
51    pub object: ObjectId,
52    pub read: bool,
53    pub placement: bool,
54}
55
56pub(crate) struct ActionJob {
57    pub requester: LocalityId,
58    pub request: RequestId,
59    pub action_id: ActionId,
60    pub domain: DomainId,
61    pub trace_id: Option<TraceId>,
62    pub input: ActionValue,
63    pub handler: RegisteredAction,
64    pub cancelled: Arc<AtomicBool>,
65    pub local: bool,
66    pub submitted_at: std::time::Instant,
67    pub object: Option<ObjectJob>,
68}
69
70pub(crate) struct ActionCompletion {
71    pub requester: LocalityId,
72    pub request: RequestId,
73    pub action_id: ActionId,
74    pub domain: DomainId,
75    pub trace_id: Option<TraceId>,
76    pub result: Result<ActionValue, ActionError>,
77    pub cancelled: bool,
78    pub local: bool,
79    pub object: Option<ObjectJob>,
80}
81
82struct DomainCounters {
83    queued: AtomicUsize,
84    running: AtomicUsize,
85    inflight: AtomicUsize,
86    completed: AtomicU64,
87    cancelled: AtomicU64,
88    queue_latency: AtomicU64,
89    max_queue_latency: AtomicU64,
90    execution: AtomicU64,
91    max_execution: AtomicU64,
92}
93
94impl DomainCounters {
95    fn snapshot(&self) -> DomainStats {
96        let queued = self.queued.load(Ordering::Acquire);
97        let running = self.running.load(Ordering::Acquire);
98        DomainStats {
99            queued,
100            running,
101            pending_completions: self
102                .inflight
103                .load(Ordering::Acquire)
104                .saturating_sub(queued.saturating_add(running)),
105            completed: self.completed.load(Ordering::Relaxed),
106            cancelled_before_start: self.cancelled.load(Ordering::Relaxed),
107            total_queue_latency_ns: self.queue_latency.load(Ordering::Relaxed),
108            max_queue_latency_ns: self.max_queue_latency.load(Ordering::Relaxed),
109            total_execution_ns: self.execution.load(Ordering::Relaxed),
110            max_execution_ns: self.max_execution.load(Ordering::Relaxed),
111        }
112    }
113}
114
115struct DomainExecutor {
116    pool: Arc<ThreadPool>,
117    completions: SyncSender<ActionCompletion>,
118    counters: Arc<DomainCounters>,
119    queue_capacity: usize,
120}
121
122pub(crate) struct DomainRegistry {
123    domains: BTreeMap<DomainId, DomainExecutor>,
124    completions: Receiver<ActionCompletion>,
125}
126
127impl DomainRegistry {
128    pub(crate) fn new(
129        configs: &[DomainConfig],
130        completion_capacity: usize,
131    ) -> Result<Self, RuntimeError> {
132        let (completion_tx, completions) = mpsc::sync_channel(completion_capacity);
133        let mut domains = BTreeMap::new();
134        for config in configs {
135            if config.workers == 0 || config.queue_capacity < config.workers {
136                return Err(RuntimeError::InvalidLimits(
137                    "domain workers must be nonzero and fit in its queue capacity",
138                ));
139            }
140            if domains.contains_key(&config.id) {
141                return Err(RuntimeError::InvalidLimits("duplicate domain id"));
142            }
143            let domain = config.id;
144            let pool = Arc::new(
145                ThreadPoolBuilder::new()
146                    .num_threads(config.workers)
147                    .thread_name(move |worker| format!("hataori-domain-{}-{worker}", domain.get()))
148                    .build()
149                    .map_err(|_| RuntimeError::InvalidLimits("failed to start domain workers"))?,
150            );
151            domains.insert(
152                config.id,
153                DomainExecutor {
154                    pool,
155                    completions: completion_tx.clone(),
156                    counters: Arc::new(DomainCounters {
157                        queued: AtomicUsize::new(0),
158                        running: AtomicUsize::new(0),
159                        inflight: AtomicUsize::new(0),
160                        completed: AtomicU64::new(0),
161                        cancelled: AtomicU64::new(0),
162                        queue_latency: AtomicU64::new(0),
163                        max_queue_latency: AtomicU64::new(0),
164                        execution: AtomicU64::new(0),
165                        max_execution: AtomicU64::new(0),
166                    }),
167                    queue_capacity: config.queue_capacity,
168                },
169            );
170        }
171        Ok(Self {
172            domains,
173            completions,
174        })
175    }
176
177    pub(crate) fn pools(&self) -> BTreeMap<DomainId, Arc<ThreadPool>> {
178        self.domains
179            .iter()
180            .map(|(id, domain)| (*id, Arc::clone(&domain.pool)))
181            .collect()
182    }
183
184    pub(crate) fn can_submit(&self, domain: DomainId) -> bool {
185        self.domains.get(&domain).is_some_and(|entry| {
186            entry.counters.queued.load(Ordering::Acquire) < entry.queue_capacity
187        })
188    }
189
190    // `fetch_update` is deprecated on newer toolchains in favor of `try_update`,
191    // which does not exist at the crate's MSRV (1.85).
192    #[allow(deprecated)]
193    pub(crate) fn submit(&self, job: ActionJob) -> Result<(), RuntimeError> {
194        let domain = self
195            .domains
196            .get(&job.domain)
197            .ok_or(RuntimeError::UnknownDomain(job.domain))?;
198        domain
199            .counters
200            .queued
201            .try_update_compat(Ordering::AcqRel, Ordering::Acquire, |queued| {
202                (queued < domain.queue_capacity).then_some(queued + 1)
203            })
204            .map_err(|_| RuntimeError::ResourceExhausted {
205                resource: ResourceKind::ActionQueue,
206                limit: domain.queue_capacity,
207            })?;
208        domain.counters.inflight.fetch_add(1, Ordering::AcqRel);
209        let completions = domain.completions.clone();
210        let counters = Arc::clone(&domain.counters);
211        domain
212            .pool
213            .spawn(move || execute_job(job, completions, counters));
214        Ok(())
215    }
216
217    pub(crate) fn try_completion(&self) -> Option<ActionCompletion> {
218        let completion = self.completions.try_recv().ok()?;
219        if let Some(domain) = self.domains.get(&completion.domain) {
220            domain.counters.inflight.fetch_sub(1, Ordering::AcqRel);
221        }
222        Some(completion)
223    }
224
225    pub(crate) fn stats(&self) -> Vec<(DomainId, DomainStats)> {
226        self.domains
227            .iter()
228            .map(|(id, domain)| (*id, domain.counters.snapshot()))
229            .collect()
230    }
231
232    pub(crate) fn idle(&self) -> bool {
233        self.domains
234            .values()
235            .all(|domain| domain.counters.inflight.load(Ordering::Acquire) == 0)
236    }
237
238    pub(crate) fn stop(&mut self) {
239        debug_assert!(self.idle(), "domain stop requires drained work");
240    }
241}
242
243fn execute_job(
244    job: ActionJob,
245    completions: SyncSender<ActionCompletion>,
246    counters: Arc<DomainCounters>,
247) {
248    counters.queued.fetch_sub(1, Ordering::AcqRel);
249    counters.running.fetch_add(1, Ordering::AcqRel);
250    let started_at = std::time::Instant::now();
251    let queue_latency = nanos(started_at.saturating_duration_since(job.submitted_at));
252    counters
253        .queue_latency
254        .fetch_add(queue_latency, Ordering::Relaxed);
255    counters
256        .max_queue_latency
257        .fetch_max(queue_latency, Ordering::Relaxed);
258    let cancelled = job.cancelled.load(Ordering::Acquire);
259    let result = if cancelled {
260        counters.cancelled.fetch_add(1, Ordering::Relaxed);
261        Err(ActionError::user("action cancelled before execution"))
262    } else {
263        job.handler.execute(job.input)
264    };
265    let execution = nanos(started_at.elapsed());
266    counters.execution.fetch_add(execution, Ordering::Relaxed);
267    counters
268        .max_execution
269        .fetch_max(execution, Ordering::Relaxed);
270    counters.running.fetch_sub(1, Ordering::AcqRel);
271    counters.completed.fetch_add(1, Ordering::Relaxed);
272    let _ = completions.send(ActionCompletion {
273        requester: job.requester,
274        request: job.request,
275        action_id: job.action_id,
276        domain: job.domain,
277        trace_id: job.trace_id,
278        result,
279        cancelled,
280        local: job.local,
281        object: job.object,
282    });
283}
284
285fn nanos(duration: std::time::Duration) -> u64 {
286    u64::try_from(duration.as_nanos()).unwrap_or(u64::MAX)
287}
288
289#[cfg(test)]
290mod tests;
hataori_runtime/domain.rs

hataori_runtime/
domain.rs
