fix checkpoint

Author: luoluoyuyu

protocol/proto/storage.proto

@@ -58,9 +58,12 @@ message KafkaSourceSubtaskCheckpoint {
   repeated KafkaPartitionOffset partitions = 4;
 }
 
-// Generic source checkpoint payload envelope (enum-like via oneof).
-message SourceCheckpointPayload {
-  oneof checkpoint {
+// Aggregated, epoch-aligned per-source checkpoint info persisted by the coordinator
+// (one entry per source subtask). The envelope decouples the catalog/meta-store from
+// any specific source implementation so new source types can be added by registering a
+// new `CheckpointAggregator` in the runtime without touching storage schema layouts.
+message SourceCheckpointInfo {
+  oneof info {
     KafkaSourceSubtaskCheckpoint kafka = 1;
   }
 }
@@ -81,10 +84,10 @@ message StreamingTableDefinition {
   // Updated by JobManager after all operators ACK. Used for crash recovery.
   uint64 latest_checkpoint_epoch = 6;
 
-  // Kafka source per-subtask offsets at the same committed epoch as `latest_checkpoint_epoch`.
-  // Populated by the runtime coordinator from source checkpoint ACKs. Optional `.bin` files under
-  // the job state dir may exist only for local recovery materialization from this field.
-  repeated KafkaSourceSubtaskCheckpoint kafka_source_checkpoints = 7;
+  // Source-agnostic per-subtask checkpoint entries aligned with `latest_checkpoint_epoch`.
+  // Populated by the runtime CheckpointCoordinator via the configured aggregator registry;
+  // storage no longer knows the internal layout of any particular source type.
+  repeated SourceCheckpointInfo source_checkpoints = 7;
 }
 
 // =============================================================================

src/coordinator/execution/executor.rs

@@ -368,6 +368,7 @@ impl PlanVisitor for Executor {
                     fs_program,
                     custom_interval,
                     None,
+                    vec![],
                 ))
             })
             .map_err(|e| ExecuteError::Internal(format!("Failed to submit streaming job: {e}")))?;

src/runtime/streaming/api/context.rs

@@ -16,7 +16,7 @@ use std::time::{Duration, SystemTime};
 
 use anyhow::{Context, Result, anyhow};
 use arrow_array::RecordBatch;
-use protocol::storage::SourceCheckpointPayload;
+use protocol::storage::SourceCheckpointInfo;
 use tokio::sync::mpsc;
 
 use crate::runtime::memory::{MemoryBlock, MemoryPool, get_array_memory_size};
@@ -78,7 +78,7 @@ pub struct TaskContext {
     /// Last globally-committed safe epoch for crash recovery.
     safe_epoch: u64,
 
-    /// When set, pipelines report checkpoint completion (and optional Kafka offsets) to the job coordinator.
+    /// When set, pipelines report checkpoint completion to the job coordinator.
     checkpoint_ack_tx: Option<mpsc::Sender<JobMasterEvent>>,
 }
 
@@ -128,17 +128,13 @@ impl TaskContext {
     }
 
     /// Notify the job checkpoint coordinator that this pipeline has finished the barrier for `epoch`.
-    pub async fn send_checkpoint_ack(
-        &self,
-        epoch: u64,
-        source_payloads: Vec<SourceCheckpointPayload>,
-    ) {
+    pub async fn send_checkpoint_ack(&self, epoch: u64, source_infos: Vec<SourceCheckpointInfo>) {
         if let Some(tx) = &self.checkpoint_ack_tx {
             let _ = tx
                 .send(JobMasterEvent::CheckpointAck {
                     pipeline_id: self.pipeline_id,
                     epoch,
-                    source_payloads,
+                    source_infos,
                 })
                 .await;
         }

src/runtime/streaming/api/operator.rs

@@ -66,7 +66,7 @@ pub trait Operator: Send + 'static {
     /// `commit_transaction` on the producer stashed during [`Self::snapshot_state`].
     async fn commit_checkpoint(
         &mut self,
-        epoch: u32,
+        epoch: u64,
         _ctx: &mut TaskContext,
     ) -> anyhow::Result<()> {
         let _ = epoch;
@@ -76,7 +76,7 @@ pub trait Operator: Send + 'static {
     /// Global checkpoint **rollback** when phase 2 must not commit (e.g. catalog persist failed).
     ///
     /// Default is no-op. Transactional Kafka sink overrides with `abort_transaction` on the stashed producer.
-    async fn abort_checkpoint(&mut self, epoch: u32, _ctx: &mut TaskContext) -> anyhow::Result<()> {
+    async fn abort_checkpoint(&mut self, epoch: u64, _ctx: &mut TaskContext) -> anyhow::Result<()> {
         let _ = epoch;
         Ok(())
     }

src/runtime/streaming/api/source.rs

@@ -15,7 +15,7 @@ use crate::sql::common::{CheckpointBarrier, Watermark};
 use arrow_array::RecordBatch;
 use async_trait::async_trait;
 use protocol::storage::{
-    KafkaSourceSubtaskCheckpoint, SourceCheckpointPayload, source_checkpoint_payload,
+    KafkaSourceSubtaskCheckpoint, SourceCheckpointInfo, source_checkpoint_info,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -34,17 +34,19 @@ pub enum SourceEvent {
     EndOfStream,
 }
 
-/// Optional metadata returned when a source completes a checkpoint barrier snapshot.
+/// Checkpoint metadata produced by a source subtask during a barrier snapshot.
+/// Sources fill this directly with [`SourceCheckpointInfo`] — the coordinator collects
+/// and persists these entries without any further translation step.
 #[derive(Debug, Default, Clone)]
 pub struct SourceCheckpointReport {
-    pub payloads: Vec<SourceCheckpointPayload>,
+    pub infos: Vec<SourceCheckpointInfo>,
 }
 
 impl SourceCheckpointReport {
     pub fn from_kafka_checkpoint(kafka: KafkaSourceSubtaskCheckpoint) -> Self {
         Self {
-            payloads: vec![SourceCheckpointPayload {
-                checkpoint: Some(source_checkpoint_payload::Checkpoint::Kafka(kafka)),
+            infos: vec![SourceCheckpointInfo {
+                info: Some(source_checkpoint_info::Info::Kafka(kafka)),
             }],
         }
     }
@@ -54,6 +56,11 @@ impl SourceCheckpointReport {
 pub trait SourceOperator: Send + 'static {
     fn name(&self) -> &str;
 
+    /// Inject persisted checkpoint records before the source is started.
+    /// Called by the engine after the operator is constructed and before [`Self::on_start`].
+    /// Default implementation is a no-op; sources with stateful recovery override this.
+    fn set_recovery_checkpoint(&mut self, _infos: Vec<SourceCheckpointInfo>) {}
+
     async fn on_start(&mut self, _ctx: &mut TaskContext) -> anyhow::Result<()> {
         Ok(())
     }
@@ -74,15 +81,15 @@ pub trait SourceOperator: Send + 'static {
     /// Kafka source keeps the default: offsets are reported at the barrier in [`Self::snapshot_state`].
     async fn commit_checkpoint(
         &mut self,
-        epoch: u32,
+        epoch: u64,
         _ctx: &mut TaskContext,
     ) -> anyhow::Result<()> {
         let _ = epoch;
         Ok(())
     }
 
     /// Same rollback hook as [`super::operator::Operator::abort_checkpoint`].
-    async fn abort_checkpoint(&mut self, epoch: u32, _ctx: &mut TaskContext) -> anyhow::Result<()> {
+    async fn abort_checkpoint(&mut self, epoch: u64, _ctx: &mut TaskContext) -> anyhow::Result<()> {
         let _ = epoch;
         Ok(())
     }

src/runtime/streaming/execution/pipeline.rs

@@ -110,7 +110,7 @@ impl Pipeline {
                                         }
                                     }
                                     AlignmentStatus::Complete => {
-                                        let epoch = barrier.epoch as u64;
+                                        let epoch = barrier.epoch;
                                         self.chain_head
                                             .process_event(
                                                 idx,

src/runtime/streaming/execution/source_driver.rs

@@ -161,7 +161,7 @@ impl SourceDriver {
                 let b: CheckpointBarrier = barrier.clone().into();
                 let report = self.operator.snapshot_state(b, &mut self.ctx).await?;
                 self.dispatch_event(StreamEvent::Barrier(b)).await?;
-                pending_source_checkpoint = Some((b.epoch as u64, report));
+                pending_source_checkpoint = Some((b.epoch, report));
             }
             ControlCommand::Commit { epoch } => {
                 self.operator
@@ -186,7 +186,7 @@ impl SourceDriver {
         }
 
         if let Some((epoch, report)) = pending_source_checkpoint {
-            self.ctx.send_checkpoint_ack(epoch, report.payloads).await;
+            self.ctx.send_checkpoint_ack(epoch, report.infos).await;
         }
 
         Ok(stop)

src/runtime/streaming/execution/tracker/barrier_aligner.rs

@@ -23,7 +23,7 @@ pub enum AlignmentStatus {
 #[derive(Debug)]
 pub struct BarrierAligner {
     input_count: usize,
-    current_epoch: Option<u32>,
+    current_epoch: Option<u64>,
     reached_inputs: HashSet<usize>,
 }