fix(zk): checkpoint reconciliation after staged entries

Author: wei

dt-connector/src/extractor/zk/zk_extractor.rs

@@ -190,21 +190,31 @@ impl ZkExtractor {
             scan_result.entries.push(entry);
         }
 
-        let position = self.build_position(&scan_result.versions, scan_result.high_water_zxid);
+        let final_position =
+            self.build_position(&scan_result.versions, scan_result.high_water_zxid);
+        let emitted_entries = scan_result.entries.len();
         for entry in scan_result.entries {
             self.extract_state
-                .push_dt_data(&self.base_extractor, DtData::Zk { entry }, position.clone())
+                .push_dt_data(
+                    &self.base_extractor,
+                    DtData::Zk { entry },
+                    Self::staged_entry_position(),
+                )
                 .await?;
         }
+        self.extract_state
+            .push_dt_data(&self.base_extractor, DtData::Heartbeat {}, final_position)
+            .await?;
 
         *path_versions = scan_result.versions;
         *high_water_zxid = scan_result.high_water_zxid;
 
         log_info!(
-            "ZkExtractor reconciliation complete: reason={}, {} paths, high_water_zxid={}",
+            "ZkExtractor reconciliation complete: reason={}, {} paths, high_water_zxid={}, emitted_entries={}",
             reason,
             path_versions.len(),
-            high_water_zxid
+            high_water_zxid,
+            emitted_entries
         );
         Ok(())
     }
@@ -235,6 +245,10 @@ impl ZkExtractor {
         root == "/" || path == root || path.starts_with(&format!("{}/", root.trim_end_matches('/')))
     }
 
+    fn staged_entry_position() -> Position {
+        Position::None
+    }
+
     async fn full_scan(
         &mut self,
         client: &ZkClient,
@@ -688,4 +702,9 @@ mod tests {
         assert!(scan_result.versions.is_empty());
         assert!(scan_result.entries.is_empty());
     }
+
+    #[test]
+    fn staged_reconciliation_entries_do_not_carry_final_position() {
+        assert_eq!(ZkExtractor::staged_entry_position(), Position::None);
+    }
 }

dt-pipeline/src/base_pipeline.rs

@@ -663,6 +663,7 @@ mod tests {
         dt_data::{DtData, DtItem},
         position::Position,
         redis::redis_entry::RedisEntry,
+        zk::{zk_entry::ZkEntry, zk_event_type::ZkEventType, zk_stat::ZkStat},
     };
     use dt_connector::extractor::resumer::utils::ResumerUtil;
 
@@ -690,6 +691,34 @@ mod tests {
         }
     }
 
+    fn zk_item(position: Position) -> DtItem {
+        DtItem {
+            dt_data: DtData::Zk {
+                entry: ZkEntry {
+                    path: "/app/service".to_string(),
+                    data: Some(b"value".to_vec()),
+                    stat: ZkStat::default(),
+                    ephemeral: false,
+                    event_type: ZkEventType::Updated,
+                    source_id: "node-a".to_string(),
+                    source_order_millis: 1,
+                    source_zxid: 1,
+                    order_origin: Default::default(),
+                },
+            },
+            position,
+            data_origin_node: String::new(),
+        }
+    }
+
+    fn heartbeat_item(position: Position) -> DtItem {
+        DtItem {
+            dt_data: DtData::Heartbeat {},
+            position,
+            data_origin_node: String::new(),
+        }
+    }
+
     #[test]
     fn fetch_raw_collects_latest_position_per_redis_node() {
         let mut pending_snapshot_finished = HashMap::new();
@@ -713,4 +742,31 @@ mod tests {
         assert_eq!(by_key.get("redis-node-node-1"), Some(&node_1_new));
         assert_eq!(by_key.get("redis-node-node-2"), Some(&node_2));
     }
+
+    #[test]
+    fn fetch_raw_advances_zk_position_only_on_trailing_heartbeat() {
+        let mut pending_snapshot_finished = HashMap::new();
+        let final_position = Position::Zk {
+            path_versions: HashMap::from([("/app/service".to_string(), (1, 10))]),
+            last_scan_timestamp: 100,
+            high_water_zxid: 10,
+            total_paths: 1,
+        };
+        let partial_data = vec![zk_item(Position::None)];
+
+        let (_, partial_position, partial_commits) =
+            BasePipeline::fetch_raw(&partial_data, &mut pending_snapshot_finished);
+        assert_eq!(partial_position, Some(Position::None));
+        assert!(partial_commits.is_empty());
+
+        let complete_data = vec![
+            zk_item(Position::None),
+            heartbeat_item(final_position.clone()),
+        ];
+        let (_, complete_position, complete_commits) =
+            BasePipeline::fetch_raw(&complete_data, &mut pending_snapshot_finished);
+
+        assert_eq!(complete_position, Some(final_position.clone()));
+        assert_eq!(complete_commits, vec![final_position]);
+    }
 }