servers/grpc/
database.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 api::v1::greptime_database_server::GreptimeDatabase;
16use api::v1::greptime_response::Response as RawResponse;
17use api::v1::{AffectedRows, GreptimeRequest, GreptimeResponse, ResponseHeader};
18use async_trait::async_trait;
19use common_error::status_code::StatusCode;
20use common_query::OutputData;
21use common_telemetry::{debug, warn};
22use futures::StreamExt;
23use prost::Message;
24use session::context::Channel;
25use tonic::{Request, Response, Status, Streaming};
26
27use crate::grpc::greptime_handler::GreptimeRequestHandler;
28use crate::grpc::memory_limit::PreDecodeMemoryReservation;
29use crate::grpc::{TonicResult, cancellation};
30use crate::hint_headers;
31use crate::request_memory_limiter::ServerMemoryLimiter;
32
33pub(crate) struct DatabaseService {
34    handler: GreptimeRequestHandler,
35}
36
37impl DatabaseService {
38    pub(crate) fn new(handler: GreptimeRequestHandler) -> Self {
39        Self { handler }
40    }
41}
42
43#[async_trait]
44impl GreptimeDatabase for DatabaseService {
45    async fn handle(
46        &self,
47        request: Request<GreptimeRequest>,
48    ) -> TonicResult<Response<GreptimeResponse>> {
49        let remote_addr = request.remote_addr();
50        let hints = hint_headers::extract_hints(request.metadata());
51        let channel = request
52            .extensions()
53            .get::<Channel>()
54            .copied()
55            .unwrap_or(Channel::Grpc);
56        debug!(
57            "GreptimeDatabase::Handle: request from {:?} with hints: {:?}",
58            remote_addr, hints
59        );
60
61        // Retain the pre-decode reservation for the whole request: the
62        // extension holding the guard would be dropped when the request is
63        // consumed below, but the post-decode charge is skipped while it is
64        // active.
65        let _pre_reservation = request
66            .extensions()
67            .get::<PreDecodeMemoryReservation>()
68            .cloned();
69        let _guard = if _pre_reservation.is_some() {
70            // Compressed requests already reserved the worst-case decoded
71            // size before tonic decompressed the message; skip the exact
72            // post-decode charge to avoid double accounting.
73            None
74        } else if let Some(limiter) = request.extensions().get::<ServerMemoryLimiter>() {
75            let message_size = request.get_ref().encoded_len() as u64;
76            Some(limiter.acquire(message_size).await?)
77        } else {
78            None
79        };
80
81        let handler = self.handler.clone();
82        let request_future = async move {
83            let request = request.into_inner();
84            let output = handler.handle_request(request, hints, channel).await?;
85            let message = match output.data {
86                OutputData::AffectedRows(rows) => GreptimeResponse {
87                    header: Some(ResponseHeader {
88                        status: Some(api::v1::Status {
89                            status_code: StatusCode::Success as _,
90                            ..Default::default()
91                        }),
92                    }),
93                    response: Some(RawResponse::AffectedRows(AffectedRows { value: rows as _ })),
94                },
95                OutputData::Stream(_) | OutputData::RecordBatches(_) => {
96                    return Err(Status::unimplemented("GreptimeDatabase::Handle for query"));
97                }
98            };
99
100            Ok(Response::new(message))
101        };
102
103        let cancellation_future = async move {
104            warn!(
105                "GreptimeDatabase::Handle: request from {:?} cancelled by client",
106                remote_addr
107            );
108            // If this future is executed it means the request future was dropped,
109            // so it doesn't actually matter what is returned here
110            Err(Status::cancelled(
111                "GreptimeDatabase::Handle: request cancelled by client",
112            ))
113        };
114        cancellation::with_cancellation_handler(request_future, cancellation_future).await
115    }
116
117    async fn handle_requests(
118        &self,
119        request: Request<Streaming<GreptimeRequest>>,
120    ) -> Result<Response<GreptimeResponse>, Status> {
121        let remote_addr = request.remote_addr();
122        let hints = hint_headers::extract_hints(request.metadata());
123        let channel = request
124            .extensions()
125            .get::<Channel>()
126            .copied()
127            .unwrap_or(Channel::Grpc);
128        debug!(
129            "GreptimeDatabase::HandleRequests: request from {:?} with hints: {:?}",
130            remote_addr, hints
131        );
132
133        let limiter = request.extensions().get::<ServerMemoryLimiter>().cloned();
134        // For compressed streams the whole stream's decoding memory was
135        // reserved before tonic started decompressing; messages are decoded
136        // one at a time, so the reservation covers each of them. The
137        // reservation is retained below for the whole stream: the extension
138        // holding the guard would otherwise be dropped when the request is
139        // consumed, while per-message charges stay skipped.
140        let reservation = request
141            .extensions()
142            .get::<PreDecodeMemoryReservation>()
143            .cloned();
144        let pre_reserved = reservation.is_some();
145
146        let handler = self.handler.clone();
147        let request_future = async move {
148            let mut affected_rows = 0;
149
150            // Hold the pre-decode reservation until the stream is exhausted.
151            let _reservation = reservation;
152            let mut stream = request.into_inner();
153            while let Some(request) = stream.next().await {
154                let request = request?;
155
156                let _guard = if pre_reserved {
157                    None
158                } else if let Some(limiter_ref) = &limiter {
159                    let message_size = request.encoded_len() as u64;
160                    Some(limiter_ref.acquire(message_size).await?)
161                } else {
162                    None
163                };
164                let output = handler
165                    .handle_request(request, hints.clone(), channel)
166                    .await?;
167                match output.data {
168                    OutputData::AffectedRows(rows) => affected_rows += rows,
169                    OutputData::Stream(_) | OutputData::RecordBatches(_) => {
170                        return Err(Status::unimplemented(
171                            "GreptimeDatabase::HandleRequests for query",
172                        ));
173                    }
174                }
175            }
176            let message = GreptimeResponse {
177                header: Some(ResponseHeader {
178                    status: Some(api::v1::Status {
179                        status_code: StatusCode::Success as _,
180                        ..Default::default()
181                    }),
182                }),
183                response: Some(RawResponse::AffectedRows(AffectedRows {
184                    value: affected_rows as u32,
185                })),
186            };
187
188            Ok(Response::new(message))
189        };
190
191        let cancellation_future = async move {
192            warn!(
193                "GreptimeDatabase::HandleRequests: request from {:?} cancelled by client",
194                remote_addr
195            );
196            // If this future is executed it means the request future was dropped,
197            // so it doesn't actually matter what is returned here
198            Err(Status::cancelled(
199                "GreptimeDatabase::HandleRequests: request cancelled by client",
200            ))
201        };
202        cancellation::with_cancellation_handler(request_future, cancellation_future).await
203    }
204}
servers/grpc/database.rs

servers/grpc/
database.rs
