feat: Add: stream based snapshot API install_snapshot_v2() to meta-service (#15395)

feature: Add: stream based snapshot API install_snapshot_v2() to meta-service

Author: drmingdrmer

Cargo.lock

@@ -123,7 +123,7 @@ sled = { git = "https://github.com/datafuse-extras/sled", tag = "v0.34.7-datafus
 
 # openraft for debugging
 #openraft = { git = "https://github.com/drmingdrmer/openraft", branch = "release-0.9", features = [
-openraft = { version = "0.9.8", features = [
+openraft = { version = "0.9.9", features = [
     "serde",
     "tracing-log",
     "generic-snapshot-data",

Cargo.toml

@@ -18,6 +18,7 @@
 use std::sync::Arc;
 use std::time::Duration;
 
+use databend_common_base::base::tokio::io::AsyncWriteExt;
 use databend_common_base::base::tokio::sync::Mutex;
 use databend_common_base::future::TimingFutureExt;
 use databend_common_meta_client::MetaGrpcReadReq;
@@ -27,22 +28,28 @@ use databend_common_meta_types::protobuf::raft_service_server::RaftService;
 use databend_common_meta_types::protobuf::RaftReply;
 use databend_common_meta_types::protobuf::RaftRequest;
 use databend_common_meta_types::protobuf::SnapshotChunkRequest;
+use databend_common_meta_types::protobuf::SnapshotChunkRequestV2;
 use databend_common_meta_types::protobuf::StreamItem;
 use databend_common_meta_types::GrpcHelper;
 use databend_common_meta_types::InstallSnapshotError;
 use databend_common_meta_types::InstallSnapshotRequest;
 use databend_common_meta_types::InstallSnapshotResponse;
 use databend_common_meta_types::RaftError;
+use databend_common_meta_types::Snapshot;
 use databend_common_meta_types::SnapshotMeta;
 use databend_common_meta_types::TypeConfig;
 use databend_common_meta_types::Vote;
 use databend_common_metrics::count::Count;
+use futures::TryStreamExt;
+use log::debug;
+use log::info;
 use minitrace::full_name;
 use minitrace::prelude::*;
 use tonic::codegen::BoxStream;
 use tonic::Request;
 use tonic::Response;
 use tonic::Status;
+use tonic::Streaming;
 
 use crate::message::ForwardRequest;
 use crate::message::ForwardRequestBody;
@@ -154,6 +161,101 @@ impl RaftServiceImpl {
             Ok(resp)
         }
     }
+
+    async fn do_install_snapshot_v2(
+        &self,
+        request: Request<Streaming<SnapshotChunkRequestV2>>,
+    ) -> Result<Response<RaftReply>, Status> {
+        let addr = remote_addr(&request);
+
+        let _g = snapshot_recv_inflight(&addr).counter_guard();
+
+        let mut strm = request.into_inner();
+
+        // Extract the first chunk to get the rpc_meta.
+        let (format, req_vote, snapshot_meta) = {
+            let Some(chunk) = strm.try_next().await? else {
+                return Err(GrpcHelper::invalid_arg("empty snapshot stream"));
+            };
+
+            let Some(meta) = &chunk.rpc_meta else {
+                return Err(GrpcHelper::invalid_arg(
+                    "first SnapshotChunkRequestV2.rpc_meta is None",
+                ));
+            };
+
+            if !chunk.chunk.is_empty() {
+                return Err(GrpcHelper::invalid_arg(
+                    "first SnapshotChunkRequestV2.chunk should not contain any data",
+                ));
+            }
+
+            let rpc_mta: (String, Vote, SnapshotMeta) = GrpcHelper::parse(meta)?;
+            rpc_mta
+        };
+
+        // Snapshot format is not used for now.
+        let _ = format;
+
+        info!(
+            format :% = &format,
+            req_vote :% = &req_vote,
+            snapshot_meta :% = &snapshot_meta;
+            "Begin receiving snapshot v2 stream from: {}",
+            addr
+        );
+
+        let mut snapshot_data = self
+            .meta_node
+            .raft
+            .begin_receiving_snapshot()
+            .await
+            .map_err(GrpcHelper::internal_err)?;
+
+        let mut ith = 0;
+        let mut total_len = 0;
+        while let Some(chunk) = strm.try_next().await? {
+            let data_len = chunk.chunk.len() as u64;
+            total_len += data_len;
+            debug!(
+                len = data_len,
+                total_len = total_len;
+                "received {ith}-th snapshot chunk from {addr}"
+            );
+
+            ith += 1;
+            if ith % 100 == 0 {
+                info!(
+                    total_len = total_len;
+                    "received {ith}-th snapshot chunk from {addr}"
+                );
+            }
+
+            raft_metrics::network::incr_recvfrom_bytes(addr.clone(), data_len);
+
+            snapshot_data.write_all(&chunk.chunk).await?;
+        }
+
+        snapshot_data.shutdown().await?;
+
+        let snapshot = Snapshot {
+            meta: snapshot_meta,
+            snapshot: snapshot_data,
+        };
+
+        let res = self
+            .meta_node
+            .raft
+            .install_full_snapshot(req_vote, snapshot)
+            .await
+            .map_err(GrpcHelper::internal_err);
+
+        raft_metrics::network::incr_snapshot_recvfrom_result(addr.clone(), res.is_ok());
+
+        let resp = res?;
+
+        GrpcHelper::ok_response(&resp)
+    }
 }
 
 #[async_trait::async_trait]
@@ -241,6 +343,14 @@ impl RaftService for RaftServiceImpl {
         self.do_install_snapshot_v1(request).in_span(root).await
     }
 
+    async fn install_snapshot_v2(
+        &self,
+        request: Request<Streaming<SnapshotChunkRequestV2>>,
+    ) -> Result<Response<RaftReply>, Status> {
+        let root = databend_common_tracing::start_trace_for_remote_request(full_name!(), &request);
+        self.do_install_snapshot_v2(request).in_span(root).await
+    }
+
     async fn vote(&self, request: Request<RaftRequest>) -> Result<Response<RaftReply>, Status> {
         let root = databend_common_tracing::start_trace_for_remote_request(full_name!(), &request);
 

src/meta/service/src/meta_service/raft_service_impl.rs

@@ -12,19 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use databend_common_arrow::arrow::array::ViewType;
 use databend_common_meta_raft_store::state_machine::MetaSnapshotId;
+use databend_common_meta_sled_store::openraft::error::SnapshotMismatch;
+use databend_common_meta_sled_store::openraft::testing::log_id;
 use databend_common_meta_sled_store::openraft::LogIdOptionExt;
 use databend_common_meta_sled_store::openraft::ServerState;
 use databend_common_meta_types::protobuf::SnapshotChunkRequest;
+use databend_common_meta_types::protobuf::SnapshotChunkRequestV2;
 use databend_common_meta_types::Cmd;
+use databend_common_meta_types::InstallSnapshotError;
 use databend_common_meta_types::InstallSnapshotRequest;
 use databend_common_meta_types::LogEntry;
+use databend_common_meta_types::RaftError;
 use databend_common_meta_types::SeqV;
 use databend_common_meta_types::SnapshotMeta;
+use databend_common_meta_types::SnapshotResponse;
 use databend_common_meta_types::StoredMembership;
 use databend_common_meta_types::UpsertKV;
 use databend_common_meta_types::Vote;
 use databend_meta::meta_service::MetaNode;
+use futures::stream;
 use log::info;
 use maplit::btreeset;
 use test_harness::test;
@@ -171,8 +179,67 @@ async fn test_raft_service_snapshot_id_mismatch() -> anyhow::Result<()> {
     r1.meta.snapshot_id = MetaSnapshotId::new(None, 2).to_string();
     r1.offset = 3;
     let req = SnapshotChunkRequest::new_v1(r1);
-    let res = client0.install_snapshot_v1(req).await;
-    println!("res: {:?}", res);
+    let resp = client0.install_snapshot_v1(req).await?;
+
+    let reply = resp.into_inner();
+
+    let err: RaftError<InstallSnapshotError> = serde_json::from_str(&reply.error)?;
+
+    assert_eq!(
+        err.api_error().unwrap(),
+        &InstallSnapshotError::SnapshotMismatch(SnapshotMismatch {
+            expect: (MetaSnapshotId::new(None, 2), 0).into(),
+            got: (MetaSnapshotId::new(None, 2), 3).into(),
+        })
+    );
+
+    Ok(())
+}
+
+#[test(harness = meta_service_test_harness)]
+#[minitrace::trace]
+async fn test_raft_service_install_snapshot_v2() -> anyhow::Result<()> {
+    // Transmit snapshot in one-piece in a stream via API install_snapshot_v2.
+
+    let (_nlog, mut tcs) = start_meta_node_cluster(btreeset![0], btreeset![]).await?;
+    let tc0 = tcs.remove(0);
+
+    let mut client0 = tc0.raft_client().await?;
+
+    let last = log_id(10, 2, 4);
+
+    let snapshot_meta = SnapshotMeta {
+        last_log_id: Some(last),
+        last_membership: StoredMembership::default(),
+        snapshot_id: MetaSnapshotId::new(Some(last), 1).to_string(),
+    };
+
+    let data = [
+        r#"{"DataHeader":{"key":"header","value":{"version":"V002","upgrading":null}}}"#,
+        r#"{"StateMachineMeta":{"key":"LastApplied","value":{"LogId":{"leader_id":{"term":10,"node_id":2},"index":4}}}}"#,
+        r#"{"StateMachineMeta":{"key":"LastMembership","value":{"Membership":{"log_id":{"leader_id":{"term":3,"node_id":3},"index":3},"membership":{"configs":[],"nodes":{}}}}}}"#,
+    ];
+
+    let strm_data = [
+        SnapshotChunkRequestV2::new_head(Vote::new_committed(10, 2), snapshot_meta),
+        SnapshotChunkRequestV2::new_chunk(data[0].to_bytes().to_vec()),
+        SnapshotChunkRequestV2::new_chunk("\n".as_bytes().to_vec()),
+        SnapshotChunkRequestV2::new_chunk(data[1].to_bytes().to_vec()),
+        SnapshotChunkRequestV2::new_chunk("\n".as_bytes().to_vec()),
+        SnapshotChunkRequestV2::new_chunk(data[2].to_bytes().to_vec()),
+    ];
+
+    let resp = client0.install_snapshot_v2(stream::iter(strm_data)).await?;
+    let reply = resp.into_inner();
+
+    let resp: SnapshotResponse = serde_json::from_str(&reply.data)?;
+
+    assert_eq!(resp.vote, Vote::new_committed(10, 2));
+
+    let meta_node = tc0.meta_node.as_ref().unwrap();
+    let m = meta_node.raft.metrics().borrow().clone();
+
+    assert_eq!(Some(last), m.snapshot);
 
     Ok(())
 }

src/meta/service/tests/it/meta_node/meta_node_replication.rs

@@ -58,7 +58,10 @@ fn setup_test() {
         let t = tempfile::tempdir().expect("create temp dir to sled db");
         databend_common_meta_sled_store::init_temp_sled_db(t);
 
-        let guards = init_logging("meta_unittests", &Config::new_testing(), BTreeMap::new());
+        let mut config = Config::new_testing();
+        config.file.prefix_filter = "".to_string();
+
+        let guards = init_logging("meta_unittests", &config, BTreeMap::new());
         Box::leak(Box::new(guards));
     });
 }

src/meta/service/tests/it/testing.rs

@@ -198,7 +198,7 @@ message SnapshotChunkRequest {
 
   // json serialized meta data, including vote and snapshot_meta.
   // ```text
-  // { vote: Vote, meta: SnapshotMeta }
+  // (Vote, SnapshotMeta)
   // ```
   string rpc_meta = 10;
 
@@ -211,6 +211,25 @@ message SnapshotChunkV1 {
   bytes data = 12;
 }
 
+// The item of snapshot chunk stream.
+//
+// The first item contains `rpc_meta`,
+// including the application defined format of this snapshot data,
+// the leader vote and snapshot-meta.
+//
+// Since the second item, the `rpc_meta` should be empty and will be ignored by the receiving end.
+message SnapshotChunkRequestV2 {
+
+  // json serialized meta data, including vote and snapshot_meta.
+  // ```text
+  // (SnapshotFormat, Vote, SnapshotMeta)
+  // ```
+  optional string rpc_meta = 10;
+
+  // Snapshot data chunk
+  bytes chunk = 20;
+}
+
 service RaftService {
 
   // Forward a request to another node.
@@ -225,6 +244,7 @@ service RaftService {
   rpc AppendEntries(RaftRequest) returns (RaftReply);
   rpc InstallSnapshot(RaftRequest) returns (RaftReply);
   rpc InstallSnapshotV1(SnapshotChunkRequest) returns (RaftReply);
+  rpc InstallSnapshotV2(stream SnapshotChunkRequestV2) returns (RaftReply);
   rpc Vote(RaftRequest) returns (RaftReply);
 }
 

src/meta/types/proto/meta.proto

@@ -147,6 +147,7 @@ pub use crate::raft_types::RaftMetrics;
 pub use crate::raft_types::RemoteError;
 pub use crate::raft_types::Snapshot;
 pub use crate::raft_types::SnapshotMeta;
+pub use crate::raft_types::SnapshotMismatch;
 pub use crate::raft_types::SnapshotResponse;
 pub use crate::raft_types::StorageError;
 pub use crate::raft_types::StorageIOError;

src/meta/types/src/lib.rs

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 use crate::protobuf::SnapshotChunkRequest;
+use crate::protobuf::SnapshotChunkRequestV2;
 use crate::protobuf::SnapshotChunkV1;
 use crate::InstallSnapshotRequest;
+use crate::SnapshotMeta;
+use crate::Vote;
 
 impl SnapshotChunkRequest {
     /// Return the length of the data in the chunk.
@@ -39,3 +42,24 @@ impl SnapshotChunkRequest {
         }
     }
 }
+
+impl SnapshotChunkRequestV2 {
+    /// Build the first chunk of a snapshot stream, which contains vote and snapshot meta, without data.
+    pub fn new_head(vote: Vote, snapshot_meta: SnapshotMeta) -> Self {
+        let meta = ("ndjson".to_string(), vote, snapshot_meta);
+        let rpc_meta = serde_json::to_string(&meta).unwrap();
+
+        SnapshotChunkRequestV2 {
+            rpc_meta: Some(rpc_meta),
+            chunk: vec![],
+        }
+    }
+
+    /// Build a chunk item with data.
+    pub fn new_chunk(chunk: Vec<u8>) -> Self {
+        SnapshotChunkRequestV2 {
+            rpc_meta: None,
+            chunk,
+        }
+    }
+}

src/meta/types/src/proto_ext/snapshot_chunk_request_ext.rs

@@ -77,6 +77,7 @@ pub type InstallSnapshotRequest = openraft::raft::InstallSnapshotRequest<TypeCon
 pub type InstallSnapshotResponse = openraft::raft::InstallSnapshotResponse<NodeId>;
 pub type SnapshotResponse = openraft::raft::SnapshotResponse<NodeId>;
 pub type InstallSnapshotError = openraft::error::InstallSnapshotError;
+pub type SnapshotMismatch = openraft::error::SnapshotMismatch;
 pub type VoteRequest = openraft::raft::VoteRequest<NodeId>;
 pub type VoteResponse = openraft::raft::VoteResponse<NodeId>;