meta_srv/

cluster.rs

// Copyright 2023 Greptime Team
//
// 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::any::Any;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;

use api::v1::meta::cluster_client::ClusterClient;
use api::v1::meta::{
    BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse,
    RangeRequest as PbRangeRequest, RangeResponse as PbRangeResponse, ResponseHeader,
};
use common_grpc::channel_manager::ChannelManager;
use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue};
use common_meta::kv_backend::{KvBackend, ResettableKvBackendRef, TxnService};
use common_meta::rpc::KeyValue;
use common_meta::rpc::store::{
    BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest,
    BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest,
    DeleteRangeResponse, PutRequest, PutResponse, RangeRequest, RangeResponse,
};
use common_meta::util;
use common_telemetry::warn;
use derive_builder::Builder;
use snafu::{OptionExt, ResultExt, ensure};

use crate::error::{self, Result, match_for_io_error};
use crate::metasrv::ElectionRef;

pub type MetaPeerClientRef = Arc<MetaPeerClient>;

#[derive(Builder)]
pub struct MetaPeerClient {
    election: Option<ElectionRef>,
    in_memory: ResettableKvBackendRef,
    #[builder(default = "ChannelManager::default()")]
    channel_manager: ChannelManager,
    #[builder(default = "3")]
    max_retry_count: usize,
    #[builder(default = "1000")]
    retry_interval_ms: u64,
}

#[async_trait::async_trait]
impl TxnService for MetaPeerClient {
    type Error = error::Error;
}

#[async_trait::async_trait]
impl KvBackend for MetaPeerClient {
    fn name(&self) -> &str {
        "MetaPeerClient"
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
        if self.is_leader() {
            return self
                .in_memory
                .range(req)
                .await
                .context(error::KvBackendSnafu);
        }

        let max_retry_count = self.max_retry_count;
        let retry_interval_ms = self.retry_interval_ms;

        for _ in 0..max_retry_count {
            match self.remote_range(req.clone()).await {
                Ok(res) => return Ok(res),
                Err(e) => {
                    if need_retry(&e) {
                        warn!(e; "Encountered an error that need to retry");
                        tokio::time::sleep(Duration::from_millis(retry_interval_ms)).await;
                    } else {
                        return Err(e);
                    }
                }
            }
        }

        error::ExceededRetryLimitSnafu {
            func_name: "range",
            retry_num: max_retry_count,
        }
        .fail()
    }

    // MetaPeerClient does not support mutable methods listed below.
    async fn put(&self, _req: PutRequest) -> Result<PutResponse> {
        error::UnsupportedSnafu {
            operation: "put".to_string(),
        }
        .fail()
    }

    async fn batch_put(&self, _req: BatchPutRequest) -> Result<BatchPutResponse> {
        error::UnsupportedSnafu {
            operation: "batch put".to_string(),
        }
        .fail()
    }

    // Get kv information from the leader's in_mem kv store
    async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse> {
        if self.is_leader() {
            return self
                .in_memory
                .batch_get(req)
                .await
                .context(error::KvBackendSnafu);
        }

        let max_retry_count = self.max_retry_count;
        let retry_interval_ms = self.retry_interval_ms;

        for _ in 0..max_retry_count {
            match self.remote_batch_get(req.keys.clone()).await {
                Ok(res) => return Ok(res),
                Err(e) => {
                    if need_retry(&e) {
                        warn!(e; "Encountered an error that need to retry");
                        tokio::time::sleep(Duration::from_millis(retry_interval_ms)).await;
                    } else {
                        return Err(e);
                    }
                }
            }
        }

        error::ExceededRetryLimitSnafu {
            func_name: "batch_get",
            retry_num: max_retry_count,
        }
        .fail()
    }

    async fn delete_range(&self, _req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
        error::UnsupportedSnafu {
            operation: "delete range".to_string(),
        }
        .fail()
    }

    async fn batch_delete(&self, _req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
        error::UnsupportedSnafu {
            operation: "batch delete".to_string(),
        }
        .fail()
    }

    async fn compare_and_put(&self, _req: CompareAndPutRequest) -> Result<CompareAndPutResponse> {
        error::UnsupportedSnafu {
            operation: "compare and put".to_string(),
        }
        .fail()
    }

    async fn put_conditionally(
        &self,
        _key: Vec<u8>,
        _value: Vec<u8>,
        _if_not_exists: bool,
    ) -> Result<bool> {
        error::UnsupportedSnafu {
            operation: "put conditionally".to_string(),
        }
        .fail()
    }

    async fn delete(&self, _key: &[u8], _prev_kv: bool) -> Result<Option<KeyValue>> {
        error::UnsupportedSnafu {
            operation: "delete".to_string(),
        }
        .fail()
    }
}

impl MetaPeerClient {
    async fn get_dn_key_value(&self, keys_only: bool) -> Result<Vec<KeyValue>> {
        let key = DatanodeStatKey::prefix_key();
        let range_end = util::get_prefix_end_key(&key);
        let range_request = RangeRequest {
            key,
            range_end,
            keys_only,
            ..Default::default()
        };
        self.range(range_request).await.map(|res| res.kvs)
    }

    // Get all datanode stat kvs from leader meta.
    pub async fn get_all_dn_stat_kvs(&self) -> Result<HashMap<DatanodeStatKey, DatanodeStatValue>> {
        let kvs = self.get_dn_key_value(false).await?;
        to_stat_kv_map(kvs)
    }

    pub async fn get_node_cnt(&self) -> Result<i32> {
        let kvs = self.get_dn_key_value(true).await?;
        kvs.into_iter()
            .map(|kv| {
                kv.key
                    .try_into()
                    .context(error::InvalidDatanodeStatFormatSnafu {})
            })
            .collect::<Result<HashSet<DatanodeStatKey>>>()
            .map(|hash_set| hash_set.len() as i32)
    }

    // Get datanode stat kvs from leader meta by input keys.
    pub async fn get_dn_stat_kvs(
        &self,
        keys: Vec<DatanodeStatKey>,
    ) -> Result<HashMap<DatanodeStatKey, DatanodeStatValue>> {
        let stat_keys = keys.into_iter().map(|key| key.into()).collect();
        let batch_get_req = BatchGetRequest { keys: stat_keys };

        let res = self.batch_get(batch_get_req).await?;

        to_stat_kv_map(res.kvs)
    }

    async fn remote_range(&self, req: RangeRequest) -> Result<RangeResponse> {
        // Safety: when self.is_leader() == false, election must not empty.
        let election = self.election.as_ref().unwrap();

        let leader_addr = election.leader().await.context(error::KvBackendSnafu)?.0;

        let channel = self
            .channel_manager
            .get(&leader_addr)
            .context(error::CreateChannelSnafu)?;

        let request = tonic::Request::new(PbRangeRequest::from(req));

        let mut client =
            common_grpc::configure_tonic_client!(ClusterClient::new(channel), self.channel_manager);
        let response: PbRangeResponse = client
            .range(request)
            .await
            .context(error::RangeSnafu)?
            .into_inner();

        check_resp_header(&response.header, Context { addr: &leader_addr })?;

        Ok(RangeResponse {
            kvs: response.kvs.into_iter().map(KeyValue::new).collect(),
            more: response.more,
        })
    }

    async fn remote_batch_get(&self, keys: Vec<Vec<u8>>) -> Result<BatchGetResponse> {
        // Safety: when self.is_leader() == false, election must not empty.
        let election = self.election.as_ref().unwrap();

        let leader_addr = election.leader().await.context(error::KvBackendSnafu)?.0;

        let channel = self
            .channel_manager
            .get(&leader_addr)
            .context(error::CreateChannelSnafu)?;

        let request = tonic::Request::new(PbBatchGetRequest {
            keys,
            ..Default::default()
        });

        let mut client =
            common_grpc::configure_tonic_client!(ClusterClient::new(channel), self.channel_manager);
        let response: PbBatchGetResponse = client
            .batch_get(request)
            .await
            .context(error::BatchGetSnafu)?
            .into_inner();

        check_resp_header(&response.header, Context { addr: &leader_addr })?;

        Ok(BatchGetResponse {
            kvs: response.kvs.into_iter().map(KeyValue::new).collect(),
        })
    }

    // Check if the meta node is a leader node.
    // Note: when self.election is None, we also consider the meta node is leader
    pub(crate) fn is_leader(&self) -> bool {
        self.election
            .as_ref()
            .map(|election| election.is_leader())
            .unwrap_or(true)
    }

    #[cfg(test)]
    pub(crate) fn memory_backend(&self) -> ResettableKvBackendRef {
        self.in_memory.clone()
    }
}

fn to_stat_kv_map(kvs: Vec<KeyValue>) -> Result<HashMap<DatanodeStatKey, DatanodeStatValue>> {
    let mut map = HashMap::with_capacity(kvs.len());
    for kv in kvs {
        let _ = map.insert(
            kv.key
                .try_into()
                .context(error::InvalidDatanodeStatFormatSnafu {})?,
            kv.value
                .try_into()
                .context(error::InvalidDatanodeStatFormatSnafu {})?,
        );
    }
    Ok(map)
}

struct Context<'a> {
    addr: &'a str,
}

fn check_resp_header(header: &Option<ResponseHeader>, ctx: Context) -> Result<()> {
    let header = header
        .as_ref()
        .context(error::ResponseHeaderNotFoundSnafu)?;

    ensure!(
        !header.is_not_leader(),
        error::IsNotLeaderSnafu {
            node_addr: ctx.addr
        }
    );

    Ok(())
}

fn need_retry(error: &error::Error) -> bool {
    match error {
        error::Error::IsNotLeader { .. } => true,
        error::Error::Range { error, .. } | error::Error::BatchGet { error, .. } => {
            match_for_io_error(error).is_some()
        }
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicI64, Ordering};

    use api::v1::meta::cluster_server::{Cluster, ClusterServer};
    use api::v1::meta::{
        BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse, Error,
        ErrorCode, MetasrvPeersRequest, MetasrvPeersResponse, RangeRequest as PbRangeRequest,
        RangeResponse as PbRangeResponse, ResponseHeader,
    };
    use common_grpc::channel_manager::ChannelManager;
    use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat};
    use common_meta::election::{Election, LeaderChangeMessage, LeaderValue, MetasrvNodeInfo};
    use common_meta::kv_backend::KvBackend;
    use common_meta::kv_backend::memory::MemoryKvBackend;
    use common_meta::rpc::KeyValue;
    use common_meta::rpc::store::RangeRequest;
    use hyper_util::rt::TokioIo;
    use tonic::{Request, Response, Status};
    use tower::service_fn;

    use super::{Context, MetaPeerClientBuilder, check_resp_header, to_stat_kv_map};
    use crate::error;

    struct FollowerElection {
        leader_addr: String,
    }

    #[async_trait::async_trait]
    impl Election for FollowerElection {
        type Leader = LeaderValue;

        fn is_leader(&self) -> bool {
            false
        }

        fn in_leader_infancy(&self) -> bool {
            false
        }

        async fn register_candidate(&self, _: &MetasrvNodeInfo) -> common_meta::error::Result<()> {
            Ok(())
        }

        async fn all_candidates(&self) -> common_meta::error::Result<Vec<MetasrvNodeInfo>> {
            Ok(vec![])
        }

        async fn campaign(&self) -> common_meta::error::Result<()> {
            Ok(())
        }

        async fn leader(&self) -> common_meta::error::Result<Self::Leader> {
            Ok(LeaderValue(self.leader_addr.clone()))
        }

        async fn resign(&self) -> common_meta::error::Result<()> {
            Ok(())
        }

        fn subscribe_leader_change(&self) -> tokio::sync::broadcast::Receiver<LeaderChangeMessage> {
            let (_, receiver) = tokio::sync::broadcast::channel(1);
            receiver
        }
    }

    struct RangeServer {
        requested_limit: Arc<AtomicI64>,
    }

    #[async_trait::async_trait]
    impl Cluster for RangeServer {
        async fn batch_get(
            &self,
            _: Request<PbBatchGetRequest>,
        ) -> std::result::Result<Response<PbBatchGetResponse>, Status> {
            Err(Status::unimplemented("batch_get is not used in this test"))
        }

        async fn range(
            &self,
            request: Request<PbRangeRequest>,
        ) -> std::result::Result<Response<PbRangeResponse>, Status> {
            self.requested_limit
                .store(request.into_inner().limit, Ordering::Relaxed);
            Ok(Response::new(PbRangeResponse {
                header: Some(ResponseHeader::success()),
                kvs: vec![api::v1::meta::KeyValue {
                    key: b"key".to_vec(),
                    value: b"value".to_vec(),
                }],
                more: true,
            }))
        }

        async fn metasrv_peers(
            &self,
            _: Request<MetasrvPeersRequest>,
        ) -> std::result::Result<Response<MetasrvPeersResponse>, Status> {
            Err(Status::unimplemented(
                "metasrv_peers is not used in this test",
            ))
        }
    }

    #[tokio::test]
    async fn test_follower_range_forwards_limit() {
        let requested_limit = Arc::new(AtomicI64::new(0));
        let range_server = RangeServer {
            requested_limit: requested_limit.clone(),
        };
        let (client, server) = tokio::io::duplex(1024);
        let _server_handle = tokio::spawn(async move {
            tonic::transport::Server::builder()
                .add_service(
                    ClusterServer::new(range_server)
                        .accept_compressed(tonic::codec::CompressionEncoding::Zstd)
                        .send_compressed(tonic::codec::CompressionEncoding::Zstd),
                )
                .serve_with_incoming(futures::stream::iter([Ok::<_, std::io::Error>(server)]))
                .await
        });

        let channel_manager = ChannelManager::new();
        let mut client = Some(client);
        channel_manager
            .reset_with_connector(
                "leader:0",
                service_fn(move |_| {
                    let client = client.take();
                    async move {
                        client
                            .map(TokioIo::new)
                            .ok_or_else(|| std::io::Error::other("client already taken"))
                    }
                }),
            )
            .unwrap();

        let follower = MetaPeerClientBuilder::default()
            .election(Some(Arc::new(FollowerElection {
                leader_addr: "leader:0".to_string(),
            })))
            .in_memory(Arc::new(MemoryKvBackend::new()))
            .channel_manager(channel_manager)
            .max_retry_count(1)
            .build()
            .unwrap();

        let response = follower
            .range(RangeRequest {
                key: b"key".to_vec(),
                limit: 1,
                ..Default::default()
            })
            .await
            .unwrap();

        assert_eq!(1, requested_limit.load(Ordering::Relaxed));
        assert_eq!(1, response.kvs.len());
        assert!(response.more);
    }

    #[test]
    fn test_to_stat_kv_map() {
        let stat_key = DatanodeStatKey { node_id: 100 };

        let stat = Stat {
            id: 100,
            addr: "127.0.0.1:3001".to_string(),
            ..Default::default()
        };
        let stat_val = DatanodeStatValue { stats: vec![stat] }.try_into().unwrap();

        let kv = KeyValue {
            key: stat_key.into(),
            value: stat_val,
        };

        let kv_map = to_stat_kv_map(vec![kv]).unwrap();
        assert_eq!(1, kv_map.len());
        let _ = kv_map.get(&stat_key).unwrap();

        let stat_val = kv_map.get(&stat_key).unwrap();
        let stat = stat_val.stats.first().unwrap();

        assert_eq!(100, stat.id);
        assert_eq!("127.0.0.1:3001", stat.addr);
    }

    #[test]
    fn test_check_resp_header() {
        let header = Some(ResponseHeader {
            error: None,
            ..Default::default()
        });
        check_resp_header(&header, mock_ctx()).unwrap();

        let result = check_resp_header(&None, mock_ctx());
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            error::Error::ResponseHeaderNotFound { .. }
        ));

        let header = Some(ResponseHeader {
            error: Some(Error {
                code: ErrorCode::NotLeader as i32,
                err_msg: "The current meta is not leader".to_string(),
            }),
            ..Default::default()
        });
        let result = check_resp_header(&header, mock_ctx());
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            error::Error::IsNotLeader { .. }
        ));
    }

    fn mock_ctx<'a>() -> Context<'a> {
        Context { addr: "addr" }
    }
}