common_frontend/
selector.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Debug;
16use std::time::Duration;
17
18use common_grpc::channel_manager::{ChannelConfig, ChannelManager};
19use common_meta::peer::{Peer, PeerDiscovery};
20use greptime_proto::v1::frontend::{
21    KillProcessRequest, KillProcessResponse, ListProcessRequest, ListProcessResponse,
22    frontend_client,
23};
24use meta_client::MetaClientRef;
25use snafu::ResultExt;
26use tonic::Response;
27
28use crate::error;
29use crate::error::{MetaSnafu, Result};
30
31pub type FrontendClientPtr = Box<dyn FrontendClient>;
32
33#[async_trait::async_trait]
34pub trait FrontendClient: Send + Debug {
35    async fn list_process(&mut self, req: ListProcessRequest) -> Result<ListProcessResponse>;
36
37    async fn kill_process(&mut self, req: KillProcessRequest) -> Result<KillProcessResponse>;
38}
39
40#[async_trait::async_trait]
41impl FrontendClient for frontend_client::FrontendClient<tonic::transport::channel::Channel> {
42    async fn list_process(&mut self, req: ListProcessRequest) -> Result<ListProcessResponse> {
43        frontend_client::FrontendClient::<tonic::transport::channel::Channel>::list_process(
44            self, req,
45        )
46        .await
47        .context(error::InvokeFrontendSnafu)
48        .map(Response::into_inner)
49    }
50
51    async fn kill_process(&mut self, req: KillProcessRequest) -> Result<KillProcessResponse> {
52        frontend_client::FrontendClient::<tonic::transport::channel::Channel>::kill_process(
53            self, req,
54        )
55        .await
56        .context(error::InvokeFrontendSnafu)
57        .map(Response::into_inner)
58    }
59}
60
61#[async_trait::async_trait]
62pub trait FrontendSelector {
63    async fn select<F>(&self, predicate: F) -> Result<Vec<FrontendClientPtr>>
64    where
65        F: Fn(&Peer) -> bool + Send;
66}
67
68#[derive(Debug, Clone)]
69pub struct MetaClientSelector {
70    meta_client: MetaClientRef,
71    channel_manager: ChannelManager,
72}
73
74#[async_trait::async_trait]
75impl FrontendSelector for MetaClientSelector {
76    async fn select<F>(&self, predicate: F) -> Result<Vec<FrontendClientPtr>>
77    where
78        F: Fn(&Peer) -> bool + Send,
79    {
80        let peers = self
81            .meta_client
82            .active_frontends()
83            .await
84            .map_err(Box::new)
85            .context(MetaSnafu)?;
86
87        peers
88            .into_iter()
89            .filter(predicate)
90            .map(|peer| {
91                let channel = self
92                    .channel_manager
93                    .get(peer.addr)
94                    .map_err(Box::new)
95                    .context(error::CreateChannelSnafu)?;
96                let client = frontend_client::FrontendClient::new(channel);
97                Ok(Box::new(client) as FrontendClientPtr)
98            })
99            .collect::<Result<Vec<_>>>()
100    }
101}
102
103impl MetaClientSelector {
104    pub fn new(meta_client: MetaClientRef) -> Self {
105        let cfg = ChannelConfig::new()
106            .connect_timeout(Duration::from_secs(30))
107            .timeout(Some(Duration::from_secs(30)));
108        let channel_manager = ChannelManager::with_config(cfg, None);
109        Self {
110            meta_client,
111            channel_manager,
112        }
113    }
114}
common_frontend/selector.rs

common_frontend/
selector.rs
