oxide_batch_repository/
partition.rs

1//! Durable local-partition plan and result values.
2
3use std::error::Error;
4use std::fmt;
5
6use oxide_batch_core::{
7    BatchStatus, ExecutionContext, ExecutionCounts, ExecutionVersion, ExitStatus, MAX_PARTITIONS,
8    StepExecution, StepExecutionId, StepPartitionId,
9};
10
11/// Maximum UTF-8 byte length of one durable partition key.
12pub const MAX_PARTITION_KEY_BYTES: usize = 128;
13/// Maximum serialized byte length of one durable partition context.
14pub const MAX_PARTITION_CONTEXT_BYTES: usize = 4 * 1024;
15
16/// A stable byte-compared key within one partitioned step execution.
17#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
18pub struct PartitionKey(String);
19
20impl PartitionKey {
21    /// Validates a nonempty bounded UTF-8 partition key.
22    ///
23    /// Whitespace and punctuation are retained because partition identity is
24    /// byte-exact and application-defined.
25    ///
26    /// # Errors
27    ///
28    /// Returns [`PartitionValueError`] for an empty or oversized key.
29    pub fn new(value: impl Into<String>) -> Result<Self, PartitionValueError> {
30        let value = value.into();
31        if value.is_empty() {
32            return Err(PartitionValueError::EmptyKey);
33        }
34        if value.len() > MAX_PARTITION_KEY_BYTES {
35            return Err(PartitionValueError::KeyTooLong {
36                max_bytes: MAX_PARTITION_KEY_BYTES,
37            });
38        }
39        Ok(Self(value))
40    }
41
42    /// Borrows the byte-compared key.
43    #[must_use]
44    pub fn as_str(&self) -> &str {
45        &self.0
46    }
47}
48
49impl fmt::Debug for PartitionKey {
50    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
51        formatter.write_str("PartitionKey(<redacted>)")
52    }
53}
54
55impl fmt::Display for PartitionKey {
56    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
57        formatter.write_str(self.as_str())
58    }
59}
60
61/// One validated entry in a partition plan before durable identity assignment.
62#[derive(Clone, Eq, PartialEq)]
63pub struct PartitionPlanEntry {
64    key: PartitionKey,
65    context: ExecutionContext,
66}
67
68impl PartitionPlanEntry {
69    /// Constructs one bounded partition-plan entry.
70    ///
71    /// # Errors
72    ///
73    /// Returns [`PartitionValueError::ContextTooLarge`] when the serialized
74    /// context exceeds the schema-3 `4 KiB` ceiling.
75    pub fn new(key: PartitionKey, context: ExecutionContext) -> Result<Self, PartitionValueError> {
76        if context.encoded_len() > MAX_PARTITION_CONTEXT_BYTES {
77            return Err(PartitionValueError::ContextTooLarge {
78                max_bytes: MAX_PARTITION_CONTEXT_BYTES,
79            });
80        }
81        Ok(Self { key, context })
82    }
83
84    /// Borrows the stable partition key.
85    #[must_use]
86    pub const fn key(&self) -> &PartitionKey {
87        &self.key
88    }
89
90    /// Borrows the redacted durable context.
91    #[must_use]
92    pub const fn context(&self) -> &ExecutionContext {
93        &self.context
94    }
95}
96
97impl fmt::Debug for PartitionPlanEntry {
98    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
99        formatter
100            .debug_struct("PartitionPlanEntry")
101            .field("key", &self.key)
102            .field("context", &self.context)
103            .finish()
104    }
105}
106
107/// A durable partition plan row and its latest result snapshot.
108#[derive(Clone, Eq, PartialEq)]
109pub struct StepPartition {
110    id: StepPartitionId,
111    step_execution_id: StepExecutionId,
112    worker_step_execution_id: Option<StepExecutionId>,
113    key: PartitionKey,
114    ordinal: u32,
115    status: BatchStatus,
116    exit_status: ExitStatus,
117    counts: ExecutionCounts,
118    context: ExecutionContext,
119    version: ExecutionVersion,
120}
121
122impl StepPartition {
123    /// Reconstructs one durable partition row read by an adapter.
124    #[allow(clippy::too_many_arguments)]
125    #[doc(hidden)]
126    #[must_use]
127    pub const fn from_snapshot(
128        id: StepPartitionId,
129        step_execution_id: StepExecutionId,
130        worker_step_execution_id: Option<StepExecutionId>,
131        key: PartitionKey,
132        ordinal: u32,
133        status: BatchStatus,
134        exit_status: ExitStatus,
135        counts: ExecutionCounts,
136        context: ExecutionContext,
137        version: ExecutionVersion,
138    ) -> Self {
139        Self {
140            id,
141            step_execution_id,
142            worker_step_execution_id,
143            key,
144            ordinal,
145            status,
146            exit_status,
147            counts,
148            context,
149            version,
150        }
151    }
152
153    /// Builds the initial durable row for one planned partition.
154    #[doc(hidden)]
155    #[must_use]
156    pub fn starting(
157        id: StepPartitionId,
158        step_execution_id: StepExecutionId,
159        ordinal: u32,
160        entry: PartitionPlanEntry,
161    ) -> Self {
162        Self::from_snapshot(
163            id,
164            step_execution_id,
165            None,
166            entry.key,
167            ordinal,
168            BatchStatus::Starting,
169            ExitStatus::unknown(),
170            ExecutionCounts::default(),
171            entry.context,
172            ExecutionVersion::INITIAL,
173        )
174    }
175
176    /// Returns the durable partition-row identifier.
177    #[must_use]
178    pub const fn id(&self) -> StepPartitionId {
179        self.id
180    }
181
182    /// Returns the parent partitioned step execution.
183    #[must_use]
184    pub const fn step_execution_id(&self) -> StepExecutionId {
185        self.step_execution_id
186    }
187
188    /// Returns the assigned worker attempt, when one has started.
189    #[must_use]
190    pub const fn worker_step_execution_id(&self) -> Option<StepExecutionId> {
191        self.worker_step_execution_id
192    }
193
194    /// Borrows the stable partition key.
195    #[must_use]
196    pub const fn key(&self) -> &PartitionKey {
197        &self.key
198    }
199
200    /// Returns the one-based partition-plan ordinal.
201    #[must_use]
202    pub const fn ordinal(&self) -> u32 {
203        self.ordinal
204    }
205
206    /// Returns the current framework status.
207    #[must_use]
208    pub const fn status(&self) -> BatchStatus {
209        self.status
210    }
211
212    /// Borrows the latest stable exit status.
213    #[must_use]
214    pub const fn exit_status(&self) -> &ExitStatus {
215        &self.exit_status
216    }
217
218    /// Returns the latest durable counters.
219    #[must_use]
220    pub const fn counts(&self) -> ExecutionCounts {
221        self.counts
222    }
223
224    /// Borrows the redacted partition context.
225    #[must_use]
226    pub const fn context(&self) -> &ExecutionContext {
227        &self.context
228    }
229
230    /// Returns the optimistic-lock version.
231    #[must_use]
232    pub const fn version(&self) -> ExecutionVersion {
233        self.version
234    }
235
236    /// Assigns one worker attempt to this partition.
237    ///
238    /// # Errors
239    ///
240    /// Returns [`PartitionMutationError`] for a stale version or a status that
241    /// cannot take an assignment.
242    #[doc(hidden)]
243    pub fn assign(
244        &mut self,
245        expected_version: ExecutionVersion,
246        worker_step_execution_id: StepExecutionId,
247    ) -> Result<(), PartitionMutationError> {
248        self.ensure_version(expected_version)?;
249        let first_assignment =
250            self.status == BatchStatus::Starting && self.worker_step_execution_id.is_none();
251        let retry_assignment = matches!(self.status, BatchStatus::Failed | BatchStatus::Stopped)
252            && self.worker_step_execution_id.is_some();
253        if !first_assignment && !retry_assignment {
254            return Err(PartitionMutationError::InvalidState {
255                status: self.status,
256            });
257        }
258        self.worker_step_execution_id = Some(worker_step_execution_id);
259        self.status = BatchStatus::Started;
260        self.exit_status = ExitStatus::unknown();
261        self.counts = ExecutionCounts::default();
262        self.version = self.next_version()?;
263        Ok(())
264    }
265
266    /// Records one terminal worker result on this partition.
267    ///
268    /// # Errors
269    ///
270    /// Returns [`PartitionMutationError`] for a stale version or a partition
271    /// that is not assigned and started.
272    #[doc(hidden)]
273    pub fn complete(
274        &mut self,
275        expected_version: ExecutionVersion,
276        result: &PartitionResult,
277    ) -> Result<(), PartitionMutationError> {
278        self.ensure_version(expected_version)?;
279        if self.status != BatchStatus::Started || self.worker_step_execution_id.is_none() {
280            return Err(PartitionMutationError::InvalidState {
281                status: self.status,
282            });
283        }
284        self.status = result.status;
285        self.exit_status = result.exit_status.clone();
286        self.counts = result.counts;
287        self.version = self.next_version()?;
288        Ok(())
289    }
290
291    fn ensure_version(
292        &self,
293        expected_version: ExecutionVersion,
294    ) -> Result<(), PartitionMutationError> {
295        if expected_version != self.version {
296            return Err(PartitionMutationError::StaleVersion {
297                expected: expected_version,
298                actual: self.version,
299            });
300        }
301        Ok(())
302    }
303
304    fn next_version(&self) -> Result<ExecutionVersion, PartitionMutationError> {
305        self.version
306            .get()
307            .checked_add(1)
308            .map(ExecutionVersion::new)
309            .ok_or(PartitionMutationError::VersionExhausted)
310    }
311}
312
313impl fmt::Debug for StepPartition {
314    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
315        formatter
316            .debug_struct("StepPartition")
317            .field("id", &self.id)
318            .field("step_execution_id", &self.step_execution_id)
319            .field("worker_step_execution_id", &self.worker_step_execution_id)
320            .field("key", &self.key)
321            .field("ordinal", &self.ordinal)
322            .field("status", &self.status)
323            .field("exit_status", &self.exit_status)
324            .field("counts", &self.counts)
325            .field("context", &self.context)
326            .field("version", &self.version)
327            .finish()
328    }
329}
330
331/// A validated terminal result published by one assigned partition worker.
332#[derive(Clone, Debug, Eq, PartialEq)]
333pub struct PartitionResult {
334    status: BatchStatus,
335    exit_status: ExitStatus,
336    counts: ExecutionCounts,
337}
338
339/// The deterministic result of aggregating one complete durable partition plan.
340#[derive(Clone, Debug, Eq, PartialEq)]
341pub struct PartitionAggregate {
342    status: BatchStatus,
343    exit_status: ExitStatus,
344    counts: ExecutionCounts,
345    selected_worker_step_execution_id: StepExecutionId,
346}
347
348impl PartitionAggregate {
349    /// Returns the most severe child status.
350    #[must_use]
351    pub const fn status(&self) -> BatchStatus {
352        self.status
353    }
354
355    /// Borrows the first matching exit status in partition-key byte order.
356    #[must_use]
357    pub const fn exit_status(&self) -> &ExitStatus {
358        &self.exit_status
359    }
360
361    /// Returns the checked sum of every durable child counter.
362    #[must_use]
363    pub const fn counts(&self) -> ExecutionCounts {
364        self.counts
365    }
366
367    /// Returns the worker attempt whose result the aggregate selected.
368    #[doc(hidden)]
369    #[must_use]
370    pub const fn selected_worker_step_execution_id(&self) -> StepExecutionId {
371        self.selected_worker_step_execution_id
372    }
373}
374
375/// Aggregates a complete partition plan independently of input or completion order.
376///
377/// Children are ordered by their byte-exact partition keys before counters and
378/// exit status are selected. The status severity is fixed as
379/// `UNKNOWN > FAILED > STOPPED > COMPLETED`.
380///
381/// # Errors
382///
383/// Returns [`PartitionAggregationError`] when the plan is empty, contains a
384/// duplicate key, still has an active/non-runtime result, or a counter sum
385/// exceeds the durable representation.
386pub fn aggregate_step_partitions(
387    partitions: &[StepPartition],
388) -> Result<PartitionAggregate, PartitionAggregationError> {
389    if partitions.is_empty() {
390        return Err(PartitionAggregationError::EmptyPlan);
391    }
392    if partitions.len() > usize::from(MAX_PARTITIONS) {
393        return Err(PartitionAggregationError::PlanTooLarge {
394            max: usize::from(MAX_PARTITIONS),
395        });
396    }
397
398    let mut ordered = partitions.iter().collect::<Vec<_>>();
399    ordered.sort_by(|left, right| left.key().cmp(right.key()));
400    if ordered
401        .windows(2)
402        .any(|pair| pair[0].key() == pair[1].key())
403    {
404        return Err(PartitionAggregationError::DuplicateKey);
405    }
406
407    let mut aggregate_status = BatchStatus::Completed;
408    let mut counts = ExecutionCounts::default();
409    for partition in &ordered {
410        let status = partition.status();
411        if partition.worker_step_execution_id().is_none()
412            || !matches!(
413                status,
414                BatchStatus::Completed
415                    | BatchStatus::Failed
416                    | BatchStatus::Stopped
417                    | BatchStatus::Unknown
418            )
419        {
420            return Err(PartitionAggregationError::Incomplete { status });
421        }
422        if partition_severity(status) > partition_severity(aggregate_status) {
423            aggregate_status = status;
424        }
425        counts = checked_sum_counts(counts, partition.counts())?;
426    }
427
428    let selected = ordered
429        .iter()
430        .find(|partition| partition.status() == aggregate_status)
431        .ok_or(PartitionAggregationError::Incomplete {
432            status: aggregate_status,
433        })?;
434    let selected_worker_step_execution_id =
435        selected
436            .worker_step_execution_id()
437            .ok_or(PartitionAggregationError::Incomplete {
438                status: aggregate_status,
439            })?;
440    Ok(PartitionAggregate {
441        status: aggregate_status,
442        exit_status: selected.exit_status().clone(),
443        counts,
444        selected_worker_step_execution_id,
445    })
446}
447
448const fn partition_severity(status: BatchStatus) -> u8 {
449    match status {
450        BatchStatus::Completed => 0,
451        BatchStatus::Stopped => 1,
452        BatchStatus::Failed => 2,
453        BatchStatus::Unknown => 3,
454        _ => 4,
455    }
456}
457
458fn checked_sum_counts(
459    left: ExecutionCounts,
460    right: ExecutionCounts,
461) -> Result<ExecutionCounts, PartitionAggregationError> {
462    let counts = ExecutionCounts::new(
463        left.read()
464            .checked_add(right.read())
465            .ok_or(PartitionAggregationError::CountExhausted)?,
466        left.processed()
467            .checked_add(right.processed())
468            .ok_or(PartitionAggregationError::CountExhausted)?,
469        left.written()
470            .checked_add(right.written())
471            .ok_or(PartitionAggregationError::CountExhausted)?,
472        left.filtered()
473            .checked_add(right.filtered())
474            .ok_or(PartitionAggregationError::CountExhausted)?,
475        left.committed()
476            .checked_add(right.committed())
477            .ok_or(PartitionAggregationError::CountExhausted)?,
478        left.rolled_back()
479            .checked_add(right.rolled_back())
480            .ok_or(PartitionAggregationError::CountExhausted)?,
481    );
482    if [
483        counts.read(),
484        counts.processed(),
485        counts.written(),
486        counts.filtered(),
487        counts.committed(),
488        counts.rolled_back(),
489    ]
490    .into_iter()
491    .any(|value| value > i64::MAX as u64)
492    {
493        return Err(PartitionAggregationError::CountExhausted);
494    }
495    Ok(counts)
496}
497
498/// A deterministic partition plan could not be aggregated safely.
499#[derive(Clone, Copy, Debug, Eq, PartialEq)]
500#[non_exhaustive]
501pub enum PartitionAggregationError {
502    /// No durable child result was supplied.
503    EmptyPlan,
504    /// The supplied plan exceeded the accepted M4 partition bound.
505    PlanTooLarge {
506        /// Maximum accepted partition count.
507        max: usize,
508    },
509    /// More than one child used the same byte-exact key.
510    DuplicateKey,
511    /// At least one child did not have a durable runtime-terminal result.
512    Incomplete {
513        /// The unusable durable status.
514        status: BatchStatus,
515    },
516    /// At least one aggregate counter exceeded `u64`.
517    CountExhausted,
518}
519
520impl fmt::Display for PartitionAggregationError {
521    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
522        match self {
523            Self::EmptyPlan => formatter.write_str("partition aggregation requires a plan"),
524            Self::PlanTooLarge { max } => {
525                write!(formatter, "partition aggregation exceeds {max} children")
526            }
527            Self::DuplicateKey => {
528                formatter.write_str("partition aggregation found a duplicate key")
529            }
530            Self::Incomplete { status } => write!(
531                formatter,
532                "partition aggregation cannot use a child in {status}"
533            ),
534            Self::CountExhausted => {
535                formatter.write_str("partition aggregate counters are exhausted")
536            }
537        }
538    }
539}
540
541impl Error for PartitionAggregationError {}
542
543impl PartitionResult {
544    /// Reads one terminal worker attempt as a durable partition result.
545    ///
546    /// # Errors
547    ///
548    /// Returns [`PartitionValueError::NonTerminalResult`] when the worker has
549    /// not reached a terminal status.
550    #[doc(hidden)]
551    pub fn from_worker(worker: &StepExecution) -> Result<Self, PartitionValueError> {
552        Self::new(
553            worker.metadata().status(),
554            worker.metadata().exit_status().clone(),
555            worker.metadata().counts(),
556        )
557    }
558    /// Validates one known or explicitly ambiguous terminal worker result.
559    ///
560    /// # Errors
561    ///
562    /// Returns [`PartitionValueError::NonTerminalResult`] for an active or
563    /// abandoned status.
564    pub fn new(
565        status: BatchStatus,
566        exit_status: ExitStatus,
567        counts: ExecutionCounts,
568    ) -> Result<Self, PartitionValueError> {
569        if !matches!(
570            status,
571            BatchStatus::Completed
572                | BatchStatus::Failed
573                | BatchStatus::Stopped
574                | BatchStatus::Unknown
575        ) {
576            return Err(PartitionValueError::NonTerminalResult { status });
577        }
578        if [
579            counts.read(),
580            counts.processed(),
581            counts.written(),
582            counts.filtered(),
583            counts.committed(),
584            counts.rolled_back(),
585        ]
586        .into_iter()
587        .any(|value| value > i64::MAX as u64)
588        {
589            return Err(PartitionValueError::CountTooLarge);
590        }
591        Ok(Self {
592            status,
593            exit_status,
594            counts,
595        })
596    }
597
598    /// Returns the terminal framework status.
599    #[must_use]
600    pub const fn status(&self) -> BatchStatus {
601        self.status
602    }
603
604    /// Borrows the terminal exit status.
605    #[must_use]
606    pub const fn exit_status(&self) -> &ExitStatus {
607        &self.exit_status
608    }
609
610    /// Returns the terminal counters.
611    #[must_use]
612    pub const fn counts(&self) -> ExecutionCounts {
613        self.counts
614    }
615}
616
617/// Invalid public partition input.
618#[derive(Clone, Debug, Eq, PartialEq)]
619#[non_exhaustive]
620pub enum PartitionValueError {
621    /// A partition key was empty.
622    EmptyKey,
623    /// A partition key exceeded the durable byte bound.
624    KeyTooLong {
625        /// Maximum accepted UTF-8 bytes.
626        max_bytes: usize,
627    },
628    /// A partition context exceeded the schema-3 bound.
629    ContextTooLarge {
630        /// Maximum accepted serialized bytes.
631        max_bytes: usize,
632    },
633    /// A worker result used a non-terminal status.
634    NonTerminalResult {
635        /// Rejected status.
636        status: BatchStatus,
637    },
638    /// A counter cannot be represented by every durable adapter.
639    CountTooLarge,
640}
641
642impl fmt::Display for PartitionValueError {
643    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
644        match self {
645            Self::EmptyKey => formatter.write_str("partition key must not be empty"),
646            Self::KeyTooLong { max_bytes } => {
647                write!(formatter, "partition key exceeds {max_bytes} UTF-8 bytes")
648            }
649            Self::ContextTooLarge { max_bytes } => {
650                write!(formatter, "partition context exceeds {max_bytes} bytes")
651            }
652            Self::NonTerminalResult { status } => {
653                write!(
654                    formatter,
655                    "partition result status {status} is not terminal"
656                )
657            }
658            Self::CountTooLarge => {
659                formatter.write_str("partition result counter exceeds the portable durable bound")
660            }
661        }
662    }
663}
664
665impl Error for PartitionValueError {}
666
667#[derive(Clone, Copy, Debug, Eq, PartialEq)]
668/// A rejected durable partition mutation.
669#[doc(hidden)]
670pub enum PartitionMutationError {
671    StaleVersion {
672        expected: ExecutionVersion,
673        actual: ExecutionVersion,
674    },
675    InvalidState {
676        status: BatchStatus,
677    },
678    VersionExhausted,
679}
680
681#[cfg(test)]
682mod tests {
683    use super::*;
684    use oxide_batch_core::StateLimits;
685
686    #[test]
687    fn partition_key_and_context_bounds_fail_before_persistence() -> Result<(), Box<dyn Error>> {
688        assert_eq!(PartitionKey::new(""), Err(PartitionValueError::EmptyKey));
689        assert_eq!(
690            PartitionKey::new("x".repeat(MAX_PARTITION_KEY_BYTES + 1)),
691            Err(PartitionValueError::KeyTooLong {
692                max_bytes: MAX_PARTITION_KEY_BYTES,
693            })
694        );
695
696        let oversized = format!(
697            "{{\"format\":\"oxide-batch.execution-context\",\"format_version\":1,\
698             \"schema\":\"partition.test\",\"schema_version\":1,\
699             \"payload\":{{\"value\":\"{}\"}}}}",
700            "x".repeat(MAX_PARTITION_CONTEXT_BYTES)
701        );
702        let context =
703            ExecutionContext::from_json(oversized.as_bytes(), StateLimits::new(8 * 1024, 16)?)?;
704        assert_eq!(
705            PartitionPlanEntry::new(PartitionKey::new("partition-1")?, context),
706            Err(PartitionValueError::ContextTooLarge {
707                max_bytes: MAX_PARTITION_CONTEXT_BYTES,
708            })
709        );
710        Ok(())
711    }
712
713    #[test]
714    fn partition_result_accepts_only_runtime_terminal_outcomes() {
715        assert_eq!(
716            PartitionResult::new(
717                BatchStatus::Started,
718                ExitStatus::unknown(),
719                ExecutionCounts::default(),
720            ),
721            Err(PartitionValueError::NonTerminalResult {
722                status: BatchStatus::Started,
723            })
724        );
725        assert!(
726            PartitionResult::new(
727                BatchStatus::Unknown,
728                ExitStatus::unknown(),
729                ExecutionCounts::default(),
730            )
731            .is_ok()
732        );
733        assert_eq!(
734            PartitionResult::new(
735                BatchStatus::Completed,
736                ExitStatus::completed(),
737                ExecutionCounts::new(i64::MAX as u64 + 1, 0, 0, 0, 0, 0),
738            ),
739            Err(PartitionValueError::CountTooLarge)
740        );
741    }
742
743    #[test]
744    fn aggregate_rejects_counts_above_the_postgres_bigint_bound() -> Result<(), Box<dyn Error>> {
745        let context = ExecutionContext::from_json(
746            br#"{"format":"oxide-batch.execution-context","format_version":1,"schema":"partition.aggregate","schema_version":1,"payload":{}}"#,
747            StateLimits::new(MAX_PARTITION_CONTEXT_BYTES, 16)?,
748        )?;
749        let completed = |id: u64, key: &str, count: u64| -> Result<StepPartition, Box<dyn Error>> {
750            let mut partition = StepPartition::starting(
751                StepPartitionId::new(id)?,
752                StepExecutionId::new(1)?,
753                u32::try_from(id)?,
754                PartitionPlanEntry::new(PartitionKey::new(key)?, context.clone())?,
755            );
756            partition
757                .assign(ExecutionVersion::INITIAL, StepExecutionId::new(id + 10)?)
758                .map_err(|_| std::io::Error::other("partition assignment failed"))?;
759            partition
760                .complete(
761                    partition.version(),
762                    &PartitionResult::new(
763                        BatchStatus::Completed,
764                        ExitStatus::completed(),
765                        ExecutionCounts::new(count, 0, 0, 0, 0, 0),
766                    )?,
767                )
768                .map_err(|_| std::io::Error::other("partition completion failed"))?;
769            Ok(partition)
770        };
771        let first = completed(1, "alpha", i64::MAX as u64)?;
772        let second = completed(2, "beta", 1)?;
773        assert_eq!(
774            aggregate_step_partitions(&[first, second]),
775            Err(PartitionAggregationError::CountExhausted)
776        );
777        Ok(())
778    }
779
780    #[test]
781    fn aggregation_is_deterministic_in_partition_key_order() -> Result<(), Box<dyn Error>> {
782        let context = ExecutionContext::from_json(
783            br#"{"format":"oxide-batch.execution-context","format_version":1,"schema":"partition.aggregate","schema_version":1,"payload":{}}"#,
784            StateLimits::new(MAX_PARTITION_CONTEXT_BYTES, 16)?,
785        )?;
786        let mut alpha = StepPartition::starting(
787            StepPartitionId::new(1)?,
788            StepExecutionId::new(1)?,
789            1,
790            PartitionPlanEntry::new(PartitionKey::new("alpha")?, context.clone())?,
791        );
792        alpha
793            .assign(ExecutionVersion::INITIAL, StepExecutionId::new(2)?)
794            .map_err(|_| std::io::Error::other("alpha assignment failed"))?;
795        alpha
796            .complete(
797                alpha.version(),
798                &PartitionResult::new(
799                    BatchStatus::Failed,
800                    ExitStatus::new(oxide_batch_core::ExitCode::new("ALPHA_FAILED")?),
801                    ExecutionCounts::new(1, 2, 3, 4, 5, 6),
802                )?,
803            )
804            .map_err(|_| std::io::Error::other("alpha completion failed"))?;
805        let mut zeta = StepPartition::starting(
806            StepPartitionId::new(2)?,
807            StepExecutionId::new(1)?,
808            2,
809            PartitionPlanEntry::new(PartitionKey::new("zeta")?, context)?,
810        );
811        zeta.assign(ExecutionVersion::INITIAL, StepExecutionId::new(3)?)
812            .map_err(|_| std::io::Error::other("zeta assignment failed"))?;
813        zeta.complete(
814            zeta.version(),
815            &PartitionResult::new(
816                BatchStatus::Failed,
817                ExitStatus::new(oxide_batch_core::ExitCode::new("ZETA_FAILED")?),
818                ExecutionCounts::new(10, 20, 30, 40, 50, 60),
819            )?,
820        )
821        .map_err(|_| std::io::Error::other("zeta completion failed"))?;
822
823        let forward = aggregate_step_partitions(&[alpha.clone(), zeta.clone()])?;
824        let reverse = aggregate_step_partitions(&[zeta.clone(), alpha.clone()])?;
825        assert_eq!(forward, reverse);
826        assert_eq!(forward.status(), BatchStatus::Failed);
827        assert_eq!(forward.exit_status().code().as_str(), "ALPHA_FAILED");
828        assert_eq!(
829            forward.counts(),
830            ExecutionCounts::new(11, 22, 33, 44, 55, 66)
831        );
832
833        let context = ExecutionContext::from_json(
834            br#"{"format":"oxide-batch.execution-context","format_version":1,"schema":"partition.aggregate","schema_version":1,"payload":{}}"#,
835            StateLimits::new(MAX_PARTITION_CONTEXT_BYTES, 16)?,
836        )?;
837        let mut unknown = StepPartition::starting(
838            StepPartitionId::new(3)?,
839            StepExecutionId::new(1)?,
840            3,
841            PartitionPlanEntry::new(PartitionKey::new("middle")?, context)?,
842        );
843        unknown
844            .assign(ExecutionVersion::INITIAL, StepExecutionId::new(4)?)
845            .map_err(|_| std::io::Error::other("unknown assignment failed"))?;
846        unknown
847            .complete(
848                unknown.version(),
849                &PartitionResult::new(
850                    BatchStatus::Unknown,
851                    ExitStatus::unknown(),
852                    ExecutionCounts::default(),
853                )?,
854            )
855            .map_err(|_| std::io::Error::other("unknown completion failed"))?;
856        let ambiguous = aggregate_step_partitions(&[alpha, zeta, unknown])?;
857        assert_eq!(ambiguous.status(), BatchStatus::Unknown);
858        assert_eq!(ambiguous.exit_status(), &ExitStatus::unknown());
859        Ok(())
860    }
861}
oxide_batch_repository/partition.rs

oxide_batch_repository/
partition.rs
