servers/grpc/
region_server.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::sync::Arc;
16
17use api::v1::region::region_server::Region as RegionServer;
18use api::v1::region::{RegionRequest, RegionResponse, region_request};
19use async_trait::async_trait;
20use common_error::ext::ErrorExt;
21use common_runtime::Runtime;
22use common_runtime::runtime::RuntimeTrait;
23use common_telemetry::tracing::info_span;
24use common_telemetry::tracing_context::{FutureExt, TracingContext};
25use common_telemetry::{debug, error, warn};
26use snafu::{OptionExt, ResultExt};
27use tonic::{Request, Response, Status};
28
29use crate::error::{InvalidQuerySnafu, JoinTaskSnafu, Result};
30use crate::grpc::{TonicResult, cancellation};
31
32#[async_trait]
33pub trait RegionServerHandler: Send + Sync {
34    async fn handle(&self, request: region_request::Body) -> Result<RegionResponse>;
35
36    async fn handle_region_request(&self, request: RegionRequest) -> Result<RegionResponse> {
37        let body = request.body.context(InvalidQuerySnafu {
38            reason: "Expecting non-empty region request body.",
39        })?;
40        self.handle(body).await
41    }
42}
43
44pub type RegionServerHandlerRef = Arc<dyn RegionServerHandler>;
45
46#[derive(Clone)]
47pub struct RegionServerRequestHandler {
48    handler: Arc<dyn RegionServerHandler>,
49    runtime: Runtime,
50}
51
52impl RegionServerRequestHandler {
53    pub fn new(handler: Arc<dyn RegionServerHandler>, runtime: Runtime) -> Self {
54        Self { handler, runtime }
55    }
56
57    async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
58        let tracing_context = TracingContext::from_w3c(
59            &request
60                .header
61                .as_ref()
62                .context(InvalidQuerySnafu {
63                    reason: "Expecting non-empty region request header.",
64                })?
65                .tracing_context,
66        );
67        let handler = self.handler.clone();
68
69        // Executes requests in another runtime to
70        // 1. prevent the execution from being cancelled unexpected by Tonic runtime;
71        //   - Refer to our blog for the rational behind it:
72        //     https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html
73        //   - Obtaining a `JoinHandle` to get the panic message (if there's any).
74        //     From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped.
75        // 2. avoid the handler blocks the gRPC runtime incidentally.
76        let handle = self.runtime.spawn(async move {
77            handler
78                .handle_region_request(request)
79                .trace(tracing_context.attach(info_span!("RegionServerRequestHandler::handle")))
80                .await
81                .map_err(|e| {
82                    if e.status_code().should_log_error() {
83                        error!(e; "Failed to handle request");
84                    } else {
85                        // Currently, we still print a debug log.
86                        debug!("Failed to handle request, err: {}", e);
87                    }
88                    e
89                })
90        });
91
92        handle.await.context(JoinTaskSnafu)?
93    }
94}
95
96#[async_trait]
97impl RegionServer for RegionServerRequestHandler {
98    async fn handle(
99        &self,
100        request: Request<RegionRequest>,
101    ) -> TonicResult<Response<RegionResponse>> {
102        let remote_addr = request.remote_addr();
103        let self_cloned = self.clone();
104        let request_future = async move {
105            let request = request.into_inner();
106            let response = self_cloned.handle(request).await?;
107
108            Ok(Response::new(response))
109        };
110
111        let cancellation_future = async move {
112            warn!("Region request from {:?} cancelled by client", remote_addr);
113            // If this future is executed it means the request future was dropped,
114            // so it doesn't actually matter what is returned here
115            Err(Status::cancelled("Region request cancelled by client"))
116        };
117        cancellation::with_cancellation_handler(request_future, cancellation_future).await
118    }
119}
servers/grpc/region_server.rs

servers/grpc/
region_server.rs
