flow/batching_mode/

frontend_client.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.

//! Frontend client to run flow as batching task which is time-window-aware normal query triggered every tick set by user

use std::collections::HashMap;
use std::sync::{Arc, Mutex, Weak};

use api::v1::greptime_request::Request;
use api::v1::query_request::Query;
use api::v1::{CreateTableExpr, QueryRequest};
use client::{Client, Database};
use common_error::ext::BoxedError;
use common_grpc::channel_manager::{ChannelConfig, ChannelManager, load_client_tls_config};
use common_meta::peer::{Peer, PeerDiscovery};
use common_query::Output;
use common_telemetry::warn;
use meta_client::client::MetaClient;
use query::datafusion::QUERY_PARALLELISM_HINT;
use query::options::QueryOptions;
use rand::rng;
use rand::seq::SliceRandom;
use servers::query_handler::grpc::GrpcQueryHandler;
use session::context::{QueryContextBuilder, QueryContextRef};
use session::hints::READ_PREFERENCE_HINT;
use snafu::{OptionExt, ResultExt};
use tokio::sync::SetOnce;

use crate::batching_mode::BatchingModeOptions;
use crate::error::{
    CreateSinkTableSnafu, ExternalSnafu, InvalidClientConfigSnafu, InvalidRequestSnafu,
    NoAvailableFrontendSnafu, UnexpectedSnafu,
};
use crate::{Error, FlowAuthHeader};

/// Adapter trait for [`GrpcQueryHandler`] that boxes the underlying error into [`BoxedError`].
///
/// This is mainly used by flownode to invoke a frontend instance in standalone mode.
#[async_trait::async_trait]
pub trait GrpcQueryHandlerWithBoxedError: Send + Sync + 'static {
    async fn do_query(
        &self,
        query: Request,
        ctx: QueryContextRef,
    ) -> std::result::Result<Output, BoxedError>;
}

/// auto impl
#[async_trait::async_trait]
impl<T: GrpcQueryHandler + Send + Sync + 'static> GrpcQueryHandlerWithBoxedError for T {
    async fn do_query(
        &self,
        query: Request,
        ctx: QueryContextRef,
    ) -> std::result::Result<Output, BoxedError> {
        self.do_query(query, ctx).await.map_err(BoxedError::new)
    }
}

#[derive(Debug, Clone)]
pub struct HandlerMutable {
    handler: Arc<Mutex<Option<Weak<dyn GrpcQueryHandlerWithBoxedError>>>>,
    is_initialized: Arc<SetOnce<()>>,
}

impl HandlerMutable {
    pub async fn set_handler(&self, handler: Weak<dyn GrpcQueryHandlerWithBoxedError>) {
        *self.handler.lock().unwrap() = Some(handler);
        // Ignore the error, as we allow the handler to be set multiple times.
        let _ = self.is_initialized.set(());
    }
}

/// A simple frontend client able to execute sql using grpc protocol
///
/// This is for computation-heavy query which need to offload computation to frontend, lifting the load from flownode
#[derive(Debug, Clone)]
pub enum FrontendClient {
    Distributed {
        meta_client: Arc<MetaClient>,
        chnl_mgr: ChannelManager,
        auth: Option<FlowAuthHeader>,
        query: QueryOptions,
        batch_opts: BatchingModeOptions,
    },
    Standalone {
        /// for the sake of simplicity still use grpc even in standalone mode
        /// notice the client here should all be lazy, so that can wait after frontend is booted then make conn
        database_client: HandlerMutable,
        query: QueryOptions,
    },
}

impl FrontendClient {
    /// Create a new empty frontend client, with a `HandlerMutable` to set the grpc handler later
    pub fn from_empty_grpc_handler(query: QueryOptions) -> (Self, HandlerMutable) {
        let is_initialized = Arc::new(SetOnce::new());
        let handler = HandlerMutable {
            handler: Arc::new(Mutex::new(None)),
            is_initialized,
        };
        (
            Self::Standalone {
                database_client: handler.clone(),
                query,
            },
            handler,
        )
    }

    /// Waits until the frontend client is initialized.
    pub async fn wait_initialized(&self) {
        if let FrontendClient::Standalone {
            database_client, ..
        } = self
        {
            database_client.is_initialized.wait().await;
        }
    }

    pub fn from_meta_client(
        meta_client: Arc<MetaClient>,
        auth: Option<FlowAuthHeader>,
        query: QueryOptions,
        batch_opts: BatchingModeOptions,
    ) -> Result<Self, Error> {
        common_telemetry::info!("Frontend client build with auth={:?}", auth);
        Ok(Self::Distributed {
            meta_client,
            chnl_mgr: {
                let cfg = ChannelConfig::new()
                    .connect_timeout(batch_opts.grpc_conn_timeout)
                    .timeout(Some(batch_opts.query_timeout));

                let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone())
                    .context(InvalidClientConfigSnafu)?;
                ChannelManager::with_config(cfg, tls_config)
            },
            auth,
            query,
            batch_opts,
        })
    }

    pub fn from_grpc_handler(
        grpc_handler: Weak<dyn GrpcQueryHandlerWithBoxedError>,
        query: QueryOptions,
    ) -> Self {
        let is_initialized = Arc::new(SetOnce::new_with(Some(())));
        let handler = HandlerMutable {
            handler: Arc::new(Mutex::new(Some(grpc_handler))),
            is_initialized: is_initialized.clone(),
        };

        Self::Standalone {
            database_client: handler,
            query,
        }
    }
}

#[derive(Debug, Clone)]
pub struct DatabaseWithPeer {
    pub database: Database,
    pub peer: Peer,
}

impl DatabaseWithPeer {
    fn new(database: Database, peer: Peer) -> Self {
        Self { database, peer }
    }

    /// Try sending a "SELECT 1" to the database
    async fn try_select_one(&self) -> Result<(), Error> {
        // notice here use `sql` for `SELECT 1` return 1 row
        let _ = self
            .database
            .sql("SELECT 1")
            .await
            .with_context(|_| InvalidRequestSnafu {
                context: format!("Failed to handle `SELECT 1` request at {:?}", self.peer),
            })?;
        Ok(())
    }
}

impl FrontendClient {
    // TODO: support more fine-grained load balancing strategies for frontend
    // selection, such as AZ (availability zone) awareness, to prefer frontends
    // in the same zone as the flownode and reduce cross-AZ latency.
    /// scan for available frontend from metadata
    pub(crate) async fn scan_for_frontend(&self) -> Result<Vec<Peer>, Error> {
        let Self::Distributed { meta_client, .. } = self else {
            return Ok(vec![]);
        };

        meta_client
            .active_frontends()
            .await
            .map(|nodes| nodes.into_iter().map(|node| node.peer).collect())
            .map_err(BoxedError::new)
            .context(ExternalSnafu)
    }

    /// Get a frontend discovered by metasrv and verified with a query probe.
    async fn get_random_active_frontend(
        &self,
        catalog: &str,
        schema: &str,
    ) -> Result<DatabaseWithPeer, Error> {
        let Self::Distributed {
            meta_client: _,
            chnl_mgr,
            auth,
            query: _,
            batch_opts,
        } = self
        else {
            return UnexpectedSnafu {
                reason: "Expect distributed mode",
            }
            .fail();
        };

        let mut interval = tokio::time::interval(batch_opts.grpc_conn_timeout);
        interval.tick().await;
        for retry in 0..batch_opts.experimental_grpc_max_retries {
            let mut frontends = self.scan_for_frontend().await?;
            // shuffle the frontends to avoid always pick the same one
            frontends.shuffle(&mut rng());

            for peer in frontends {
                let addr = peer.addr.clone();
                let client = Client::with_manager_and_urls(chnl_mgr.clone(), vec![addr.clone()]);
                let database = {
                    let mut db = Database::new(catalog, schema, client);
                    if let Some(auth) = auth {
                        db.set_auth(auth.auth().clone());
                    }
                    db
                };
                let db = DatabaseWithPeer::new(database, peer);
                match db.try_select_one().await {
                    Ok(_) => return Ok(db),
                    Err(e) => {
                        warn!(
                            "Failed to connect to frontend {} on retry={}: \n{e:?}",
                            addr, retry
                        );
                    }
                }
            }
            // no available frontend
            // sleep and retry
            interval.tick().await;
        }

        NoAvailableFrontendSnafu {
            timeout: batch_opts.grpc_conn_timeout,
            context: "No available frontend found that is able to process query",
        }
        .fail()
    }

    pub async fn create(
        &self,
        create: CreateTableExpr,
        catalog: &str,
        schema: &str,
    ) -> Result<u32, Error> {
        self.handle(
            Request::Ddl(api::v1::DdlRequest {
                expr: Some(api::v1::ddl_request::Expr::CreateTable(create.clone())),
            }),
            catalog,
            schema,
            &mut None,
        )
        .await
        .map_err(BoxedError::new)
        .with_context(|_| CreateSinkTableSnafu {
            create: create.clone(),
        })
    }

    /// Execute a SQL statement on the frontend.
    pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result<Output, Error> {
        match self {
            FrontendClient::Distributed { .. } => {
                let db = self.get_random_active_frontend(catalog, schema).await?;
                db.database
                    .sql(sql)
                    .await
                    .map_err(BoxedError::new)
                    .context(ExternalSnafu)
            }
            FrontendClient::Standalone {
                database_client, ..
            } => {
                let ctx = QueryContextBuilder::default()
                    .current_catalog(catalog.to_string())
                    .current_schema(schema.to_string())
                    .build();
                let ctx = Arc::new(ctx);
                {
                    let database_client = {
                        database_client
                            .handler
                            .lock()
                            .unwrap()
                            .as_ref()
                            .context(UnexpectedSnafu {
                                reason: "Standalone's frontend instance is not set",
                            })?
                            .upgrade()
                            .context(UnexpectedSnafu {
                                reason: "Failed to upgrade database client",
                            })?
                    };
                    let req = Request::Query(QueryRequest {
                        query: Some(Query::Sql(sql.to_string())),
                    });
                    database_client
                        .do_query(req, ctx)
                        .await
                        .map_err(BoxedError::new)
                        .context(ExternalSnafu)
                }
            }
        }
    }

    /// Handle a request to frontend
    pub(crate) async fn handle(
        &self,
        req: api::v1::greptime_request::Request,
        catalog: &str,
        schema: &str,
        peer_desc: &mut Option<PeerDesc>,
    ) -> Result<u32, Error> {
        match self {
            FrontendClient::Distributed {
                query, batch_opts, ..
            } => {
                let db = self.get_random_active_frontend(catalog, schema).await?;

                *peer_desc = Some(PeerDesc::Dist {
                    peer: db.peer.clone(),
                });

                db.database
                    .handle_with_retry(
                        req.clone(),
                        batch_opts.experimental_grpc_max_retries,
                        &[
                            (QUERY_PARALLELISM_HINT, &query.parallelism.to_string()),
                            (READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref()),
                        ],
                    )
                    .await
                    .with_context(|_| InvalidRequestSnafu {
                        context: format!("Failed to handle request at {:?}: {:?}", db.peer, req),
                    })
            }
            FrontendClient::Standalone {
                database_client,
                query,
            } => {
                let ctx = QueryContextBuilder::default()
                    .current_catalog(catalog.to_string())
                    .current_schema(schema.to_string())
                    .extensions(HashMap::from([(
                        QUERY_PARALLELISM_HINT.to_string(),
                        query.parallelism.to_string(),
                    )]))
                    .build();
                let ctx = Arc::new(ctx);
                {
                    let database_client = {
                        database_client
                            .handler
                            .lock()
                            .unwrap()
                            .as_ref()
                            .context(UnexpectedSnafu {
                                reason: "Standalone's frontend instance is not set",
                            })?
                            .upgrade()
                            .context(UnexpectedSnafu {
                                reason: "Failed to upgrade database client",
                            })?
                    };
                    let resp: common_query::Output = database_client
                        .do_query(req, ctx)
                        .await
                        .map_err(BoxedError::new)
                        .context(ExternalSnafu)?;
                    match resp.data {
                        common_query::OutputData::AffectedRows(rows) => {
                            Ok(rows.try_into().map_err(|_| {
                                UnexpectedSnafu {
                                    reason: format!("Failed to convert rows to u32: {}", rows),
                                }
                                .build()
                            })?)
                        }
                        _ => UnexpectedSnafu {
                            reason: "Unexpected output data",
                        }
                        .fail(),
                    }
                }
            }
        }
    }
}

/// Describe a peer of frontend
#[derive(Debug, Default)]
pub(crate) enum PeerDesc {
    /// Distributed mode's frontend peer address
    Dist {
        /// frontend peer address
        peer: Peer,
    },
    /// Standalone mode
    #[default]
    Standalone,
}

impl std::fmt::Display for PeerDesc {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PeerDesc::Dist { peer } => write!(f, "{}", peer.addr),
            PeerDesc::Standalone => write!(f, "standalone"),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use common_query::Output;
    use tokio::time::timeout;

    use super::*;

    #[derive(Debug)]
    struct NoopHandler;

    #[async_trait::async_trait]
    impl GrpcQueryHandlerWithBoxedError for NoopHandler {
        async fn do_query(
            &self,
            _query: Request,
            _ctx: QueryContextRef,
        ) -> std::result::Result<Output, BoxedError> {
            Ok(Output::new_with_affected_rows(0))
        }
    }

    #[tokio::test]
    async fn wait_initialized() {
        let (client, handler_mut) =
            FrontendClient::from_empty_grpc_handler(QueryOptions::default());

        assert!(
            timeout(Duration::from_millis(50), client.wait_initialized())
                .await
                .is_err()
        );

        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
        handler_mut.set_handler(Arc::downgrade(&handler)).await;

        timeout(Duration::from_secs(1), client.wait_initialized())
            .await
            .expect("wait_initialized should complete after handler is set");

        timeout(Duration::from_millis(10), client.wait_initialized())
            .await
            .expect("wait_initialized should be a no-op once initialized");

        let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(NoopHandler);
        let client =
            FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default());
        assert!(
            timeout(Duration::from_millis(10), client.wait_initialized())
                .await
                .is_ok()
        );

        let meta_client = Arc::new(MetaClient::new(0, api::v1::meta::Role::Frontend));
        let client = FrontendClient::from_meta_client(
            meta_client,
            None,
            QueryOptions::default(),
            BatchingModeOptions::default(),
        )
        .unwrap();
        assert!(
            timeout(Duration::from_millis(10), client.wait_initialized())
                .await
                .is_ok()
        );
    }
}