1use std::{
2 collections::HashMap,
3 future::Future,
4 panic::{AssertUnwindSafe, catch_unwind},
5 sync::{
6 OnceLock,
7 atomic::{AtomicU64, Ordering},
8 mpsc,
9 },
10 time::{Duration, Instant, SystemTime},
11};
12
13use datum::{
14 Materializer, NotUsed, RunnableGraph, Source, StreamError, StreamInstrumentationRegistry,
15 StreamResult, Topic, TopicOverflow, TopicTryPublishError,
16};
17use ractor::{Actor, ActorProcessingErr, ActorRef};
18use tokio::runtime::{Builder as TokioRuntimeBuilder, Runtime as TokioRuntime};
19
20use crate::{
21 AgentConfig, AgentError, AgentResult, ClusterJobMetadata, DesiredJobState, JobContext,
22 JobControl, JobDrainBehavior, JobEvent, JobEventKind, JobExitReason, JobId, JobMat, JobSpec,
23 JobState, JobStatus,
24 job::{RestartCause, new_generation_kill_switch},
25};
26
27type Reply<T> = mpsc::Sender<AgentResult<T>>;
28
29pub struct JobRegistry;
31
32impl JobRegistry {
33 pub fn start(config: AgentConfig) -> AgentResult<JobRegistryHandle> {
34 let event_buffer = config.event_buffer.max(1);
35 let events = Topic::new(event_buffer, TopicOverflow::Sliding)?;
36 let args = RegistryArgs {
37 instrumentation: config.instrumentation.clone(),
38 config,
39 events: events.clone(),
40 materializer: Materializer::new(),
41 };
42 let (actor, _join) =
43 block_on_agent_runtime(RegistryActor::spawn(None, RegistryActor, args)).and_then(
44 |result| {
45 result.map_err(|error| {
46 AgentError::Actor(format!("job registry failed to spawn: {error}"))
47 })
48 },
49 )?;
50 Ok(JobRegistryHandle { actor, events })
51 }
52}
53
54#[derive(Clone)]
56pub struct JobRegistryHandle {
57 actor: ActorRef<RegistryMessage>,
58 events: Topic<JobEvent>,
59}
60
61impl JobRegistryHandle {
62 pub fn submit(&self, spec: JobSpec) -> AgentResult<JobStatus> {
63 self.request(|reply| RegistryMessage::Submit {
64 spec: Box::new(spec),
65 reply,
66 })
67 }
68
69 pub fn start(&self, name: impl Into<String>) -> AgentResult<JobStatus> {
70 let name = name.into();
71 self.request(|reply| RegistryMessage::Start { name, reply })
72 }
73
74 pub fn restart(&self, name: impl Into<String>) -> AgentResult<JobStatus> {
75 let name = name.into();
76 self.request(|reply| RegistryMessage::Restart { name, reply })
77 }
78
79 pub fn drain(&self, name: impl Into<String>) -> AgentResult<JobStatus> {
80 let name = name.into();
81 self.request(|reply| RegistryMessage::Drain { name, reply })
82 }
83
84 pub fn stop(&self, name: impl Into<String>) -> AgentResult<JobStatus> {
85 let name = name.into();
86 self.request(|reply| RegistryMessage::Stop { name, reply })
87 }
88
89 pub fn status(&self, name: impl Into<String>) -> AgentResult<JobStatus> {
90 let name = name.into();
91 self.request(|reply| RegistryMessage::Status { name, reply })
92 }
93
94 pub fn list(&self) -> AgentResult<Vec<JobStatus>> {
95 self.request(|reply| RegistryMessage::List { reply })
96 }
97
98 pub(crate) fn update_cluster_metadata(
99 &self,
100 name: impl Into<String>,
101 metadata: ClusterJobMetadata,
102 ) -> AgentResult<JobStatus> {
103 self.request(|reply| RegistryMessage::UpdateClusterMetadata {
104 name: name.into(),
105 metadata: Box::new(metadata),
106 reply,
107 })
108 }
109
110 #[must_use]
111 pub fn events(&self) -> Source<JobEvent> {
112 self.events.subscribe()
113 }
114
115 pub fn shutdown(&self) -> AgentResult<()> {
116 self.request(RegistryMessage::Shutdown)
117 }
118
119 fn request<T, F>(&self, make: F) -> AgentResult<T>
120 where
121 T: Send + 'static,
122 F: FnOnce(Reply<T>) -> RegistryMessage,
123 {
124 let (reply, receiver) = mpsc::channel();
125 self.actor
126 .send_message(make(reply))
127 .map_err(|_| AgentError::RegistryClosed)?;
128 receiver.recv().map_err(|_| AgentError::RegistryClosed)?
129 }
130}
131
132struct RegistryArgs {
133 instrumentation: StreamInstrumentationRegistry,
134 config: AgentConfig,
135 events: Topic<JobEvent>,
136 materializer: Materializer,
137}
138
139struct RegistryActor;
140
141impl Actor for RegistryActor {
142 type Msg = RegistryMessage;
143 type State = RegistryState;
144 type Arguments = RegistryArgs;
145
146 async fn pre_start(
147 &self,
148 _myself: ActorRef<Self::Msg>,
149 args: Self::Arguments,
150 ) -> Result<Self::State, ActorProcessingErr> {
151 Ok(RegistryState {
152 config: args.config,
153 materializer: args.materializer,
154 events: args.events,
155 instrumentation: args.instrumentation,
156 next_job_id: 1,
157 next_event_sequence: 1,
158 jobs: HashMap::new(),
159 })
160 }
161
162 async fn post_start(
163 &self,
164 myself: ActorRef<Self::Msg>,
165 state: &mut Self::State,
166 ) -> Result<(), ActorProcessingErr> {
167 let interval = state.config.poll_interval.max(Duration::from_millis(1));
168 spawn_completion_poller(myself, interval);
169 Ok(())
170 }
171
172 async fn post_stop(
173 &self,
174 _myself: ActorRef<Self::Msg>,
175 state: &mut Self::State,
176 ) -> Result<(), ActorProcessingErr> {
177 abort_all_jobs(state);
178 let materializer = state.materializer.clone();
179 let events = state.events.clone();
180 materializer.shutdown();
181 let _ = events.close();
182 Ok(())
183 }
184
185 async fn handle(
186 &self,
187 myself: ActorRef<Self::Msg>,
188 message: Self::Msg,
189 state: &mut Self::State,
190 ) -> Result<(), ActorProcessingErr> {
191 match message {
192 RegistryMessage::Submit { spec, reply } => {
193 send_reply(reply, submit_job(state, *spec));
194 }
195 RegistryMessage::Start { name, reply } => {
196 send_reply(reply, start_job(state, &myself, &name, StartKind::Initial));
197 }
198 RegistryMessage::Restart { name, reply } => {
199 send_reply(reply, restart_job(state, &myself, &name));
200 }
201 RegistryMessage::Drain { name, reply } => {
202 send_reply(reply, drain_job(state, &name));
203 }
204 RegistryMessage::Stop { name, reply } => {
205 send_reply(reply, stop_job(state, &name, JobExitReason::Stopped));
206 }
207 RegistryMessage::Status { name, reply } => {
208 send_reply(reply, status_by_name(state, &name));
209 }
210 RegistryMessage::List { reply } => {
211 send_reply(reply, Ok(list_statuses(state)));
212 }
213 RegistryMessage::UpdateClusterMetadata {
214 name,
215 metadata,
216 reply,
217 } => {
218 send_reply(reply, update_cluster_metadata(state, &name, *metadata));
219 }
220 RegistryMessage::Tick => {
221 poll_jobs(state, &myself);
222 }
223 RegistryMessage::RestartReady {
224 name,
225 expected_generation,
226 } => {
227 restart_ready(state, &myself, &name, expected_generation);
228 }
229 RegistryMessage::Shutdown(reply) => {
230 abort_all_jobs(state);
231 let _ = reply.send(Ok(()));
232 myself.stop(Some("datum-agent shutdown".to_owned()));
233 }
234 }
235 Ok(())
236 }
237}
238
239enum RegistryMessage {
240 Submit {
241 spec: Box<JobSpec>,
242 reply: Reply<JobStatus>,
243 },
244 Start {
245 name: String,
246 reply: Reply<JobStatus>,
247 },
248 Restart {
249 name: String,
250 reply: Reply<JobStatus>,
251 },
252 Drain {
253 name: String,
254 reply: Reply<JobStatus>,
255 },
256 Stop {
257 name: String,
258 reply: Reply<JobStatus>,
259 },
260 Status {
261 name: String,
262 reply: Reply<JobStatus>,
263 },
264 List {
265 reply: Reply<Vec<JobStatus>>,
266 },
267 UpdateClusterMetadata {
268 name: String,
269 metadata: Box<ClusterJobMetadata>,
270 reply: Reply<JobStatus>,
271 },
272 Tick,
273 RestartReady {
274 name: String,
275 expected_generation: u64,
276 },
277 Shutdown(Reply<()>),
278}
279
280struct RegistryState {
281 config: AgentConfig,
282 materializer: Materializer,
283 events: Topic<JobEvent>,
284 instrumentation: StreamInstrumentationRegistry,
285 next_job_id: u64,
286 next_event_sequence: u64,
287 jobs: HashMap<String, JobEntry>,
288}
289
290struct JobEntry {
291 id: JobId,
292 spec: JobSpec,
293 state: JobState,
294 desired_state: DesiredJobState,
295 generation: u64,
296 run: Option<JobRun>,
297 backoff: RestartBackoffState,
298 starts_total: u64,
299 restarts_total: u64,
300 last_start_at: Option<SystemTime>,
301 last_exit_at: Option<SystemTime>,
302 last_exit_reason: Option<JobExitReason>,
303 backoff_until: Option<Instant>,
304 drain_deadline: Option<Instant>,
305 drain_supported: bool,
306}
307
308impl JobEntry {
309 fn new(id: JobId, spec: JobSpec) -> Self {
310 Self {
311 id,
312 spec,
313 state: JobState::Submitted,
314 desired_state: DesiredJobState::Stopped,
315 generation: 0,
316 run: None,
317 backoff: RestartBackoffState::default(),
318 starts_total: 0,
319 restarts_total: 0,
320 last_start_at: None,
321 last_exit_at: None,
322 last_exit_reason: None,
323 backoff_until: None,
324 drain_deadline: None,
325 drain_supported: true,
326 }
327 }
328}
329
330struct JobRun {
331 generation: u64,
332 control: JobControl,
333 completion: datum::StreamCompletion<NotUsed>,
334}
335
336#[derive(Default)]
337struct RestartBackoffState {
338 window_start: Option<Instant>,
339 restarts_in_window: usize,
340}
341
342impl RestartBackoffState {
343 fn reset(&mut self) {
344 self.window_start = None;
345 self.restarts_in_window = 0;
346 }
347
348 fn next_delay(&mut self, settings: &datum::RestartSettings) -> Option<Duration> {
349 let now = Instant::now();
350 match self.window_start {
351 Some(window_start)
352 if now.duration_since(window_start) <= settings.max_restarts_within() => {}
353 _ => {
354 self.window_start = Some(now);
355 self.restarts_in_window = 0;
356 }
357 }
358
359 if self.restarts_in_window >= settings.max_restarts() {
360 return None;
361 }
362
363 let exponent = (self.restarts_in_window.min(31) as i32).max(0);
364 let base = settings
365 .min_backoff()
366 .mul_f64(2_f64.powi(exponent))
367 .min(settings.max_backoff());
368 self.restarts_in_window += 1;
369
370 if settings.random_factor() == 0.0 || base.is_zero() {
371 Some(base)
372 } else {
373 let jitter = base.mul_f64(next_random_fraction() * settings.random_factor());
374 Some(base.saturating_add(jitter))
375 }
376 }
377}
378
379#[derive(Clone, Copy)]
380enum StartKind {
381 Initial,
382 AutomaticRestart { previous_generation: u64 },
383 ManualRestart,
384}
385
386fn submit_job(state: &mut RegistryState, spec: JobSpec) -> AgentResult<JobStatus> {
387 if spec.name.trim().is_empty() {
388 return Err(AgentError::InvalidJobName);
389 }
390 if state.jobs.contains_key(&spec.name) {
391 return Err(AgentError::JobAlreadyExists(spec.name));
392 }
393
394 let name = spec.name.clone();
395 let id = JobId(state.next_job_id);
396 state.next_job_id += 1;
397 state.jobs.insert(name.clone(), JobEntry::new(id, spec));
398 emit_event(state, &name, 0, JobEventKind::Submitted);
399 status_by_name(state, &name)
400}
401
402fn update_cluster_metadata(
403 state: &mut RegistryState,
404 name: &str,
405 metadata: ClusterJobMetadata,
406) -> AgentResult<JobStatus> {
407 let entry = state
408 .jobs
409 .get_mut(name)
410 .ok_or_else(|| AgentError::JobNotFound(name.to_owned()))?;
411 entry.spec.cluster = Some(metadata);
412 status_by_name(state, name)
413}
414
415fn start_job(
416 state: &mut RegistryState,
417 myself: &ActorRef<RegistryMessage>,
418 name: &str,
419 kind: StartKind,
420) -> AgentResult<JobStatus> {
421 let (spec, job_id, generation) = {
422 let entry = state
423 .jobs
424 .get_mut(name)
425 .ok_or_else(|| AgentError::JobNotFound(name.to_owned()))?;
426
427 if matches!(
428 entry.state,
429 JobState::Starting | JobState::Running | JobState::Draining | JobState::BackingOff
430 ) && !matches!(kind, StartKind::AutomaticRestart { .. })
431 {
432 return Err(AgentError::JobAlreadyRunning(name.to_owned()));
433 }
434
435 entry.state = JobState::Starting;
436 entry.desired_state = DesiredJobState::Running;
437 entry.generation += 1;
438 entry.backoff_until = None;
439 entry.drain_deadline = None;
440 entry.run = None;
441 (entry.spec.clone(), entry.id, entry.generation)
442 };
443
444 let kill_switch = new_generation_kill_switch(&spec.name, generation);
445 let context = JobContext::new(
446 spec.name.clone(),
447 job_id,
448 generation,
449 kill_switch,
450 state.instrumentation.clone(),
451 );
452 let graph = invoke_factory(&spec, context);
453 let materializer = state
454 .materializer
455 .with_name_prefix(format!("datum-agent-{}", spec.name));
456 let materialized = graph.and_then(|graph| materialize_graph(&graph, &materializer));
457
458 match materialized {
459 Ok(mat) => {
460 let (completion, control) = mat.into_parts();
461 let mut restart_event = None;
462 {
463 let entry = state
464 .jobs
465 .get_mut(name)
466 .ok_or_else(|| AgentError::JobNotFound(name.to_owned()))?;
467 entry.run = Some(JobRun {
468 generation,
469 control: control.clone(),
470 completion,
471 });
472 entry.state = JobState::Running;
473 entry.desired_state = DesiredJobState::Running;
474 entry.starts_total += 1;
475 entry.last_start_at = Some(SystemTime::now());
476 entry.drain_supported = control.drain_supported();
477 if let StartKind::AutomaticRestart {
478 previous_generation,
479 } = kind
480 {
481 entry.restarts_total += 1;
482 restart_event = Some(JobEventKind::Restarted {
483 previous_generation,
484 });
485 }
486 }
487 if let Some(kind) = restart_event {
488 emit_event(state, name, generation, kind);
489 }
490 emit_event(state, name, generation, JobEventKind::Started);
491 status_by_name(state, name)
492 }
493 Err(error) => {
494 handle_start_failure(state, myself, name, generation, error.clone());
495 Err(AgentError::Stream(error))
496 }
497 }
498}
499
500fn restart_job(
501 state: &mut RegistryState,
502 myself: &ActorRef<RegistryMessage>,
503 name: &str,
504) -> AgentResult<JobStatus> {
505 if !state.jobs.contains_key(name) {
506 return Err(AgentError::JobNotFound(name.to_owned()));
507 }
508 let _ = stop_job(state, name, JobExitReason::Stopped)?;
509 if let Some(entry) = state.jobs.get_mut(name) {
510 entry.backoff.reset();
511 }
512 start_job(state, myself, name, StartKind::ManualRestart)
513}
514
515fn drain_job(state: &mut RegistryState, name: &str) -> AgentResult<JobStatus> {
516 let generation = {
517 let entry = state
518 .jobs
519 .get_mut(name)
520 .ok_or_else(|| AgentError::JobNotFound(name.to_owned()))?;
521
522 let Some(run) = entry.run.as_ref() else {
523 entry.state = JobState::Drained;
524 entry.desired_state = DesiredJobState::Stopped;
525 entry.last_exit_at = Some(SystemTime::now());
526 entry.last_exit_reason = Some(JobExitReason::Drained);
527 return status_by_name(state, name);
528 };
529
530 let generation = run.generation;
531 let drain_supported = run.control.drain_supported();
532 let control = run.control.clone();
533 let timeout = match entry.spec.drain_behavior {
534 JobDrainBehavior::Graceful { timeout } if drain_supported => timeout,
535 JobDrainBehavior::Graceful { .. } | JobDrainBehavior::CancelOnly => {
536 return Err(AgentError::DrainUnsupported(name.to_owned()));
537 }
538 };
539
540 if !control.shutdown() {
541 return Err(AgentError::DrainUnsupported(name.to_owned()));
542 }
543
544 entry.state = JobState::Draining;
545 entry.desired_state = DesiredJobState::Draining;
546 let deadline = Instant::now() + timeout;
547 entry.drain_deadline = Some(deadline);
548 generation
549 };
550
551 emit_event(state, name, generation, JobEventKind::Draining);
552 status_by_name(state, name)
553}
554
555fn stop_job(
556 state: &mut RegistryState,
557 name: &str,
558 reason: JobExitReason,
559) -> AgentResult<JobStatus> {
560 let generation = {
561 let entry = state
562 .jobs
563 .get_mut(name)
564 .ok_or_else(|| AgentError::JobNotFound(name.to_owned()))?;
565
566 let generation = entry.generation;
567 if let Some(run) = entry.run.take() {
568 let _ = run.control.abort(StreamError::Cancelled);
569 drop(run.completion);
570 }
571 entry.state = JobState::Stopped;
572 entry.desired_state = DesiredJobState::Stopped;
573 entry.backoff_until = None;
574 entry.drain_deadline = None;
575 entry.last_exit_at = Some(SystemTime::now());
576 entry.last_exit_reason = Some(reason.clone());
577 generation
578 };
579
580 emit_event(state, name, generation, JobEventKind::Stopped { reason });
581 status_by_name(state, name)
582}
583
584fn status_by_name(state: &RegistryState, name: &str) -> AgentResult<JobStatus> {
585 state
586 .jobs
587 .get(name)
588 .map(|entry| snapshot(state, entry))
589 .ok_or_else(|| AgentError::JobNotFound(name.to_owned()))
590}
591
592fn list_statuses(state: &RegistryState) -> Vec<JobStatus> {
593 let mut statuses = state
594 .jobs
595 .values()
596 .map(|entry| snapshot(state, entry))
597 .collect::<Vec<_>>();
598 statuses.sort_by(|left, right| left.name.cmp(&right.name));
599 statuses
600}
601
602fn poll_jobs(state: &mut RegistryState, myself: &ActorRef<RegistryMessage>) {
603 let names = state.jobs.keys().cloned().collect::<Vec<_>>();
604 for name in names {
605 let completed = state.jobs.get_mut(&name).and_then(|entry| {
606 entry.run.as_mut().and_then(|run| {
607 run.completion
608 .try_wait()
609 .map(|result| (run.generation, result))
610 })
611 });
612
613 if let Some((generation, result)) = completed {
614 finish_generation(state, myself, &name, generation, result);
615 }
616 }
617
618 let now = Instant::now();
619 let timed_out = state
620 .jobs
621 .iter()
622 .filter_map(|(name, entry)| match (entry.state, entry.drain_deadline) {
623 (JobState::Draining, Some(deadline)) if deadline <= now => Some(name.clone()),
624 _ => None,
625 })
626 .collect::<Vec<_>>();
627
628 for name in timed_out {
629 let _ = stop_job(state, &name, JobExitReason::DrainTimedOut);
630 }
631}
632
633fn finish_generation(
634 state: &mut RegistryState,
635 myself: &ActorRef<RegistryMessage>,
636 name: &str,
637 generation: u64,
638 result: StreamResult<NotUsed>,
639) {
640 let mut restart = None;
641 let event;
642 {
643 let Some(entry) = state.jobs.get_mut(name) else {
644 return;
645 };
646 if entry.run.as_ref().map(|run| run.generation) != Some(generation) {
647 return;
648 }
649 let _run = entry.run.take();
650 entry.drain_deadline = None;
651 entry.last_exit_at = Some(SystemTime::now());
652
653 match (entry.desired_state, result) {
654 (DesiredJobState::Draining, Ok(_)) => {
655 entry.state = JobState::Drained;
656 entry.desired_state = DesiredJobState::Stopped;
657 entry.last_exit_reason = Some(JobExitReason::Drained);
658 event = JobEventKind::Drained;
659 }
660 (DesiredJobState::Draining | DesiredJobState::Stopped, Err(error)) => {
661 let reason = if matches!(error, StreamError::Cancelled) {
662 JobExitReason::Stopped
663 } else {
664 JobExitReason::Failed(error)
665 };
666 entry.state = JobState::Stopped;
667 entry.desired_state = DesiredJobState::Stopped;
668 entry.last_exit_reason = Some(reason.clone());
669 event = JobEventKind::Stopped { reason };
670 }
671 (DesiredJobState::Stopped, Ok(_)) => {
672 entry.state = JobState::Stopped;
673 entry.last_exit_reason = Some(JobExitReason::Stopped);
674 event = JobEventKind::Stopped {
675 reason: JobExitReason::Stopped,
676 };
677 }
678 (DesiredJobState::Running, Ok(_)) => {
679 entry.last_exit_reason = Some(JobExitReason::Completed);
680 entry.state = JobState::Completed;
681 event = JobEventKind::Completed;
682 if let Some(settings) = entry
683 .spec
684 .restart_policy
685 .settings_for(RestartCause::Completion)
686 {
687 restart = Some((settings, RestartCause::Completion));
688 }
689 }
690 (DesiredJobState::Running, Err(error)) => {
691 let reason = JobExitReason::Failed(error);
692 entry.last_exit_reason = Some(reason.clone());
693 entry.state = JobState::Failed;
694 event = JobEventKind::Failed { reason };
695 if let Some(settings) = entry
696 .spec
697 .restart_policy
698 .settings_for(RestartCause::Failure)
699 {
700 restart = Some((settings, RestartCause::Failure));
701 }
702 }
703 }
704 }
705
706 emit_event(state, name, generation, event);
707 if let Some((settings, cause)) = restart {
708 schedule_restart(state, myself, name, generation, settings, cause);
709 }
710}
711
712fn handle_start_failure(
713 state: &mut RegistryState,
714 myself: &ActorRef<RegistryMessage>,
715 name: &str,
716 generation: u64,
717 error: StreamError,
718) {
719 let mut restart = None;
720 {
721 let Some(entry) = state.jobs.get_mut(name) else {
722 return;
723 };
724 let reason = JobExitReason::Failed(error);
725 entry.state = JobState::Failed;
726 entry.run = None;
727 entry.last_exit_at = Some(SystemTime::now());
728 entry.last_exit_reason = Some(reason.clone());
729 if let Some(settings) = entry
730 .spec
731 .restart_policy
732 .settings_for(RestartCause::Failure)
733 {
734 restart = Some(settings);
735 }
736 }
737
738 let reason = state
739 .jobs
740 .get(name)
741 .and_then(|entry| entry.last_exit_reason.clone())
742 .unwrap_or(JobExitReason::Failed(StreamError::AbruptTermination));
743 emit_event(state, name, generation, JobEventKind::Failed { reason });
744 if let Some(settings) = restart {
745 schedule_restart(
746 state,
747 myself,
748 name,
749 generation,
750 settings,
751 RestartCause::Failure,
752 );
753 }
754}
755
756fn schedule_restart(
757 state: &mut RegistryState,
758 myself: &ActorRef<RegistryMessage>,
759 name: &str,
760 generation: u64,
761 settings: datum::RestartSettings,
762 _cause: RestartCause,
763) {
764 let delay = {
765 let Some(entry) = state.jobs.get_mut(name) else {
766 return;
767 };
768 if entry.desired_state != DesiredJobState::Running {
769 return;
770 }
771 let Some(delay) = entry.backoff.next_delay(&settings) else {
772 entry.backoff_until = None;
773 return;
774 };
775 entry.state = JobState::BackingOff;
776 entry.backoff_until = Some(Instant::now() + delay);
777 delay
778 };
779
780 emit_event(
781 state,
782 name,
783 generation,
784 JobEventKind::RestartScheduled { delay },
785 );
786
787 let actor = myself.clone();
788 let name = name.to_owned();
789 tokio::spawn(async move {
790 tokio::time::sleep(delay).await;
791 let _ = actor.send_message(RegistryMessage::RestartReady {
792 name,
793 expected_generation: generation,
794 });
795 });
796}
797
798fn restart_ready(
799 state: &mut RegistryState,
800 myself: &ActorRef<RegistryMessage>,
801 name: &str,
802 expected_generation: u64,
803) {
804 let ready = state.jobs.get(name).is_some_and(|entry| {
805 entry.generation == expected_generation
806 && entry.state == JobState::BackingOff
807 && entry.desired_state == DesiredJobState::Running
808 });
809 if ready {
810 let _ = start_job(
811 state,
812 myself,
813 name,
814 StartKind::AutomaticRestart {
815 previous_generation: expected_generation,
816 },
817 );
818 }
819}
820
821fn abort_all_jobs(state: &mut RegistryState) {
822 let names = state.jobs.keys().cloned().collect::<Vec<_>>();
823 for name in names {
824 let _ = stop_job(state, &name, JobExitReason::Stopped);
825 }
826}
827
828fn snapshot(_state: &RegistryState, entry: &JobEntry) -> JobStatus {
829 JobStatus {
830 name: entry.spec.name.clone(),
831 job_id: entry.id,
832 state: entry.state,
833 desired_state: entry.desired_state,
834 generation: entry.generation,
835 starts_total: entry.starts_total,
836 restarts_total: entry.restarts_total,
837 last_start_at: entry.last_start_at,
838 last_exit_at: entry.last_exit_at,
839 last_exit_reason: entry.last_exit_reason.clone(),
840 backoff_until: entry.backoff_until,
841 drain_deadline: entry.drain_deadline,
842 drain_supported: entry.drain_supported,
843 active_streams: Some(usize::from(entry.run.is_some())),
844 instrumentation: None,
845 cluster: entry.spec.cluster.clone(),
846 }
847}
848
849fn emit_event(state: &mut RegistryState, name: &str, generation: u64, kind: JobEventKind) {
850 let Some(entry) = state.jobs.get(name) else {
851 return;
852 };
853 let event = JobEvent {
854 sequence: state.next_event_sequence,
855 timestamp: SystemTime::now(),
856 name: entry.spec.name.clone(),
857 job_id: entry.id,
858 generation,
859 kind,
860 };
861 state.next_event_sequence += 1;
862 match state.events.try_publish(event) {
863 Ok(()) | Err(TopicTryPublishError::Full(_)) | Err(TopicTryPublishError::Busy(_)) => {}
864 Err(TopicTryPublishError::Closed(_)) => {}
865 }
866}
867
868fn invoke_factory(spec: &JobSpec, context: JobContext) -> StreamResult<RunnableGraph<JobMat>> {
869 catch_unwind(AssertUnwindSafe(|| (spec.factory)(context)))
870 .unwrap_or_else(|_| Err(StreamError::Failed("job factory panicked".to_owned())))
871}
872
873fn materialize_graph(
874 graph: &RunnableGraph<JobMat>,
875 materializer: &Materializer,
876) -> StreamResult<JobMat> {
877 catch_unwind(AssertUnwindSafe(|| {
878 graph.run_with_materializer(materializer)
879 }))
880 .unwrap_or_else(|_| {
881 Err(StreamError::Failed(
882 "job materialization panicked".to_owned(),
883 ))
884 })
885}
886
887fn send_reply<T>(reply: Reply<T>, result: AgentResult<T>) {
888 let _ = reply.send(result);
889}
890
891fn spawn_completion_poller(actor: ActorRef<RegistryMessage>, interval: Duration) {
892 tokio::spawn(async move {
893 loop {
894 tokio::time::sleep(interval).await;
895 if actor.send_message(RegistryMessage::Tick).is_err() {
896 break;
897 }
898 }
899 });
900}
901
902fn agent_runtime() -> AgentResult<&'static TokioRuntime> {
903 static RUNTIME: OnceLock<Result<TokioRuntime, String>> = OnceLock::new();
904 match RUNTIME.get_or_init(|| {
905 TokioRuntimeBuilder::new_multi_thread()
906 .thread_name("datum-agent-ractor")
907 .enable_all()
908 .build()
909 .map_err(|error| format!("agent runtime failed to start: {error}"))
910 }) {
911 Ok(runtime) => Ok(runtime),
912 Err(error) => Err(AgentError::Actor(error.clone())),
913 }
914}
915
916fn block_on_agent_runtime<T, F>(future: F) -> AgentResult<T>
917where
918 T: Send + 'static,
919 F: Future<Output = T> + Send + 'static,
920{
921 let runtime = agent_runtime()?;
922 if tokio::runtime::Handle::try_current().is_ok() {
923 std::thread::spawn(move || runtime.block_on(future))
924 .join()
925 .map_err(|_| AgentError::Actor("agent runtime task panicked".to_owned()))
926 } else {
927 Ok(runtime.block_on(future))
928 }
929}
930
931fn next_random_fraction() -> f64 {
932 static BACKOFF_RANDOM_SEED: AtomicU64 = AtomicU64::new(0x9e37_79b9_7f4a_7c15);
933 let mut current = BACKOFF_RANDOM_SEED.load(Ordering::Relaxed);
934 loop {
935 let mut next = current;
936 next ^= next << 13;
937 next ^= next >> 7;
938 next ^= next << 17;
939 match BACKOFF_RANDOM_SEED.compare_exchange_weak(
940 current,
941 next,
942 Ordering::Relaxed,
943 Ordering::Relaxed,
944 ) {
945 Ok(_) => return ((next >> 11) as f64) / ((1_u64 << 53) as f64),
946 Err(observed) => current = observed,
947 }
948 }
949}