Skip to main content

datum_agent/
registry.rs

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
29/// Ractor-backed local job registry.
30pub 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/// Cloneable in-process handle to the registry actor.
55#[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}