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