1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::error::Error;
5use std::fmt::{Display, Formatter};
6use std::sync::Arc;
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub enum DagError {
10 DuplicateNode(String),
11 NodeNotFound(String),
12 PrecedingNodeNotFound(String),
13 CycleDetected,
14 EmptyGraph,
15 MissingNodeCompute(String),
16 NodeExecutionFailed {
17 node_id: String,
18 reason: String,
19 },
20 RetryExhausted {
21 node_id: String,
22 attempts: usize,
23 max_retries: usize,
24 retryable: bool,
25 last_error: String,
26 },
27 TaskJoinFailed {
28 node_id: String,
29 reason: String,
30 },
31 RunAlreadyInProgress {
32 lock_key: String,
33 },
34 RunGuardAcquireFailed {
35 lock_key: String,
36 reason: String,
37 },
38 RunGuardRenewFailed {
39 lock_key: String,
40 reason: String,
41 },
42 RunGuardReleaseFailed {
43 lock_key: String,
44 reason: String,
45 },
46 MissingRunFencingToken {
47 resource: String,
48 },
49 FencingTokenRejected {
50 resource: String,
51 token: u64,
52 reason: String,
53 },
54 InvalidPayloadEnvelope(String),
55 UnsupportedPayloadVersion(u8),
56 ReservedControlNode(String),
57 InvalidPrecedingControlNode(String),
58 RemoteDispatchNotConfigured(String),
59 NodeTimeout {
60 node_id: String,
61 timeout_ms: u64,
62 },
63 ExecutionTimeout {
64 run_id: String,
65 timeout_ms: u64,
66 },
67 ExecutionIncomplete {
68 run_id: String,
69 unfinished_nodes: usize,
70 },
71 InvalidStateTransition {
72 node_id: String,
73 from: DagNodeStatus,
74 to: DagNodeStatus,
75 },
76}
77
78impl Display for DagError {
79 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
80 match self {
81 DagError::DuplicateNode(id) => write!(f, "duplicate node id: {id}"),
82 DagError::NodeNotFound(id) => write!(f, "node not found: {id}"),
83 DagError::PrecedingNodeNotFound(id) => write!(f, "preceding node not found: {id}"),
84 DagError::CycleDetected => write!(f, "cycle detected in dag"),
85 DagError::EmptyGraph => write!(f, "dag graph is empty"),
86 DagError::MissingNodeCompute(id) => {
87 write!(f, "missing compute function for node: {id}")
88 }
89 DagError::NodeExecutionFailed { node_id, reason } => {
90 write!(f, "node execution failed: node={node_id} reason={reason}")
91 }
92 DagError::RetryExhausted {
93 node_id,
94 attempts,
95 max_retries,
96 retryable,
97 last_error,
98 } => {
99 write!(
100 f,
101 "node retry exhausted: node={node_id} attempts={attempts} max_retries={max_retries} retryable={retryable} last_error={last_error}"
102 )
103 }
104 DagError::TaskJoinFailed { node_id, reason } => {
105 write!(f, "node task join failed: node={node_id} reason={reason}")
106 }
107 DagError::RunAlreadyInProgress { lock_key } => {
108 write!(f, "dag run already in progress: lock_key={lock_key}")
109 }
110 DagError::RunGuardAcquireFailed { lock_key, reason } => {
111 write!(
112 f,
113 "dag run guard acquire failed: lock_key={lock_key} reason={reason}"
114 )
115 }
116 DagError::RunGuardRenewFailed { lock_key, reason } => {
117 write!(
118 f,
119 "dag run guard renew failed: lock_key={lock_key} reason={reason}"
120 )
121 }
122 DagError::RunGuardReleaseFailed { lock_key, reason } => {
123 write!(
124 f,
125 "dag run guard release failed: lock_key={lock_key} reason={reason}"
126 )
127 }
128 DagError::MissingRunFencingToken { resource } => {
129 write!(f, "missing run fencing token for resource: {resource}")
130 }
131 DagError::FencingTokenRejected {
132 resource,
133 token,
134 reason,
135 } => {
136 write!(
137 f,
138 "fencing token rejected: resource={resource} token={token} reason={reason}"
139 )
140 }
141 DagError::InvalidPayloadEnvelope(reason) => {
142 write!(f, "invalid payload envelope: {reason}")
143 }
144 DagError::UnsupportedPayloadVersion(version) => {
145 write!(f, "unsupported payload version: {version}")
146 }
147 DagError::ReservedControlNode(id) => {
148 write!(f, "node id is reserved as control node: {id}")
149 }
150 DagError::InvalidPrecedingControlNode(id) => {
151 write!(f, "invalid preceding control node: {id}")
152 }
153 DagError::RemoteDispatchNotConfigured(id) => {
154 write!(f, "remote dispatch not configured for node: {id}")
155 }
156 DagError::NodeTimeout {
157 node_id,
158 timeout_ms,
159 } => {
160 write!(
161 f,
162 "node execution timeout: node={node_id} timeout_ms={timeout_ms}"
163 )
164 }
165 DagError::ExecutionTimeout { run_id, timeout_ms } => {
166 write!(
167 f,
168 "dag execution timeout: run_id={run_id} timeout_ms={timeout_ms}"
169 )
170 }
171 DagError::ExecutionIncomplete {
172 run_id,
173 unfinished_nodes,
174 } => {
175 write!(
176 f,
177 "dag execution incomplete: run_id={run_id} unfinished_nodes={unfinished_nodes}"
178 )
179 }
180 DagError::InvalidStateTransition { node_id, from, to } => {
181 write!(
182 f,
183 "invalid node state transition: node={node_id} from={from:?} to={to:?}"
184 )
185 }
186 }
187 }
188}
189
190impl Error for DagError {}
191
192#[derive(Debug, Clone, Default, PartialEq, Eq)]
193pub struct TaskPayload {
194 pub bytes: Vec<u8>,
195 pub content_type: Option<String>,
196 pub encoding: Option<String>,
197 pub metadata: HashMap<String, String>,
198}
199
200impl TaskPayload {
201 const ENVELOPE_MAGIC: [u8; 4] = *b"MDAG";
202 const ENVELOPE_VERSION: u8 = 1;
203
204 pub fn from_bytes(bytes: Vec<u8>) -> Self {
205 Self {
206 bytes,
207 content_type: None,
208 encoding: None,
209 metadata: HashMap::new(),
210 }
211 }
212
213 pub fn with_content_type(mut self, content_type: impl AsRef<str>) -> Self {
214 self.content_type = Some(content_type.as_ref().to_string());
215 self
216 }
217
218 pub fn with_encoding(mut self, encoding: impl AsRef<str>) -> Self {
219 self.encoding = Some(encoding.as_ref().to_string());
220 self
221 }
222
223 pub fn with_meta(mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> Self {
224 self.metadata
225 .insert(key.as_ref().to_string(), value.as_ref().to_string());
226 self
227 }
228
229 pub fn to_envelope_bytes(&self) -> Result<Vec<u8>, DagError> {
230 fn to_u16_len(len: usize, field: &str) -> Result<u16, DagError> {
231 u16::try_from(len)
232 .map_err(|_| DagError::InvalidPayloadEnvelope(format!("{field} too large")))
233 }
234 fn to_u32_len(len: usize, field: &str) -> Result<u32, DagError> {
235 u32::try_from(len)
236 .map_err(|_| DagError::InvalidPayloadEnvelope(format!("{field} too large")))
237 }
238
239 let ct_bytes = self.content_type.as_deref().unwrap_or("").as_bytes();
240 let enc_bytes = self.encoding.as_deref().unwrap_or("").as_bytes();
241
242 let ct_len = to_u16_len(ct_bytes.len(), "content_type")?;
243 let enc_len = to_u16_len(enc_bytes.len(), "encoding")?;
244 let meta_count = to_u16_len(self.metadata.len(), "metadata_count")?;
245 let data_len = to_u32_len(self.bytes.len(), "bytes")?;
246
247 let mut out = Vec::new();
248 out.extend_from_slice(&Self::ENVELOPE_MAGIC);
249 out.push(Self::ENVELOPE_VERSION);
250 out.extend_from_slice(&ct_len.to_le_bytes());
251 out.extend_from_slice(&enc_len.to_le_bytes());
252 out.extend_from_slice(&meta_count.to_le_bytes());
253 out.extend_from_slice(&data_len.to_le_bytes());
254 out.extend_from_slice(ct_bytes);
255 out.extend_from_slice(enc_bytes);
256
257 let mut keys: Vec<&String> = self.metadata.keys().collect();
258 keys.sort();
259
260 for key in keys {
261 let value = self
262 .metadata
263 .get(key)
264 .ok_or_else(|| DagError::InvalidPayloadEnvelope("metadata key missing".into()))?;
265 let k = key.as_bytes();
266 let v = value.as_bytes();
267 let k_len = to_u16_len(k.len(), "metadata_key")?;
268 let v_len = to_u16_len(v.len(), "metadata_value")?;
269 out.extend_from_slice(&k_len.to_le_bytes());
270 out.extend_from_slice(&v_len.to_le_bytes());
271 out.extend_from_slice(k);
272 out.extend_from_slice(v);
273 }
274
275 out.extend_from_slice(&self.bytes);
276 Ok(out)
277 }
278
279 pub fn from_envelope_bytes(input: &[u8]) -> Result<Self, DagError> {
280 struct Cursor<'a> {
281 buf: &'a [u8],
282 pos: usize,
283 }
284 impl<'a> Cursor<'a> {
285 fn new(buf: &'a [u8]) -> Self {
286 Self { buf, pos: 0 }
287 }
288 fn read_exact(&mut self, n: usize) -> Result<&'a [u8], DagError> {
289 if self.pos + n > self.buf.len() {
290 return Err(DagError::InvalidPayloadEnvelope("unexpected eof".into()));
291 }
292 let s = &self.buf[self.pos..self.pos + n];
293 self.pos += n;
294 Ok(s)
295 }
296 fn read_u8(&mut self) -> Result<u8, DagError> {
297 Ok(self.read_exact(1)?[0])
298 }
299 fn read_u16(&mut self) -> Result<u16, DagError> {
300 let b = self.read_exact(2)?;
301 Ok(u16::from_le_bytes([b[0], b[1]]))
302 }
303 fn read_u32(&mut self) -> Result<u32, DagError> {
304 let b = self.read_exact(4)?;
305 Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
306 }
307 fn read_str(&mut self, len: usize) -> Result<String, DagError> {
308 let b = self.read_exact(len)?;
309 std::str::from_utf8(b)
310 .map(|s| s.to_string())
311 .map_err(|_| DagError::InvalidPayloadEnvelope("invalid utf8".into()))
312 }
313 }
314
315 let mut c = Cursor::new(input);
316 let magic = c.read_exact(4)?;
317 if magic != Self::ENVELOPE_MAGIC {
318 return Err(DagError::InvalidPayloadEnvelope("bad magic".into()));
319 }
320
321 let version = c.read_u8()?;
322 if version != Self::ENVELOPE_VERSION {
323 return Err(DagError::UnsupportedPayloadVersion(version));
324 }
325
326 let ct_len = c.read_u16()? as usize;
327 let enc_len = c.read_u16()? as usize;
328 let meta_count = c.read_u16()? as usize;
329 let data_len = c.read_u32()? as usize;
330
331 let content_type = {
332 let s = c.read_str(ct_len)?;
333 if s.is_empty() { None } else { Some(s) }
334 };
335 let encoding = {
336 let s = c.read_str(enc_len)?;
337 if s.is_empty() { None } else { Some(s) }
338 };
339
340 let mut metadata = HashMap::new();
341 for _ in 0..meta_count {
342 let k_len = c.read_u16()? as usize;
343 let v_len = c.read_u16()? as usize;
344 let key = c.read_str(k_len)?;
345 let value = c.read_str(v_len)?;
346 metadata.insert(key, value);
347 }
348
349 let bytes = c.read_exact(data_len)?.to_vec();
350 if c.pos != input.len() {
351 return Err(DagError::InvalidPayloadEnvelope("trailing bytes".into()));
352 }
353
354 Ok(Self {
355 bytes,
356 content_type,
357 encoding,
358 metadata,
359 })
360 }
361}
362
363#[derive(Debug, Clone, PartialEq, Eq)]
364pub enum DagNodeStatus {
365 Pending,
366 Ready,
367 Running,
368 Succeeded,
369 Failed,
370}
371
372#[derive(Debug, Clone)]
373pub struct NodeExecutionContext {
374 pub run_id: String,
375 pub run_fencing_token: Option<u64>,
376 pub node_id: String,
377 pub attempt: usize,
378 pub upstream_nodes: Vec<String>,
379 pub upstream_outputs: HashMap<String, TaskPayload>,
380 pub layer_index: usize,
381}
382
383#[async_trait]
384pub trait DagNodeTrait: Send + Sync {
385 async fn start(&self, context: NodeExecutionContext) -> Result<TaskPayload, DagError>;
386}
387
388#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
389pub enum DagErrorCode {
390 DuplicateNode,
391 NodeNotFound,
392 PrecedingNodeNotFound,
393 CycleDetected,
394 EmptyGraph,
395 MissingNodeCompute,
396 NodeExecutionFailed,
397 RetryExhausted,
398 TaskJoinFailed,
399 RunAlreadyInProgress,
400 RunGuardAcquireFailed,
401 RunGuardRenewFailed,
402 RunGuardReleaseFailed,
403 MissingRunFencingToken,
404 FencingTokenRejected,
405 InvalidPayloadEnvelope,
406 UnsupportedPayloadVersion,
407 ReservedControlNode,
408 InvalidPrecedingControlNode,
409 RemoteDispatchNotConfigured,
410 NodeTimeout,
411 ExecutionTimeout,
412 ExecutionIncomplete,
413 InvalidStateTransition,
414}
415
416#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
417pub enum DagErrorClass {
418 Retryable,
419 NonRetryable,
420}
421
422#[derive(Debug, Clone, Serialize, Deserialize)]
423pub struct DagNodeSyncState {
424 pub run_id: String,
425 pub run_fencing_token: Option<u64>,
426 pub node_id: String,
427 pub stage: String,
428 pub placement: String,
429 pub worker_group: Option<String>,
430 pub worker_id: Option<String>,
431 pub attempt: usize,
432 pub layer_index: usize,
433 pub idempotency_key: Option<String>,
434 pub deadline_ms: Option<u64>,
435 pub dispatch_latency_ms: Option<u64>,
436 pub error_code: Option<DagErrorCode>,
437 pub error_class: Option<DagErrorClass>,
438 pub error: Option<String>,
439 pub timestamp_ms: u64,
440}
441
442#[derive(Debug, Clone, Default)]
443pub struct DagRunResumeState {
444 pub run_id: String,
445 pub run_fencing_token: Option<u64>,
446 pub succeeded_outputs: HashMap<String, TaskPayload>,
447}
448
449#[derive(Debug, Clone, PartialEq, Eq)]
450pub enum NodePlacement {
451 Local,
452 Remote { worker_group: String },
453}
454
455#[derive(Debug, Clone, Copy, PartialEq, Eq)]
456pub enum DagNodeRetryMode {
457 AllErrors,
458 RetryableOnly,
459}
460
461#[derive(Debug, Clone)]
462pub struct DagNodeExecutionPolicy {
463 pub max_retries: usize,
464 pub timeout_ms: Option<u64>,
465 pub retry_backoff_ms: u64,
466 pub idempotency_key: Option<String>,
467 pub retry_mode: DagNodeRetryMode,
468 pub circuit_breaker_failure_threshold: Option<usize>,
469 pub circuit_breaker_open_ms: u64,
470}
471
472impl Default for DagNodeExecutionPolicy {
473 fn default() -> Self {
474 Self {
475 max_retries: 0,
476 timeout_ms: None,
477 retry_backoff_ms: 0,
478 idempotency_key: None,
479 retry_mode: DagNodeRetryMode::AllErrors,
480 circuit_breaker_failure_threshold: None,
481 circuit_breaker_open_ms: 0,
482 }
483 }
484}
485
486#[async_trait]
487pub trait DagNodeDispatcher: Send + Sync {
488 async fn dispatch(
489 &self,
490 node_id: &str,
491 placement: &NodePlacement,
492 executor: Arc<dyn DagNodeTrait>,
493 context: NodeExecutionContext,
494 ) -> Result<TaskPayload, DagError>;
495}
496
497#[async_trait]
498pub trait DagFencingStore: Send + Sync {
499 async fn commit(
500 &self,
501 resource: &str,
502 token: u64,
503 node_id: &str,
504 payload: &TaskPayload,
505 ) -> Result<(), DagError>;
506}
507
508#[async_trait]
509pub trait DagRunStateStore: Send + Sync {
510 async fn load(&self, run_key: &str) -> Result<Option<DagRunResumeState>, DagError>;
511 async fn save(&self, run_key: &str, state: &DagRunResumeState) -> Result<(), DagError>;
512 async fn clear(&self, run_key: &str) -> Result<(), DagError>;
513}
514
515#[derive(Debug, Clone)]
516pub struct DagRunGuardAcquireOutcome {
517 pub acquired: bool,
518 pub fencing_token: Option<u64>,
519}
520
521#[async_trait]
522pub trait DagRunGuard: Send + Sync {
523 async fn try_acquire(
524 &self,
525 lock_key: &str,
526 owner: &str,
527 ttl_ms: u64,
528 ) -> Result<DagRunGuardAcquireOutcome, DagError>;
529 async fn renew(&self, lock_key: &str, owner: &str, ttl_ms: u64) -> Result<bool, DagError>;
530 async fn release(&self, lock_key: &str, owner: &str) -> Result<(), DagError>;
531}
532
533#[derive(Default)]
534pub struct LocalNodeDispatcher;
535
536#[async_trait]
537impl DagNodeDispatcher for LocalNodeDispatcher {
538 async fn dispatch(
539 &self,
540 node_id: &str,
541 placement: &NodePlacement,
542 executor: Arc<dyn DagNodeTrait>,
543 context: NodeExecutionContext,
544 ) -> Result<TaskPayload, DagError> {
545 match placement {
546 NodePlacement::Local => executor.start(context).await,
547 NodePlacement::Remote { .. } => {
548 Err(DagError::RemoteDispatchNotConfigured(node_id.to_string()))
549 }
550 }
551 }
552}
553
554#[derive(Clone)]
555pub struct DagNodeRecord {
556 pub id: String,
557 pub predecessors: Vec<String>,
558 pub status: DagNodeStatus,
559 pub placement: NodePlacement,
560 pub execution_policy: DagNodeExecutionPolicy,
561 pub executor: Option<Arc<dyn DagNodeTrait>>,
562 pub result: Option<TaskPayload>,
563 pub metadata: HashMap<String, String>,
564 pub error: Option<String>,
565}