refactor: in snapshot-v2 protocol, move head to end to inform explicitly streaming end (#15399)

Author: drmingdrmer

src/meta/service/src/meta_service/mod.rs

@@ -12,6 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod errors;
+mod forwarder;
+mod meta_node_kv_api_impl;
+
+pub(crate) mod snapshot_receiver;
+
+pub mod meta_leader;
+pub mod meta_node;
+pub mod raft_service_impl;
+
 pub use forwarder::MetaForwarder;
 pub use meta_node::MetaNode;
 pub use raft_service_impl::RaftServiceImpl;
@@ -20,10 +30,3 @@ pub use crate::message::ForwardRequest;
 pub use crate::message::ForwardRequestBody;
 pub use crate::message::JoinRequest;
 pub use crate::message::LeaveRequest;
-
-mod errors;
-mod forwarder;
-pub mod meta_leader;
-pub mod meta_node;
-mod meta_node_kv_api_impl;
-pub mod raft_service_impl;

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

@@ -18,7 +18,6 @@
 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;
@@ -35,14 +34,11 @@ 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;
@@ -53,6 +49,7 @@ use tonic::Streaming;
 
 use crate::message::ForwardRequest;
 use crate::message::ForwardRequestBody;
+use crate::meta_service::snapshot_receiver::Receiver;
 use crate::meta_service::MetaNode;
 use crate::metrics::raft_metrics;
 
@@ -168,93 +165,42 @@ impl RaftServiceImpl {
     ) -> 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
-        };
+        let _guard = snapshot_recv_inflight(&addr).counter_guard();
 
-        // 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
+        let 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);
+        let mut receiver = Receiver::new(&addr, snapshot_data);
 
-            snapshot_data.write_all(&chunk.chunk).await?;
-        }
+        let mut strm = request.into_inner();
 
-        snapshot_data.shutdown().await?;
+        while let Some(chunk) = strm.try_next().await? {
+            let snapshot = receiver.receive(chunk).await?;
 
-        let snapshot = Snapshot {
-            meta: snapshot_meta,
-            snapshot: snapshot_data,
-        };
+            if let Some((_format, req_vote, snapshot)) = snapshot {
+                let res = self
+                    .meta_node
+                    .raft
+                    .install_full_snapshot(req_vote, snapshot)
+                    .await
+                    .map_err(GrpcHelper::internal_err);
 
-        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());
 
-        raft_metrics::network::incr_snapshot_recvfrom_result(addr.clone(), res.is_ok());
+                let resp = res?;
 
-        let resp = res?;
+                return GrpcHelper::ok_response(&resp);
+            }
+        }
 
-        GrpcHelper::ok_response(&resp)
+        Err(Status::invalid_argument(format!(
+            "snapshot stream is closed without finishing: {}",
+            receiver.stat_str()
+        )))
     }
 }
 

src/meta/service/src/meta_service/snapshot_receiver.rs

@@ -0,0 +1,149 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt;
+
+use databend_common_base::base::tokio::io::AsyncWriteExt;
+use databend_common_meta_types::protobuf::SnapshotChunkRequestV2;
+use databend_common_meta_types::Snapshot;
+use databend_common_meta_types::SnapshotData;
+use databend_common_meta_types::SnapshotMeta;
+use databend_common_meta_types::Vote;
+use log::debug;
+use log::info;
+use tonic::Status;
+
+use crate::metrics::raft_metrics;
+
+pub(crate) struct Receiver {
+    remote_addr: String,
+
+    snapshot_data: Option<Box<SnapshotData>>,
+
+    /// number of bytes received.
+    n_received: usize,
+
+    /// number of bytes received.
+    size_received: usize,
+}
+
+impl Receiver {
+    /// Create a new snapshot receiver with an empty snapshot.
+    pub(crate) fn new(remote_addr: impl ToString, snapshot_data: Box<SnapshotData>) -> Self {
+        let remote_addr = remote_addr.to_string();
+        info!("Begin receiving snapshot v2 stream from: {}", remote_addr);
+
+        Receiver {
+            remote_addr,
+            snapshot_data: Some(snapshot_data),
+            n_received: 0,
+            size_received: 0,
+        }
+    }
+
+    pub(crate) fn stat_str(&self) -> impl fmt::Display {
+        format!(
+            "received {} chunks, {} bytes from {}",
+            self.n_received, self.size_received, self.remote_addr
+        )
+    }
+
+    pub(crate) async fn receive(
+        &mut self,
+        chunk: SnapshotChunkRequestV2,
+    ) -> Result<Option<(String, Vote, Snapshot)>, Status> {
+        // 1. update stat
+        self.update_stat(&chunk);
+
+        // 2. write chunk to local snapshot_data
+        {
+            let snapshot_data = self.snapshot_data.as_mut().ok_or_else(|| {
+                Status::internal("snapshot_data is already shutdown when receiving snapshot chunk")
+            })?;
+
+            snapshot_data.write_all(&chunk.chunk).await.map_err(|e| {
+                Status::internal(format!(
+                    "{} when writing chunk to local temp snapshot_data",
+                    e
+                ))
+            })?;
+        }
+
+        // 3. if it is the last chunk, finish and return the snapshot.
+        {
+            let end = self.load_end(&chunk).map_err(|e| {
+                Status::invalid_argument(format!("{} when loading last chunk rpc_meta", e))
+            })?;
+
+            let Some((format, vote, snapshot_meta)) = end else {
+                return Ok(None);
+            };
+
+            info!(
+                "snapshot from {} is completely received, format: {}, vote: {:?}, meta: {:?}, size: {}",
+                self.remote_addr, format, vote, snapshot_meta, self.size_received
+            );
+
+            // Safe unwrap: snapshot_data is guaranteed to be Some in the above code.
+            let mut snapshot_data = self.snapshot_data.take().unwrap();
+
+            snapshot_data.shutdown().await.map_err(|e| {
+                Status::internal(format!("{} when shutdown local temp snapshot_data", e))
+            })?;
+
+            Ok(Some((format, vote, Snapshot {
+                meta: snapshot_meta,
+                snapshot: snapshot_data,
+            })))
+        }
+    }
+
+    fn update_stat(&mut self, chunk: &SnapshotChunkRequestV2) {
+        let data_len = chunk.chunk.len();
+        self.n_received += 1;
+        self.size_received += data_len;
+
+        debug!(
+            len = data_len,
+            total_len = self.size_received;
+            "received {}-th snapshot chunk from {}",
+            self.n_received, self.remote_addr
+        );
+
+        if self.n_received % 100 == 0 {
+            info!(
+                total_len = self.size_received;
+                "received {}-th snapshot chunk from {}",
+                self.n_received, self.remote_addr
+            );
+        }
+
+        raft_metrics::network::incr_recvfrom_bytes(self.remote_addr.clone(), data_len as u64);
+    }
+
+    /// Load meta data from the last chunk.
+    fn load_end(
+        &self,
+        chunk: &SnapshotChunkRequestV2,
+    ) -> Result<Option<(String, Vote, SnapshotMeta)>, serde_json::Error> {
+        let Some(meta) = &chunk.rpc_meta else {
+            return Ok(None);
+        };
+
+        let (format, vote, snapshot_meta): (String, Vote, SnapshotMeta) =
+            serde_json::from_str(meta)?;
+
+        Ok(Some((format, vote, snapshot_meta)))
+    }
+}

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

@@ -206,12 +206,12 @@ async fn test_raft_service_install_snapshot_v2() -> anyhow::Result<()> {
 
     let mut client0 = tc0.raft_client().await?;
 
-    let last = log_id(10, 2, 4);
+    let last_log_id = log_id(10, 2, 4);
 
     let snapshot_meta = SnapshotMeta {
-        last_log_id: Some(last),
+        last_log_id: Some(last_log_id),
         last_membership: StoredMembership::default(),
-        snapshot_id: MetaSnapshotId::new(Some(last), 1).to_string(),
+        snapshot_id: MetaSnapshotId::new(Some(last_log_id), 1).to_string(),
     };
 
     let data = [
@@ -220,13 +220,15 @@ async fn test_raft_service_install_snapshot_v2() -> anyhow::Result<()> {
         r#"{"StateMachineMeta":{"key":"LastMembership","value":{"Membership":{"log_id":{"leader_id":{"term":3,"node_id":3},"index":3},"membership":{"configs":[],"nodes":{}}}}}}"#,
     ];
 
+    // Complete transmit
+
     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()),
+        SnapshotChunkRequestV2::new_end_chunk(Vote::new_committed(10, 2), snapshot_meta),
     ];
 
     let resp = client0.install_snapshot_v2(stream::iter(strm_data)).await?;
@@ -239,7 +241,21 @@ async fn test_raft_service_install_snapshot_v2() -> anyhow::Result<()> {
     let meta_node = tc0.meta_node.as_ref().unwrap();
     let m = meta_node.raft.metrics().borrow().clone();
 
-    assert_eq!(Some(last), m.snapshot);
+    assert_eq!(Some(last_log_id), m.snapshot);
+
+    // Incomplete
+
+    let strm_data = [
+        SnapshotChunkRequestV2::new_chunk(data[0].to_bytes().to_vec()),
+        SnapshotChunkRequestV2::new_chunk("\n".as_bytes().to_vec()),
+    ];
+
+    let err = client0
+        .install_snapshot_v2(stream::iter(strm_data))
+        .await
+        .unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::InvalidArgument);
 
     Ok(())
 }

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

@@ -44,8 +44,8 @@ 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 {
+    /// Build the last chunk of a snapshot stream, which contains vote and snapshot meta, without data.
+    pub fn new_end_chunk(vote: Vote, snapshot_meta: SnapshotMeta) -> Self {
         let meta = ("ndjson".to_string(), vote, snapshot_meta);
         let rpc_meta = serde_json::to_string(&meta).unwrap();