krishiv_plan/
task_fragment.rs

1//! Typed task fragments with explicit execution kind (unified batch/streaming).
2
3use crate::{ExecutionKind, NodeOp, PlanError, PlanNode};
4
5const FRAGMENT_PREFIX: &str = "krishiv-fragment:";
6
7/// Current typed task-fragment envelope version.
8pub const TASK_FRAGMENT_VERSION: u32 = 1;
9
10const fn default_task_fragment_version() -> u32 {
11    TASK_FRAGMENT_VERSION
12}
13
14/// Wire- and log-stable task fragment envelope.
15#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
16pub struct TypedTaskFragment {
17    /// Envelope format version. Version changes are independent from operator/state versions.
18    #[serde(default = "default_task_fragment_version")]
19    pub version: u32,
20    pub execution_kind: ExecutionKind,
21    pub body: String,
22}
23
24impl TypedTaskFragment {
25    pub fn new(execution_kind: ExecutionKind, body: impl Into<String>) -> Self {
26        Self {
27            version: TASK_FRAGMENT_VERSION,
28            execution_kind,
29            body: body.into(),
30        }
31    }
32
33    pub fn encode(&self) -> Result<String, PlanError> {
34        let json = serde_json::to_string(self)
35            .map_err(|e| PlanError::Encode(format!("fragment json: {e}")))?;
36        Ok(format!("{FRAGMENT_PREFIX}{json}"))
37    }
38
39    pub fn decode(fragment: &str) -> Option<Self> {
40        Self::decode_versioned(fragment).ok()
41    }
42
43    /// Decode and validate the fragment envelope version.
44    pub fn decode_versioned(fragment: &str) -> Result<Self, PlanError> {
45        let payload = fragment.strip_prefix(FRAGMENT_PREFIX).ok_or_else(|| {
46            PlanError::Parse("task fragment does not use the typed envelope".into())
47        })?;
48        let decoded: Self = serde_json::from_str(payload)
49            .map_err(|e| PlanError::Parse(format!("fragment json: {e}")))?;
50        if decoded.version != TASK_FRAGMENT_VERSION {
51            return Err(PlanError::Validation(format!(
52                "unsupported task fragment version {}; supported version is {}",
53                decoded.version, TASK_FRAGMENT_VERSION
54            )));
55        }
56        Ok(decoded)
57    }
58
59    pub fn execution_kind_from_legacy(fragment: &str) -> ExecutionKind {
60        if fragment.starts_with("stream:") {
61            return ExecutionKind::Streaming;
62        }
63        // IVM fragments shipped by the coordinator's distributed dispatch:
64        // the legacy stateless `delta:step:` tick plus the Phase 57 resident
65        // protocol (`delta:attach:` / `delta:tick:` / `delta:ckpt:` /
66        // `delta:detach:`). Without this the executor infers Batch and the
67        // batch dispatcher rejects the fragment, forcing central fallback.
68        if fragment.starts_with("delta:") {
69            return ExecutionKind::DeltaBatch;
70        }
71        if let Some(op) = crate::lowering::decode_task_fragment(fragment) {
72            return match op {
73                NodeOp::Window { .. }
74                | NodeOp::StreamSource { .. }
75                | NodeOp::Watermark { .. }
76                | NodeOp::KeyBy { .. }
77                | NodeOp::StateTtl { .. } => ExecutionKind::Streaming,
78                _ => ExecutionKind::Batch,
79            };
80        }
81        ExecutionKind::Batch
82    }
83
84    pub fn decode_or_legacy(fragment: &str) -> Self {
85        Self::decode(fragment).unwrap_or_else(|| {
86            Self::new(
87                Self::execution_kind_from_legacy(fragment),
88                fragment.to_string(),
89            )
90        })
91    }
92
93    /// Decode a fragment, rejecting legacy untyped strings in durable/production profiles.
94    pub fn decode_for_profile(
95        fragment: &str,
96        profile: krishiv_common::DurabilityProfile,
97    ) -> Result<Self, PlanError> {
98        if fragment.starts_with(FRAGMENT_PREFIX) {
99            return Self::decode_versioned(fragment);
100        }
101        if krishiv_common::allow_legacy_task_fragments(profile) {
102            return Ok(Self::decode_or_legacy(fragment));
103        }
104        Err(PlanError::Validation(format!(
105            "legacy untyped task fragment rejected for profile '{}': {}",
106            profile,
107            fragment.chars().take(120).collect::<String>()
108        )))
109    }
110}
111
112/// Validate and return the task body for executor/scheduler routing.
113pub fn task_body_for_profile(
114    fragment: &str,
115    profile: krishiv_common::DurabilityProfile,
116) -> Result<String, PlanError> {
117    let body = TypedTaskFragment::decode_for_profile(fragment, profile)?.body;
118    if body.len() == body.trim().len() {
119        return Ok(body);
120    }
121    Ok(body.trim().to_owned())
122}
123
124/// Validate every task fragment in a job spec.
125pub fn validate_job_fragments(
126    spec: &krishiv_proto::JobSpec,
127    profile: krishiv_common::DurabilityProfile,
128) -> Result<(), PlanError> {
129    for stage in spec.stages() {
130        for task in stage.tasks() {
131            let typed = TypedTaskFragment::decode_for_profile(task.description(), profile)?;
132            // Validate window specs embedded in the fragment body.
133            if typed
134                .body
135                .starts_with(crate::window::WINDOW_EXECUTION_SPEC_PREFIX)
136                || typed.body.starts_with("stream:tw:")
137                || typed.body.starts_with("stream:sw:")
138                || typed.body.starts_with("stream:ses:")
139            {
140                crate::window::decode_window_execution_spec(&typed.body)?;
141            } else if let Some(NodeOp::Window { spec }) =
142                crate::lowering::decode_task_fragment(&typed.body)
143            {
144                crate::window::validate_window_execution_spec(&spec)?;
145            }
146        }
147    }
148    Ok(())
149}
150
151/// Encode a plan node as a typed task fragment string for the scheduler/executor.
152pub fn encode_typed_task_fragment(node: &PlanNode) -> Result<String, PlanError> {
153    let body = crate::lowering::encode_task_fragment(node);
154    TypedTaskFragment::new(node.kind(), body).encode()
155}
156
157/// Decode execution kind from any fragment representation.
158pub fn execution_kind_from_fragment(fragment: &str) -> ExecutionKind {
159    TypedTaskFragment::decode_or_legacy(fragment).execution_kind
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use crate::{NodeOp, PlanNode};
166
167    #[test]
168    fn round_trip_typed_fragment() {
169        use crate::window::WindowExecutionSpec;
170        let spec = WindowExecutionSpec::tumbling("user_id", "ts", 1_000);
171        let node = PlanNode::new("w", "win", ExecutionKind::Streaming).with_op(NodeOp::Window {
172            spec: Box::new(spec),
173        });
174        let encoded = encode_typed_task_fragment(&node).expect("encode");
175        let decoded = TypedTaskFragment::decode_or_legacy(&encoded);
176        assert_eq!(decoded.execution_kind, ExecutionKind::Streaming);
177        assert!(decoded.body.starts_with("stream:"));
178    }
179
180    #[test]
181    fn legacy_stream_prefix_is_streaming() {
182        assert_eq!(
183            execution_kind_from_fragment("stream:tw:key=u"),
184            ExecutionKind::Streaming
185        );
186    }
187
188    #[test]
189    fn legacy_delta_step_prefix_is_delta_batch() {
190        // Distributed IVM tick fragment must classify as DeltaBatch so the
191        // executor routes it to the IVM handler, not the batch dispatcher.
192        assert_eq!(
193            execution_kind_from_fragment("delta:step:orders|d|s|st"),
194            ExecutionKind::DeltaBatch
195        );
196    }
197
198    #[test]
199    fn rejects_unknown_fragment_version() {
200        let encoded = format!(
201            "{FRAGMENT_PREFIX}{}",
202            serde_json::json!({"version": 99, "execution_kind": "Batch", "body": "scan"})
203        );
204        let err = TypedTaskFragment::decode_versioned(&encoded).unwrap_err();
205        assert!(
206            err.to_string()
207                .contains("unsupported task fragment version")
208        );
209    }
210
211    #[test]
212    fn durable_profile_rejects_legacy_fragment() {
213        use krishiv_common::DurabilityProfile;
214        let err = TypedTaskFragment::decode_for_profile(
215            "stream:tw:key=u",
216            DurabilityProfile::SingleNodeDurable,
217        )
218        .unwrap_err();
219        assert!(err.to_string().contains("legacy untyped"));
220    }
221}
krishiv_plan/task_fragment.rs

krishiv_plan/
task_fragment.rs
