standalone/
datanode_manager.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::region::{RegionRequest, RegionResponse as RegionResponseV1};
19use async_trait::async_trait;
20use client::region;
21use common_error::ext::BoxedError;
22use common_meta::error::{self as meta_error, Result as MetaResult};
23use common_meta::node_manager::{
24    Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef,
25};
26use common_meta::peer::Peer;
27use common_query::request::QueryRequest;
28use common_recordbatch::SendableRecordBatchStream;
29use common_telemetry::tracing;
30use common_telemetry::tracing_context::{FutureExt, TracingContext};
31use datanode::region_server::RegionServer;
32use servers::grpc::region_server::RegionServerHandler;
33use snafu::ResultExt;
34
35use crate::error::{InvalidRegionRequestSnafu, InvokeRegionServerSnafu, Result};
36
37pub struct StandaloneDatanodeManager {
38    pub region_server: RegionServer,
39    pub flow_server: FlownodeRef,
40}
41
42#[async_trait]
43impl DatanodeManager for StandaloneDatanodeManager {
44    async fn datanode(&self, _datanode: &Peer) -> DatanodeRef {
45        RegionInvoker::arc(self.region_server.clone())
46    }
47}
48
49#[async_trait]
50impl FlownodeManager for StandaloneDatanodeManager {
51    async fn flownode(&self, _node: &Peer) -> FlownodeRef {
52        self.flow_server.clone()
53    }
54}
55
56/// Relative to [client::region::RegionRequester]
57struct RegionInvoker {
58    region_server: RegionServer,
59}
60
61impl RegionInvoker {
62    fn arc(region_server: RegionServer) -> Arc<Self> {
63        Arc::new(Self { region_server })
64    }
65
66    async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponseV1> {
67        self.region_server
68            .handle_region_request(request)
69            .await
70            .context(InvokeRegionServerSnafu)
71    }
72}
73
74#[async_trait]
75impl Datanode for RegionInvoker {
76    async fn handle_request(&self, request: Arc<RegionRequest>) -> MetaResult<RegionResponse> {
77        // Local execution consumes the request; standalone does not retain it for retries.
78        let request = Arc::try_unwrap(request)
79            .map_err(|_| {
80                InvalidRegionRequestSnafu {
81                    reason: "local region request has multiple owners",
82                }
83                .build()
84            })
85            .map_err(BoxedError::new)
86            .context(meta_error::ExternalSnafu)?;
87        let span = request
88            .header
89            .as_ref()
90            .map(|h| TracingContext::from_w3c(&h.tracing_context))
91            .unwrap_or_default()
92            .attach(tracing::info_span!("RegionInvoker::handle_region_request"));
93        let response = self
94            .handle_inner(request)
95            .trace(span)
96            .await
97            .map_err(BoxedError::new)
98            .context(meta_error::ExternalSnafu)?;
99        region::decode_region_response(response).map_err(Into::into)
100    }
101
102    async fn handle_query(&self, request: QueryRequest) -> MetaResult<SendableRecordBatchStream> {
103        let region_id = request.region_id.to_string();
104        let span = request
105            .header
106            .as_ref()
107            .map(|h| TracingContext::from_w3c(&h.tracing_context))
108            .unwrap_or_default()
109            .attach(tracing::info_span!(
110                "RegionInvoker::handle_query",
111                region_id = region_id
112            ));
113        self.region_server
114            .handle_read(request)
115            .trace(span)
116            .await
117            .map_err(BoxedError::new)
118            .context(meta_error::ExternalSnafu)
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use api::v1::Rows;
125    use api::v1::region::{InsertRequest, InsertRequests, region_request};
126    use common_error::ext::ErrorExt;
127    use common_error::mock::MockError;
128    use common_error::status_code::StatusCode;
129    use datanode::tests::{MockRegionEngine, mock_region_server};
130    use mito2::engine::MITO_ENGINE_NAME;
131    use mito2::test_util::CreateRequestBuilder;
132    use store_api::region_request::RegionRequest as StoreRegionRequest;
133    use store_api::storage::RegionId;
134
135    use super::*;
136
137    #[tokio::test]
138    async fn test_write_preserves_original_error() {
139        for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] {
140            let mut server = mock_region_server();
141            let (engine, _receiver) = MockRegionEngine::with_mock_fn(
142                MITO_ENGINE_NAME,
143                Box::new(move |region_id, request| {
144                    if matches!(request, StoreRegionRequest::Put(_)) {
145                        Err(BoxedError::new(MockError::new(code)))
146                            .context(datanode::error::HandleRegionRequestSnafu { region_id })
147                    } else {
148                        Ok(0)
149                    }
150                }),
151            );
152            server.register_engine(engine);
153            let mut requests = Vec::new();
154            for number in 0..2 {
155                let region_id = RegionId::new(1, number);
156                server
157                    .handle_request(
158                        region_id,
159                        StoreRegionRequest::Create(CreateRequestBuilder::new().build()),
160                    )
161                    .await
162                    .unwrap();
163                requests.push(InsertRequest {
164                    region_id: region_id.as_u64(),
165                    rows: Some(Rows::default()),
166                    ..Default::default()
167                });
168            }
169            let datanode: DatanodeRef = RegionInvoker::arc(server);
170            let error = datanode
171                .handle(RegionRequest {
172                    header: Some(Default::default()),
173                    body: Some(region_request::Body::Inserts(InsertRequests { requests })),
174                })
175                .await
176                .unwrap_err();
177            assert_eq!(code, error.status_code());
178            assert_eq!(MockError::new(code).retry_hint(), error.retry_hint());
179            assert!(!matches!(error, meta_error::Error::PartialRegion { .. }));
180        }
181    }
182}
standalone/datanode_manager.rs

standalone/
datanode_manager.rs
