datum_agent/
job.rs

1use std::{
2    fmt,
3    sync::Arc,
4    time::{Duration, Instant, SystemTime},
5};
6
7use datum::{
8    Flow, KillSwitches, NotUsed, RestartSettings, RunnableGraph, SharedKillSwitch,
9    StreamCompletion, StreamError, StreamInstrumentationRegistry, StreamResult,
10};
11
12/// Stable identifier assigned by the local registry when a job is submitted.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
14pub struct JobId(pub u64);
15
16/// Factory used by [`JobSpec`] to build a fresh job blueprint for one generation.
17///
18/// The factory receives a [`JobContext`] containing the generation identity and an agent-owned
19/// `SharedKillSwitch`. The factory should build a `RunnableGraph<JobMat>` and should not
20/// materialize it.
21pub type JobGraphFactory =
22    dyn Fn(JobContext) -> StreamResult<RunnableGraph<JobMat>> + Send + Sync + 'static;
23
24/// A named, supervised stream program.
25#[derive(Clone)]
26pub struct JobSpec {
27    pub name: String,
28    pub factory: Arc<JobGraphFactory>,
29    pub restart_policy: JobRestartPolicy,
30    pub drain_behavior: JobDrainBehavior,
31}
32
33impl JobSpec {
34    #[must_use]
35    pub fn new<F>(name: impl Into<String>, factory: F) -> Self
36    where
37        F: Fn(JobContext) -> StreamResult<RunnableGraph<JobMat>> + Send + Sync + 'static,
38    {
39        Self {
40            name: name.into(),
41            factory: Arc::new(factory),
42            restart_policy: JobRestartPolicy::Never,
43            drain_behavior: JobDrainBehavior::default(),
44        }
45    }
46
47    #[must_use]
48    pub fn with_restart_policy(mut self, restart_policy: JobRestartPolicy) -> Self {
49        self.restart_policy = restart_policy;
50        self
51    }
52
53    #[must_use]
54    pub fn with_drain_behavior(mut self, drain_behavior: JobDrainBehavior) -> Self {
55        self.drain_behavior = drain_behavior;
56        self
57    }
58}
59
60impl fmt::Debug for JobSpec {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        f.debug_struct("JobSpec")
63            .field("name", &self.name)
64            .field("restart_policy", &self.restart_policy)
65            .field("drain_behavior", &self.drain_behavior)
66            .finish_non_exhaustive()
67    }
68}
69
70/// Context passed to a job blueprint factory for one materialized generation.
71#[derive(Clone, Debug)]
72pub struct JobContext {
73    name: Arc<str>,
74    job_id: JobId,
75    generation: u64,
76    kill_switch: SharedKillSwitch,
77    instrumentation: StreamInstrumentationRegistry,
78}
79
80impl JobContext {
81    pub(crate) fn new(
82        name: impl Into<Arc<str>>,
83        job_id: JobId,
84        generation: u64,
85        kill_switch: SharedKillSwitch,
86        instrumentation: StreamInstrumentationRegistry,
87    ) -> Self {
88        Self {
89            name: name.into(),
90            job_id,
91            generation,
92            kill_switch,
93            instrumentation,
94        }
95    }
96
97    #[must_use]
98    pub fn name(&self) -> &str {
99        &self.name
100    }
101
102    #[must_use]
103    pub fn job_id(&self) -> JobId {
104        self.job_id
105    }
106
107    #[must_use]
108    pub fn generation(&self) -> u64 {
109        self.generation
110    }
111
112    #[must_use]
113    pub fn kill_switch(&self) -> SharedKillSwitch {
114        self.kill_switch.clone()
115    }
116
117    #[must_use]
118    pub fn instrumentation_registry(&self) -> &StreamInstrumentationRegistry {
119        &self.instrumentation
120    }
121
122    /// Return a flow backed by the agent-owned shared kill switch.
123    ///
124    /// Graph factories should wire this flow into the job path when graceful drain is supported.
125    #[must_use]
126    pub fn drain_flow<T: Send + 'static>(&self) -> Flow<T, T, SharedKillSwitch> {
127        self.kill_switch.flow()
128    }
129
130    #[must_use]
131    pub fn control(&self) -> JobControl {
132        JobControl::graceful(self.kill_switch.clone())
133    }
134}
135
136/// Standardized materialized value for daemon-managed jobs.
137pub struct JobMat {
138    completion: StreamCompletion<NotUsed>,
139    control: JobControl,
140}
141
142impl JobMat {
143    #[must_use]
144    pub fn new(completion: StreamCompletion<NotUsed>, control: JobControl) -> Self {
145        Self {
146            completion,
147            control,
148        }
149    }
150
151    #[must_use]
152    pub fn graceful(completion: StreamCompletion<NotUsed>, kill_switch: SharedKillSwitch) -> Self {
153        Self::new(completion, JobControl::graceful(kill_switch))
154    }
155
156    #[must_use]
157    pub fn cancel_only(completion: StreamCompletion<NotUsed>) -> Self {
158        Self::new(completion, JobControl::cancel_only())
159    }
160
161    pub(crate) fn into_parts(self) -> (StreamCompletion<NotUsed>, JobControl) {
162        (self.completion, self.control)
163    }
164}
165
166impl fmt::Debug for JobMat {
167    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
168        f.debug_struct("JobMat")
169            .field("control", &self.control)
170            .finish_non_exhaustive()
171    }
172}
173
174/// Control handles associated with a running job generation.
175#[derive(Clone, Debug)]
176pub struct JobControl {
177    kill_switch: Option<SharedKillSwitch>,
178}
179
180impl JobControl {
181    #[must_use]
182    pub fn graceful(kill_switch: SharedKillSwitch) -> Self {
183        Self {
184            kill_switch: Some(kill_switch),
185        }
186    }
187
188    #[must_use]
189    pub fn cancel_only() -> Self {
190        Self { kill_switch: None }
191    }
192
193    #[must_use]
194    pub fn drain_supported(&self) -> bool {
195        self.kill_switch.is_some()
196    }
197
198    #[must_use]
199    pub fn kill_switch(&self) -> Option<SharedKillSwitch> {
200        self.kill_switch.clone()
201    }
202
203    pub(crate) fn shutdown(&self) -> bool {
204        if let Some(kill_switch) = &self.kill_switch {
205            kill_switch.shutdown();
206            true
207        } else {
208            false
209        }
210    }
211
212    pub(crate) fn abort(&self, error: StreamError) -> bool {
213        if let Some(kill_switch) = &self.kill_switch {
214            kill_switch.abort(error);
215            true
216        } else {
217            false
218        }
219    }
220}
221
222/// Registry-visible restart policy for a job.
223#[derive(Clone, Debug)]
224pub enum JobRestartPolicy {
225    Never,
226    OnFailure(RestartSettings),
227    Always(RestartSettings),
228    Manual,
229}
230
231impl JobRestartPolicy {
232    #[must_use]
233    pub fn never() -> Self {
234        Self::Never
235    }
236
237    #[must_use]
238    pub fn on_failure(settings: RestartSettings) -> Self {
239        Self::OnFailure(settings)
240    }
241
242    #[must_use]
243    pub fn always(settings: RestartSettings) -> Self {
244        Self::Always(settings)
245    }
246
247    #[must_use]
248    pub fn manual() -> Self {
249        Self::Manual
250    }
251
252    pub(crate) fn settings_for(&self, cause: RestartCause) -> Option<RestartSettings> {
253        match (self, cause) {
254            (Self::Always(settings), RestartCause::Failure | RestartCause::Completion)
255            | (Self::OnFailure(settings), RestartCause::Failure) => Some(settings.clone()),
256            (Self::Never | Self::Manual | Self::OnFailure(_), _) => None,
257        }
258    }
259}
260
261#[derive(Clone, Copy, Debug, PartialEq, Eq)]
262pub(crate) enum RestartCause {
263    Failure,
264    Completion,
265}
266
267/// Drain behavior declared by a job spec.
268#[derive(Clone, Debug, PartialEq, Eq)]
269pub enum JobDrainBehavior {
270    Graceful { timeout: Duration },
271    CancelOnly,
272}
273
274impl JobDrainBehavior {
275    #[must_use]
276    pub fn graceful(timeout: Duration) -> Self {
277        Self::Graceful { timeout }
278    }
279
280    #[must_use]
281    pub fn cancel_only() -> Self {
282        Self::CancelOnly
283    }
284
285    #[must_use]
286    pub fn timeout(&self) -> Option<Duration> {
287        match self {
288            Self::Graceful { timeout } => Some(*timeout),
289            Self::CancelOnly => None,
290        }
291    }
292}
293
294impl Default for JobDrainBehavior {
295    fn default() -> Self {
296        Self::Graceful {
297            timeout: Duration::from_secs(30),
298        }
299    }
300}
301
302/// Desired state recorded by the registry.
303#[derive(Debug, Clone, Copy, PartialEq, Eq)]
304pub enum DesiredJobState {
305    Running,
306    Draining,
307    Stopped,
308}
309
310/// Observed lifecycle state of a job.
311#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312pub enum JobState {
313    Submitted,
314    Starting,
315    Running,
316    Draining,
317    BackingOff,
318    Completed,
319    Drained,
320    Stopped,
321    Failed,
322}
323
324/// Why the latest generation exited.
325#[derive(Debug, Clone, PartialEq, Eq)]
326pub enum JobExitReason {
327    Completed,
328    Failed(StreamError),
329    Drained,
330    Stopped,
331    DrainTimedOut,
332}
333
334/// Reserved integration point for WP-A1 stream instrumentation.
335///
336/// WP-A1 is not present in the current main branch, so WP-A2 intentionally exposes no per-element
337/// counters yet. Future instrumentation handles should fill this snapshot without changing the
338/// registry control-plane API.
339#[derive(Debug, Clone, Default, PartialEq, Eq)]
340#[non_exhaustive]
341pub struct JobInstrumentationSnapshot {}
342
343/// Point-in-time registry snapshot for one job.
344#[derive(Debug, Clone, PartialEq, Eq)]
345pub struct JobStatus {
346    pub name: String,
347    pub job_id: JobId,
348    pub state: JobState,
349    pub desired_state: DesiredJobState,
350    pub generation: u64,
351    pub starts_total: u64,
352    pub restarts_total: u64,
353    pub last_start_at: Option<SystemTime>,
354    pub last_exit_at: Option<SystemTime>,
355    pub last_exit_reason: Option<JobExitReason>,
356    pub backoff_until: Option<Instant>,
357    pub drain_deadline: Option<Instant>,
358    pub drain_supported: bool,
359    pub active_streams: Option<usize>,
360    pub instrumentation: Option<JobInstrumentationSnapshot>,
361}
362
363/// Lifecycle event emitted by the registry.
364#[derive(Debug, Clone, PartialEq, Eq)]
365pub struct JobEvent {
366    pub sequence: u64,
367    pub timestamp: SystemTime,
368    pub name: String,
369    pub job_id: JobId,
370    pub generation: u64,
371    pub kind: JobEventKind,
372}
373
374/// Kind-specific lifecycle event details.
375#[derive(Debug, Clone, PartialEq, Eq)]
376pub enum JobEventKind {
377    Submitted,
378    Started,
379    Failed { reason: JobExitReason },
380    RestartScheduled { delay: Duration },
381    Restarted { previous_generation: u64 },
382    Draining,
383    Drained,
384    Stopped { reason: JobExitReason },
385    Completed,
386}
387
388pub(crate) fn new_generation_kill_switch(name: &str, generation: u64) -> SharedKillSwitch {
389    KillSwitches::shared(format!("job:{name}:{generation}"))
390}
datum_agent/job.rs

datum_agent/
job.rs
