client/
region.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::region::RegionResponse;
18use api::v1::ResponseHeader;
19use api::v1::region::RegionRequest;
20use arc_swap::ArcSwapOption;
21use arrow_flight::Ticket;
22use async_stream::stream;
23use async_trait::async_trait;
24use common_error::ext::BoxedError;
25use common_error::status_code::StatusCode;
26use common_grpc::flight::{FlightDecoder, FlightMessage};
27use common_meta::error::{self as meta_error, Result as MetaResult};
28use common_meta::node_manager::Datanode;
29use common_query::request::QueryRequest;
30use common_recordbatch::error::ExternalSnafu;
31use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, SendableRecordBatchStream};
32use common_telemetry::error;
33use common_telemetry::tracing_context::TracingContext;
34use prost::Message;
35use query::query_engine::DefaultSerializer;
36use snafu::{OptionExt, ResultExt, location};
37use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
38use tokio_stream::StreamExt;
39
40use crate::error::{
41    self, ConvertFlightDataSnafu, FlightGetSnafu, IllegalDatabaseResponseSnafu,
42    IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu,
43};
44use crate::{Client, Error, metrics};
45
46#[derive(Debug)]
47pub struct RegionRequester {
48    client: Client,
49    send_compression: bool,
50    accept_compression: bool,
51}
52
53#[async_trait]
54impl Datanode for RegionRequester {
55    async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
56        self.handle_inner(request).await.map_err(|err| {
57            if err.should_retry() {
58                meta_error::Error::RetryLater {
59                    source: BoxedError::new(err),
60                    clean_poisons: false,
61                }
62            } else {
63                meta_error::Error::External {
64                    source: BoxedError::new(err),
65                    location: location!(),
66                }
67            }
68        })
69    }
70
71    async fn handle_query(&self, request: QueryRequest) -> MetaResult<SendableRecordBatchStream> {
72        let plan = DFLogicalSubstraitConvertor
73            .encode(&request.plan, DefaultSerializer)
74            .map_err(BoxedError::new)
75            .context(meta_error::ExternalSnafu)?
76            .to_vec();
77        let request = api::v1::region::QueryRequest {
78            header: request.header,
79            region_id: request.region_id.as_u64(),
80            plan,
81        };
82
83        let ticket = Ticket {
84            ticket: request.encode_to_vec().into(),
85        };
86        self.do_get_inner(ticket)
87            .await
88            .map_err(BoxedError::new)
89            .context(meta_error::ExternalSnafu)
90    }
91}
92
93impl RegionRequester {
94    pub fn new(client: Client, send_compression: bool, accept_compression: bool) -> Self {
95        Self {
96            client,
97            send_compression,
98            accept_compression,
99        }
100    }
101
102    pub async fn do_get_inner(&self, ticket: Ticket) -> Result<SendableRecordBatchStream> {
103        let mut flight_client = self
104            .client
105            .make_flight_client(self.send_compression, self.accept_compression)?;
106        let response = flight_client
107            .mut_inner()
108            .do_get(ticket)
109            .await
110            .or_else(|e| {
111                let tonic_code = e.code();
112                let e: error::Error = e.into();
113                error!(
114                    e; "Failed to do Flight get, addr: {}, code: {}",
115                    flight_client.addr(),
116                    tonic_code
117                );
118                Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu {
119                    addr: flight_client.addr().to_string(),
120                    tonic_code,
121                })
122            })?;
123
124        let flight_data_stream = response.into_inner();
125        let mut decoder = FlightDecoder::default();
126
127        let mut flight_message_stream = flight_data_stream.map(move |flight_data| {
128            flight_data
129                .map_err(Error::from)
130                .and_then(|data| decoder.try_decode(&data).context(ConvertFlightDataSnafu))?
131                .context(IllegalFlightMessagesSnafu {
132                    reason: "none message",
133                })
134        });
135
136        let Some(first_flight_message) = flight_message_stream.next().await else {
137            return IllegalFlightMessagesSnafu {
138                reason: "Expect the response not to be empty",
139            }
140            .fail();
141        };
142        let FlightMessage::Schema(schema) = first_flight_message? else {
143            return IllegalFlightMessagesSnafu {
144                reason: "Expect schema to be the first flight message",
145            }
146            .fail();
147        };
148
149        let metrics = Arc::new(ArcSwapOption::from(None));
150        let metrics_ref = metrics.clone();
151
152        let tracing_context = TracingContext::from_current_span();
153
154        let schema = Arc::new(
155            datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?,
156        );
157        let schema_cloned = schema.clone();
158        let stream = Box::pin(stream!({
159            let _span = tracing_context.attach(common_telemetry::tracing::info_span!(
160                "poll_flight_data_stream"
161            ));
162
163            let mut buffered_message: Option<FlightMessage> = None;
164            let mut stream_ended = false;
165
166            while !stream_ended {
167                // get the next message from the buffered message or read from the flight message stream
168                let flight_message_item = if let Some(msg) = buffered_message.take() {
169                    Some(Ok(msg))
170                } else {
171                    flight_message_stream.next().await
172                };
173
174                let flight_message = match flight_message_item {
175                    Some(Ok(message)) => message,
176                    Some(Err(e)) => {
177                        yield Err(BoxedError::new(e)).context(ExternalSnafu);
178                        break;
179                    }
180                    None => break,
181                };
182
183                match flight_message {
184                    FlightMessage::RecordBatch(record_batch) => {
185                        let result_to_yield = RecordBatch::try_from_df_record_batch(
186                            schema_cloned.clone(),
187                            record_batch,
188                        );
189
190                        // get the next message from the stream. normally it should be a metrics message.
191                        if let Some(next_flight_message_result) = flight_message_stream.next().await
192                        {
193                            match next_flight_message_result {
194                                Ok(FlightMessage::Metrics(s)) => {
195                                    let m = serde_json::from_str(&s).ok().map(Arc::new);
196                                    metrics_ref.swap(m);
197                                }
198                                Ok(FlightMessage::RecordBatch(rb)) => {
199                                    // for some reason it's not a metrics message, so we need to buffer this record batch
200                                    // and yield it in the next iteration.
201                                    buffered_message = Some(FlightMessage::RecordBatch(rb));
202                                }
203                                Ok(_) => {
204                                    yield IllegalFlightMessagesSnafu {
205                                        reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message"
206                                    }
207                                    .fail()
208                                    .map_err(BoxedError::new)
209                                    .context(ExternalSnafu);
210                                    break;
211                                }
212                                Err(e) => {
213                                    yield Err(BoxedError::new(e)).context(ExternalSnafu);
214                                    break;
215                                }
216                            }
217                        } else {
218                            // the stream has ended
219                            stream_ended = true;
220                        }
221
222                        yield result_to_yield;
223                    }
224                    FlightMessage::Metrics(s) => {
225                        // just a branch in case of some metrics message comes after other things.
226                        let m = serde_json::from_str(&s).ok().map(Arc::new);
227                        metrics_ref.swap(m);
228                        break;
229                    }
230                    _ => {
231                        yield IllegalFlightMessagesSnafu {
232                            reason: "A Schema message must be succeeded exclusively by a set of RecordBatch messages"
233                        }
234                        .fail()
235                        .map_err(BoxedError::new)
236                        .context(ExternalSnafu);
237                        break;
238                    }
239                }
240            }
241        }));
242        let record_batch_stream = RecordBatchStreamWrapper {
243            schema,
244            stream,
245            output_ordering: None,
246            metrics,
247        };
248        Ok(Box::pin(record_batch_stream))
249    }
250
251    async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponse> {
252        let request_type = request
253            .body
254            .as_ref()
255            .with_context(|| MissingFieldSnafu { field: "body" })?
256            .as_ref()
257            .to_string();
258        let _timer = metrics::METRIC_REGION_REQUEST_GRPC
259            .with_label_values(&[request_type.as_str()])
260            .start_timer();
261
262        let (addr, mut client) = self.client.raw_region_client()?;
263
264        let response = client
265            .handle(request)
266            .await
267            .map_err(|e| {
268                let code = e.code();
269                // Uses `Error::RegionServer` instead of `Error::Server`
270                error::Error::RegionServer {
271                    addr,
272                    code,
273                    source: BoxedError::new(error::Error::from(e)),
274                    location: location!(),
275                }
276            })?
277            .into_inner();
278
279        check_response_header(&response.header)?;
280
281        Ok(RegionResponse::from_region_response(response))
282    }
283
284    pub async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
285        self.handle_inner(request).await
286    }
287}
288
289pub fn check_response_header(header: &Option<ResponseHeader>) -> Result<()> {
290    let status = header
291        .as_ref()
292        .and_then(|header| header.status.as_ref())
293        .context(IllegalDatabaseResponseSnafu {
294            err_msg: "either response header or status is missing",
295        })?;
296
297    if StatusCode::is_success(status.status_code) {
298        Ok(())
299    } else {
300        let code =
301            StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu {
302                err_msg: format!("unknown server status: {:?}", status),
303            })?;
304        ServerSnafu {
305            code,
306            msg: status.err_msg.clone(),
307        }
308        .fail()
309    }
310}
311
312#[cfg(test)]
313mod test {
314    use api::v1::Status as PbStatus;
315
316    use super::*;
317    use crate::Error::{IllegalDatabaseResponse, Server};
318
319    #[test]
320    fn test_check_response_header() {
321        let result = check_response_header(&None);
322        assert!(matches!(
323            result.unwrap_err(),
324            IllegalDatabaseResponse { .. }
325        ));
326
327        let result = check_response_header(&Some(ResponseHeader { status: None }));
328        assert!(matches!(
329            result.unwrap_err(),
330            IllegalDatabaseResponse { .. }
331        ));
332
333        let result = check_response_header(&Some(ResponseHeader {
334            status: Some(PbStatus {
335                status_code: StatusCode::Success as u32,
336                err_msg: String::default(),
337            }),
338        }));
339        assert!(result.is_ok());
340
341        let result = check_response_header(&Some(ResponseHeader {
342            status: Some(PbStatus {
343                status_code: u32::MAX,
344                err_msg: String::default(),
345            }),
346        }));
347        assert!(matches!(
348            result.unwrap_err(),
349            IllegalDatabaseResponse { .. }
350        ));
351
352        let result = check_response_header(&Some(ResponseHeader {
353            status: Some(PbStatus {
354                status_code: StatusCode::Internal as u32,
355                err_msg: "blabla".to_string(),
356            }),
357        }));
358        let Server { code, msg, .. } = result.unwrap_err() else {
359            unreachable!()
360        };
361        assert_eq!(code, StatusCode::Internal);
362        assert_eq!(msg, "blabla");
363    }
364}