servers/

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

pub mod builder;
mod cancellation;
pub mod context_auth;
mod database;
pub mod flight;
pub mod frontend_grpc_handler;
pub mod greptime_handler;
pub mod memory_limit;
pub mod prom_query_gateway;
pub mod region_server;

use std::any::Any;
use std::net::{IpAddr, SocketAddr};
use std::time::Duration;

use api::v1::health_check_server::{HealthCheck, HealthCheckServer};
use api::v1::{HealthCheckRequest, HealthCheckResponse};
use async_trait::async_trait;
use common_base::readable_size::ReadableSize;
use common_grpc::channel_manager::{
    DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE, DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
};
use common_telemetry::{error, info, warn};
use futures::FutureExt;
use http::{HeaderName, Method};
use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsServiceServer;
use serde::{Deserialize, Serialize};
use snafu::{OptionExt, ResultExt, ensure};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio::sync::oneshot::{self, Receiver, Sender};
use tonic::service::Routes;
use tonic::service::interceptor::InterceptedService;
use tonic::transport::ServerTlsConfig;
use tonic::transport::server::TcpIncoming;
use tonic::{Request, Response, Status};
use tonic_reflection::server::v1::{ServerReflection, ServerReflectionServer};
use tower_http::cors::{AllowHeaders, CorsLayer};

use crate::error::{AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, TcpBindSnafu};
use crate::grpc::memory_limit::MemoryLimiterExtensionService;
use crate::http::cors_allow_origin;
use crate::install_default_crypto_provider;
use crate::metrics::MetricsMiddlewareLayer;
use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler};
use crate::query_handler::OpenTelemetryProtocolHandlerRef;
use crate::server::Server;
use crate::tls::TlsOption;

type TonicResult<T> = std::result::Result<T, Status>;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct GrpcOptions {
    /// The address to bind the gRPC server.
    pub bind_addr: String,
    /// The address to advertise to clients.
    pub server_addr: String,
    /// Max gRPC receiving(decoding) message size
    pub max_recv_message_size: ReadableSize,
    /// Max gRPC sending(encoding) message size
    pub max_send_message_size: ReadableSize,
    /// Compression mode in Arrow Flight service.
    pub flight_compression: FlightCompression,
    pub runtime_size: usize,
    #[serde(default = "Default::default")]
    pub tls: TlsOption,
    /// Maximum time that a channel may exist.
    /// Useful when the server wants to control the reconnection of its clients.
    /// Default to `None`, means infinite.
    #[serde(with = "humantime_serde")]
    pub max_connection_age: Option<Duration>,
    /// The HTTP/2 keep-alive interval.
    #[serde(with = "humantime_serde")]
    pub http2_keep_alive_interval: Duration,
    /// The HTTP/2 keep-alive timeout.
    #[serde(with = "humantime_serde")]
    pub http2_keep_alive_timeout: Duration,
    /// Whether to enable CORS, required by gRPC-Web clients in browsers.
    /// Only the frontend's public gRPC server honors it.
    pub enable_cors: bool,
    /// Origins allowed by CORS. Empty allows any origin.
    pub cors_allowed_origins: Vec<String>,
}

impl GrpcOptions {
    /// Detect the server address.
    #[cfg(not(target_os = "android"))]
    pub fn detect_server_addr(&mut self) {
        if self.server_addr.is_empty() {
            match local_ip_address::local_ip() {
                Ok(ip) => {
                    let port = port_from_bind_addr(&self.bind_addr);
                    let detected_addr = format_server_addr(ip, port);
                    info!("Using detected: {} as server address", detected_addr);
                    self.server_addr = detected_addr;
                }
                Err(e) => {
                    error!("Failed to detect local ip address: {}", e);
                }
            }
        }
    }

    #[cfg(target_os = "android")]
    pub fn detect_server_addr(&mut self) {
        if self.server_addr.is_empty() {
            common_telemetry::debug!("detect local IP is not supported on Android");
        }
    }

    /// Create a [GrpcServerConfig] from self's options.
    pub fn as_config(&self) -> GrpcServerConfig {
        GrpcServerConfig {
            max_recv_message_size: self.max_recv_message_size.as_bytes() as usize,
            max_send_message_size: self.max_send_message_size.as_bytes() as usize,
            tls: self.tls.clone(),
            max_connection_age: self.max_connection_age,
        }
    }
}

const DEFAULT_GRPC_ADDR_PORT: u16 = 4001;

fn port_from_bind_addr(bind_addr: &str) -> u16 {
    bind_addr
        .rsplit_once(':')
        .and_then(|(_, port)| port.parse().ok())
        .unwrap_or(DEFAULT_GRPC_ADDR_PORT)
}

fn format_server_addr(ip: IpAddr, port: u16) -> String {
    SocketAddr::new(ip, port).to_string()
}

const DEFAULT_INTERNAL_GRPC_ADDR_PORT: &str = "4010";

impl Default for GrpcOptions {
    fn default() -> Self {
        Self {
            bind_addr: format!("127.0.0.1:{}", DEFAULT_GRPC_ADDR_PORT),
            // If hostname is not set, the server will use the local ip address as the hostname.
            server_addr: String::new(),
            max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
            max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
            flight_compression: FlightCompression::ArrowIpc,
            runtime_size: 8,
            tls: TlsOption::default(),
            max_connection_age: None,
            http2_keep_alive_interval: Duration::from_secs(10),
            http2_keep_alive_timeout: Duration::from_secs(3),
            enable_cors: false,
            cors_allowed_origins: Vec::new(),
        }
    }
}

impl GrpcOptions {
    /// Default options for internal gRPC server.
    /// The internal gRPC server is used for communication between different nodes in cluster.
    /// It is not exposed to the outside world.
    pub fn internal_default() -> Self {
        Self {
            bind_addr: format!("127.0.0.1:{}", DEFAULT_INTERNAL_GRPC_ADDR_PORT),
            // If hostname is not set, the server will use the local ip address as the hostname.
            server_addr: format!("127.0.0.1:{}", DEFAULT_INTERNAL_GRPC_ADDR_PORT),
            max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
            max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
            flight_compression: FlightCompression::ArrowIpc,
            runtime_size: 8,
            tls: TlsOption::default(),
            max_connection_age: None,
            http2_keep_alive_interval: Duration::from_secs(10),
            http2_keep_alive_timeout: Duration::from_secs(3),
            enable_cors: false,
            cors_allowed_origins: Vec::new(),
        }
    }

    pub fn with_bind_addr(mut self, bind_addr: &str) -> Self {
        self.bind_addr = bind_addr.to_string();
        self
    }

    pub fn with_server_addr(mut self, server_addr: &str) -> Self {
        self.server_addr = server_addr.to_string();
        self
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Default)]
#[serde(rename_all = "snake_case")]
pub enum FlightCompression {
    /// Disable all compression in Arrow Flight service.
    #[default]
    None,
    /// Enable only transport layer compression (zstd).
    Transport,
    /// Enable only payload compression (lz4)
    ArrowIpc,
    /// Enable all compression.
    All,
}

impl FlightCompression {
    pub fn transport_compression(&self) -> bool {
        self == &FlightCompression::Transport || self == &FlightCompression::All
    }

    pub fn arrow_compression(&self) -> bool {
        self == &FlightCompression::ArrowIpc || self == &FlightCompression::All
    }
}

/// The wrapped OTLP Arrow service type used by [`GrpcServer`].
type OtelArrowService = MemoryLimiterExtensionService<
    InterceptedService<
        ArrowMetricsServiceServer<OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>>,
        HeaderInterceptor,
    >,
>;

pub struct GrpcServer {
    // states
    shutdown_tx: Mutex<Option<Sender<()>>>,
    /// gRPC serving state receiver. Only present if the gRPC server is started.
    /// Used to wait for the server to stop, performing the old blocking fashion.
    serve_state: Mutex<Option<Receiver<Result<()>>>>,
    // handlers
    routes: Mutex<Option<Routes>>,
    // tls config
    tls_config: Option<ServerTlsConfig>,
    // Otel arrow service
    otel_arrow_service: Mutex<Option<OtelArrowService>>,
    bind_addr: Option<SocketAddr>,
    name: Option<String>,
    config: GrpcServerConfig,
    /// `None` disables CORS.
    cors_allowed_origins: Option<Vec<String>>,
}

/// Grpc Server configuration
#[derive(Debug, Clone)]
pub struct GrpcServerConfig {
    // Max gRPC receiving(decoding) message size
    pub max_recv_message_size: usize,
    // Max gRPC sending(encoding) message size
    pub max_send_message_size: usize,
    pub tls: TlsOption,
    /// Maximum time that a channel may exist.
    /// Useful when the server wants to control the reconnection of its clients.
    /// Default to `None`, means infinite.
    pub max_connection_age: Option<Duration>,
}

impl Default for GrpcServerConfig {
    fn default() -> Self {
        Self {
            max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE.as_bytes() as usize,
            max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE.as_bytes() as usize,
            tls: TlsOption::default(),
            max_connection_age: None,
        }
    }
}

impl GrpcServer {
    fn cors_layer(&self) -> Result<Option<CorsLayer>> {
        let Some(allowed_origins) = &self.cors_allowed_origins else {
            return Ok(None);
        };
        let layer = CorsLayer::new()
            .allow_methods([Method::POST])
            .allow_origin(cors_allow_origin(allowed_origins)?)
            .allow_headers(AllowHeaders::any())
            // Trailers-only responses (most errors) carry `grpc-status` in HTTP
            // headers, which browsers hide unless exposed.
            .expose_headers([
                HeaderName::from_static("grpc-status"),
                HeaderName::from_static("grpc-message"),
                HeaderName::from_static("grpc-status-details-bin"),
            ]);
        Ok(Some(layer))
    }

    pub fn create_healthcheck_service(&self) -> HealthCheckServer<impl HealthCheck> {
        HealthCheckServer::new(HealthCheckHandler)
    }

    pub fn create_reflection_service(&self) -> ServerReflectionServer<impl ServerReflection> {
        tonic_reflection::server::Builder::configure()
            .register_encoded_file_descriptor_set(api::v1::GREPTIME_GRPC_DESC)
            .with_service_name("greptime.v1.GreptimeDatabase")
            .with_service_name("greptime.v1.HealthCheck")
            .with_service_name("greptime.v1.RegionServer")
            .build_v1()
            .inspect_err(|e| {
                common_telemetry::error!(e; "Failed to build gRPC reflection server");
            })
            .unwrap()
    }

    pub async fn wait_for_serve(&self) -> Result<()> {
        let mut serve_state = self.serve_state.lock().await;
        let rx = serve_state.take().context(InternalSnafu {
            err_msg: "gRPC serving state is unknown, maybe the server is not started, \
                      or we have already waited for the serve result before.",
        })?;
        let Ok(result) = rx.await else {
            warn!("Background gRPC serving task is quited before we can receive the serve result.");
            return Ok(());
        };
        if let Err(e) = result {
            error!(e; "GRPC serve error");
        }
        Ok(())
    }
}

pub struct HealthCheckHandler;

#[async_trait]
impl HealthCheck for HealthCheckHandler {
    async fn health_check(
        &self,
        _req: Request<HealthCheckRequest>,
    ) -> TonicResult<Response<HealthCheckResponse>> {
        Ok(Response::new(HealthCheckResponse {}))
    }
}

pub const GRPC_SERVER: &str = "GRPC_SERVER";

#[async_trait]
impl Server for GrpcServer {
    async fn shutdown(&self) -> Result<()> {
        let mut shutdown_tx = self.shutdown_tx.lock().await;
        if let Some(tx) = shutdown_tx.take()
            && tx.send(()).is_err()
        {
            info!("Receiver dropped, the grpc server has already exited");
        }
        info!("Shutdown grpc server");

        Ok(())
    }

    async fn start(&mut self, addr: SocketAddr) -> Result<()> {
        let routes = {
            let mut routes = self.routes.lock().await;
            let Some(routes) = routes.take() else {
                return AlreadyStartedSnafu {
                    server: self.name(),
                }
                .fail();
            };
            routes
        };

        let (tx, rx) = oneshot::channel();
        let (incoming, addr) = {
            let mut shutdown_tx = self.shutdown_tx.lock().await;
            ensure!(
                shutdown_tx.is_none(),
                AlreadyStartedSnafu { server: "gRPC" }
            );

            let listener = TcpListener::bind(addr)
                .await
                .context(TcpBindSnafu { addr })?;
            let addr = listener.local_addr().context(TcpBindSnafu { addr })?;
            let incoming = TcpIncoming::from(listener).with_nodelay(Some(true));
            info!("gRPC server(name={}) is bound to {}", self.name(), addr);

            *shutdown_tx = Some(tx);

            (incoming, addr)
        };

        let cors_layer = self.cors_layer()?;

        // `Server::layer` prepends, so the first layer added is outermost. CORS must
        // sit outside `GrpcWebLayer`, which rejects the `OPTIONS` preflight with a 400.
        let middleware_layer = tower::ServiceBuilder::new()
            .layer(MetricsMiddlewareLayer)
            .option_layer(cors_layer)
            .into_inner();

        let mut builder = tonic::transport::Server::builder()
            .accept_http1(true)
            .layer(middleware_layer)
            .layer(tonic_web::GrpcWebLayer::new());

        if let Some(tls_config) = self.tls_config.clone() {
            // tonic builds the underlying rustls server config here, which requires a
            // process-level crypto provider to be installed first.
            if let Err(err) = install_default_crypto_provider() {
                warn!("Failed to install default rustls crypto provider: {err}");
            }
            builder = builder.tls_config(tls_config).context(StartGrpcSnafu)?;
        }

        if let Some(max_connection_age) = self.config.max_connection_age {
            builder = builder.max_connection_age(max_connection_age);
        }

        let mut builder = builder
            .add_routes(routes)
            .add_service(self.create_healthcheck_service())
            .add_service(self.create_reflection_service());

        if let Some(otel_arrow_service) = self.otel_arrow_service.lock().await.take() {
            builder = builder.add_service(otel_arrow_service);
        }

        let (serve_state_tx, serve_state_rx) = oneshot::channel();
        let mut serve_state = self.serve_state.lock().await;
        *serve_state = Some(serve_state_rx);

        let _handle = common_runtime::spawn_global(async move {
            let result = builder
                .serve_with_incoming_shutdown(incoming, rx.map(drop))
                .await
                .context(StartGrpcSnafu);
            serve_state_tx.send(result)
        });

        self.bind_addr = Some(addr);
        Ok(())
    }

    fn name(&self) -> &str {
        if let Some(name) = &self.name {
            name
        } else {
            GRPC_SERVER
        }
    }

    fn bind_addr(&self) -> Option<SocketAddr> {
        self.bind_addr
    }

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

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};

    use common_runtime::Runtime;
    use common_runtime::runtime::{BuilderBuild, RuntimeTrait};
    use http::header::{
        ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN,
        ACCESS_CONTROL_EXPOSE_HEADERS, ACCESS_CONTROL_REQUEST_HEADERS,
        ACCESS_CONTROL_REQUEST_METHOD, CONTENT_TYPE, ORIGIN,
    };
    use http::{HeaderName, Method, StatusCode};

    use super::{
        DEFAULT_GRPC_ADDR_PORT, GrpcServer, GrpcServerConfig, format_server_addr,
        port_from_bind_addr,
    };
    use crate::grpc::builder::GrpcServerBuilder;
    use crate::server::Server;

    #[test]
    fn test_port_from_bind_addr() {
        assert_eq!(3002, port_from_bind_addr("127.0.0.1:3002"));
        assert_eq!(3002, port_from_bind_addr("[::]:3002"));
        assert_eq!(
            3002,
            port_from_bind_addr("greptimedb-metasrv.default.svc.cluster.local:3002")
        );
        assert_eq!(
            DEFAULT_GRPC_ADDR_PORT,
            port_from_bind_addr("invalid-bind-addr")
        );
    }

    #[test]
    fn test_format_server_addr() {
        assert_eq!(
            "127.0.0.1:3002",
            format_server_addr(IpAddr::V4(Ipv4Addr::LOCALHOST), 3002)
        );
        assert_eq!(
            "[::1]:3002",
            format_server_addr(IpAddr::V6(Ipv6Addr::LOCALHOST), 3002)
        );
    }

    /// The returned server owns the shutdown sender and must outlive the requests.
    async fn start_test_grpc_server(
        cors_allowed_origins: Option<Vec<String>>,
    ) -> (GrpcServer, SocketAddr) {
        let runtime = Runtime::builder().build().unwrap();
        let mut builder = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime);
        if let Some(origins) = cors_allowed_origins {
            builder = builder.with_cors(origins);
        }
        let mut server = builder.build();
        server
            .start(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
            .await
            .unwrap();
        let addr = server.bind_addr().unwrap();
        (server, addr)
    }

    const TEST_PATH: &str = "greptime.v1.HealthCheck/Check";

    async fn send_preflight(addr: SocketAddr, origin: &str) -> reqwest::Response {
        reqwest::Client::new()
            .request(Method::OPTIONS, format!("http://{addr}/{TEST_PATH}"))
            .header(ORIGIN, origin)
            .header(ACCESS_CONTROL_REQUEST_METHOD, "POST")
            .header(ACCESS_CONTROL_REQUEST_HEADERS, "content-type,x-grpc-web")
            .send()
            .await
            .unwrap()
    }

    /// The path is not served, so the response is a trailers-only `Unimplemented`.
    async fn send_grpc_web_call(addr: SocketAddr, origin: &str) -> reqwest::Response {
        reqwest::Client::new()
            .post(format!("http://{addr}/{TEST_PATH}"))
            .header(ORIGIN, origin)
            .header(CONTENT_TYPE, "application/grpc-web+proto")
            .header("x-grpc-web", "1")
            // Empty gRPC-Web frame: 1-byte compression flag + 4-byte length.
            .body(vec![0u8; 5])
            .send()
            .await
            .unwrap()
    }

    fn header(response: &reqwest::Response, name: HeaderName) -> Option<&str> {
        response.headers().get(name).map(|v| v.to_str().unwrap())
    }

    #[tokio::test]
    async fn test_grpc_web_cors_preflight() {
        let (_server, addr) = start_test_grpc_server(Some(Vec::new())).await;

        let response = send_preflight(addr, "https://example.com").await;

        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), Some("*"));
        assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_HEADERS), Some("*"));
        // Lists the method of the actual call, not `OPTIONS`.
        assert_eq!(
            header(&response, ACCESS_CONTROL_ALLOW_METHODS),
            Some("POST")
        );
    }

    #[tokio::test]
    async fn test_grpc_web_cors_expose_headers() {
        let (_server, addr) = start_test_grpc_server(Some(Vec::new())).await;

        let response = send_grpc_web_call(addr, "https://example.com").await;

        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), Some("*"));
        // Premise: the response is trailers-only.
        assert_eq!(
            header(&response, HeaderName::from_static("grpc-status")),
            Some("12")
        );
        let expose_headers = header(&response, ACCESS_CONTROL_EXPOSE_HEADERS).unwrap();
        for name in ["grpc-status", "grpc-message", "grpc-status-details-bin"] {
            assert!(expose_headers.contains(name), "{expose_headers}");
        }
    }

    #[tokio::test]
    async fn test_grpc_web_cors_custom_origins() {
        let (_server, addr) =
            start_test_grpc_server(Some(vec!["https://example.com".to_string()])).await;

        let response = send_preflight(addr, "https://example.com").await;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            header(&response, ACCESS_CONTROL_ALLOW_ORIGIN),
            Some("https://example.com")
        );

        // A disallowed origin still gets a 200, just without allow-origin; the
        // browser does the blocking.
        let response = send_preflight(addr, "https://notallowed.com").await;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), None);
    }

    #[tokio::test]
    async fn test_grpc_web_cors_disabled() {
        let (_server, addr) = start_test_grpc_server(None).await;

        let response = send_preflight(addr, "https://example.com").await;

        // Without CORS, `GrpcWebLayer` rejects the preflight.
        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
        assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), None);
    }
}