update

Author: luoluoyuyu

protocol/proto/storage.proto

@@ -43,6 +43,21 @@ message CatalogSourceTable {
 // Streaming table storage (CREATE STREAMING TABLE persistence)
 // =============================================================================
 
+// Partition offset for one Kafka partition at a completed checkpoint.
+message KafkaPartitionOffset {
+  int32 partition = 1;
+  int64 offset = 2;
+}
+
+// Kafka source subtask checkpoint: one file / one TaskContext (pipeline + subtask).
+message KafkaSourceSubtaskCheckpoint {
+  uint32 pipeline_id = 1;
+  uint32 subtask_index = 2;
+  // Epoch of the barrier when this snapshot was taken (aligns with latest_checkpoint_epoch on commit).
+  uint64 checkpoint_epoch = 3;
+  repeated KafkaPartitionOffset partitions = 4;
+}
+
 // Persisted record for one streaming table (CREATE STREAMING TABLE).
 // On restart, the engine re-submits each record to JobManager to resume the pipeline.
 message StreamingTableDefinition {
@@ -58,6 +73,11 @@ message StreamingTableDefinition {
   // Last globally-committed checkpoint epoch.
   // 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;
 }
 
 // =============================================================================

src/runtime/streaming/api/context.rs

@@ -16,9 +16,12 @@ use std::time::{Duration, SystemTime};
 
 use anyhow::{Context, Result, anyhow};
 use arrow_array::RecordBatch;
+use protocol::storage::KafkaSourceSubtaskCheckpoint;
+use tokio::sync::mpsc;
 
 use crate::runtime::memory::{MemoryBlock, MemoryPool, get_array_memory_size};
 use crate::runtime::streaming::network::endpoint::PhysicalSender;
+use crate::runtime::streaming::protocol::control::JobMasterEvent;
 use crate::runtime::streaming::protocol::event::{StreamEvent, TrackedEvent};
 use crate::runtime::streaming::state::IoManager;
 
@@ -74,6 +77,9 @@ 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.
+    checkpoint_ack_tx: Option<mpsc::Sender<JobMasterEvent>>,
 }
 
 impl TaskContext {
@@ -90,6 +96,7 @@ impl TaskContext {
         pipeline_state_memory_block: Option<Arc<MemoryBlock>>,
         operator_state_memory_bytes: u64,
         safe_epoch: u64,
+        checkpoint_ack_tx: Option<mpsc::Sender<JobMasterEvent>>,
     ) -> Self {
         let task_name = format!(
             "Task-[{}]-Pipe[{}]-Sub[{}/{}]",
@@ -111,6 +118,7 @@ impl TaskContext {
             pipeline_state_memory_block,
             operator_state_memory_bytes,
             safe_epoch,
+            checkpoint_ack_tx,
         }
     }
 
@@ -119,6 +127,23 @@ impl TaskContext {
         self.safe_epoch
     }
 
+    /// Notify the job checkpoint coordinator that this pipeline has finished the barrier for `epoch`.
+    pub async fn send_checkpoint_ack(
+        &self,
+        epoch: u64,
+        kafka_subtask: Option<KafkaSourceSubtaskCheckpoint>,
+    ) {
+        if let Some(tx) = &self.checkpoint_ack_tx {
+            let _ = tx
+                .send(JobMasterEvent::CheckpointAck {
+                    pipeline_id: self.pipeline_id,
+                    epoch,
+                    kafka_subtask,
+                })
+                .await;
+        }
+    }
+
     #[inline]
     pub fn config(&self) -> &TaskContextConfig {
         &self.config

src/runtime/streaming/api/operator.rs

@@ -54,11 +54,24 @@ pub trait Operator: Send + 'static {
         ctx: &mut TaskContext,
     ) -> anyhow::Result<()>;
 
+    /// Global checkpoint **phase 2** (after metadata is durable): finalize external side effects.
+    ///
+    /// Default is no-op. Examples of overrides: transactional Kafka sink calls
+    /// `commit_transaction` on the producer stashed during [`Self::snapshot_state`].
     async fn commit_checkpoint(
         &mut self,
-        _epoch: u32,
+        epoch: u32,
         _ctx: &mut TaskContext,
     ) -> anyhow::Result<()> {
+        let _ = epoch;
+        Ok(())
+    }
+
+    /// 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<()> {
+        let _ = epoch;
         Ok(())
     }
 

src/runtime/streaming/api/source.rs

@@ -14,6 +14,7 @@ use crate::runtime::streaming::api::context::TaskContext;
 use crate::sql::common::{CheckpointBarrier, Watermark};
 use arrow_array::RecordBatch;
 use async_trait::async_trait;
+use protocol::storage::KafkaSourceSubtaskCheckpoint;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 pub enum SourceOffset {
@@ -31,6 +32,12 @@ pub enum SourceEvent {
     EndOfStream,
 }
 
+/// Optional metadata returned when a source completes a checkpoint barrier snapshot.
+#[derive(Debug, Default, Clone)]
+pub struct SourceCheckpointReport {
+    pub kafka_subtask: Option<KafkaSourceSubtaskCheckpoint>,
+}
+
 #[async_trait]
 pub trait SourceOperator: Send + 'static {
     fn name(&self) -> &str;
@@ -49,13 +56,22 @@ pub trait SourceOperator: Send + 'static {
         &mut self,
         barrier: CheckpointBarrier,
         ctx: &mut TaskContext,
-    ) -> anyhow::Result<()>;
+    ) -> anyhow::Result<SourceCheckpointReport>;
 
+    /// Same checkpoint **phase 2** hook as [`super::operator::Operator::commit_checkpoint`].
+    /// Kafka source keeps the default: offsets are reported at the barrier in [`Self::snapshot_state`].
     async fn commit_checkpoint(
         &mut self,
-        _epoch: u32,
+        epoch: u32,
         _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<()> {
+        let _ = epoch;
         Ok(())
     }
 

src/runtime/streaming/execution/operator_chain.rs

@@ -158,6 +158,9 @@ impl OperatorDrive for IntermediateDriver {
             ControlCommand::Commit { epoch } => {
                 self.operator.commit_checkpoint(*epoch, ctx).await?;
             }
+            ControlCommand::AbortCheckpoint { epoch } => {
+                self.operator.abort_checkpoint(*epoch, ctx).await?;
+            }
             ControlCommand::Stop { mode } if *mode == StopMode::Immediate => {
                 stop = true;
             }
@@ -273,6 +276,9 @@ impl OperatorDrive for TailDriver {
             ControlCommand::Commit { epoch } => {
                 self.operator.commit_checkpoint(*epoch, ctx).await?;
             }
+            ControlCommand::AbortCheckpoint { epoch } => {
+                self.operator.abort_checkpoint(*epoch, ctx).await?;
+            }
             ControlCommand::Stop { mode } if *mode == StopMode::Immediate => {
                 stop = true;
             }

src/runtime/streaming/execution/pipeline.rs

@@ -110,6 +110,7 @@ impl Pipeline {
                                         }
                                     }
                                     AlignmentStatus::Complete => {
+                                        let epoch = barrier.epoch as u64;
                                         self.chain_head
                                             .process_event(
                                                 idx,
@@ -123,6 +124,7 @@ impl Pipeline {
                                                 active_streams.insert(i, stream);
                                             }
                                         }
+                                        self.ctx.send_checkpoint_ack(epoch, None).await;
                                     }
                                 }
                             }

src/runtime/streaming/execution/source_driver.rs

@@ -15,7 +15,7 @@ use tokio::time::{Instant, sleep};
 use tracing::{Instrument, info, info_span, warn};
 
 use crate::runtime::streaming::api::context::TaskContext;
-use crate::runtime::streaming::api::source::{SourceEvent, SourceOperator};
+use crate::runtime::streaming::api::source::{SourceCheckpointReport, SourceEvent, SourceOperator};
 use crate::runtime::streaming::error::RunError;
 use crate::runtime::streaming::execution::OperatorDrive;
 use crate::runtime::streaming::protocol::{
@@ -154,18 +154,25 @@ impl SourceDriver {
 
     async fn handle_control(&mut self, cmd: ControlCommand) -> Result<bool, RunError> {
         let mut stop = false;
+        let mut pending_source_checkpoint: Option<(u64, SourceCheckpointReport)> = None;
 
         match &cmd {
             ControlCommand::TriggerCheckpoint { barrier } => {
                 let b: CheckpointBarrier = barrier.clone().into();
-                self.operator.snapshot_state(b, &mut self.ctx).await?;
+                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));
             }
             ControlCommand::Commit { epoch } => {
                 self.operator
                     .commit_checkpoint(*epoch, &mut self.ctx)
                     .await?;
             }
+            ControlCommand::AbortCheckpoint { epoch } => {
+                self.operator
+                    .abort_checkpoint(*epoch, &mut self.ctx)
+                    .await?;
+            }
             ControlCommand::Stop { .. } => {
                 stop = true;
             }
@@ -178,6 +185,12 @@ impl SourceDriver {
             stop = true;
         }
 
+        if let Some((epoch, report)) = pending_source_checkpoint {
+            self.ctx
+                .send_checkpoint_ack(epoch, report.kafka_subtask)
+                .await;
+        }
+
         Ok(stop)
     }
 

src/runtime/streaming/factory/connector/kafka.rs

@@ -200,7 +200,13 @@ impl KafkaConnectorDispatcher {
         let client_configs = merge_client_configs(&cfg.auth, &cfg.client_configs);
 
         let consistency = match cfg.commit_mode() {
-            KafkaSinkCommitMode::KafkaSinkExactlyOnce => ConsistencyMode::ExactlyOnce,
+            KafkaSinkCommitMode::KafkaSinkExactlyOnce => {
+                info!(
+                    topic = %cfg.topic,
+                    "Kafka sink exactly-once: transactional producer + checkpoint 2PC. Downstream Kafka consumers of this topic should set isolation.level=read_committed."
+                );
+                ConsistencyMode::ExactlyOnce
+            }
             KafkaSinkCommitMode::KafkaSinkAtLeastOnce => ConsistencyMode::AtLeastOnce,
         };