datum_agent/
job.rs

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

datum_agent/
job.rs
