feat: define region server and related requests (#2160)

* define region server and related requests

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fill request body

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* change mito2's request type

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix clippy

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* chore: bump greptime-proto to d9167cab (row insert/delete)

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix test compile

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* remove name_to_index

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* address cr comments

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* finilise

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

src/datanode/Cargo.toml

@@ -30,6 +30,7 @@ common-recordbatch = { workspace = true }
 common-runtime = { workspace = true }
 common-telemetry = { workspace = true }
 common-time = { workspace = true }
+dashmap = "5.4"
 datafusion-common.workspace = true
 datafusion-expr.workspace = true
 datafusion.workspace = true

src/datanode/src/error.rs

@@ -19,7 +19,7 @@ use common_error::status_code::StatusCode;
 use common_procedure::ProcedureId;
 use serde_json::error::Error as JsonError;
 use snafu::{Location, Snafu};
-use store_api::storage::RegionNumber;
+use store_api::storage::{RegionId, RegionNumber};
 use table::error::Error as TableError;
 
 /// Business error of datanode.
@@ -482,6 +482,30 @@ pub enum Error {
         violated: String,
         location: Location,
     },
+
+    #[snafu(display(
+        "Failed to handle request for region {}, source: {}, location: {}",
+        region_id,
+        source,
+        location
+    ))]
+    HandleRegionRequest {
+        region_id: RegionId,
+        location: Location,
+        source: BoxedError,
+    },
+
+    #[snafu(display("RegionId {} not found, location: {}", region_id, location))]
+    RegionNotFound {
+        region_id: RegionId,
+        location: Location,
+    },
+
+    #[snafu(display("Region engine {} is not registered, location: {}", name, location))]
+    RegionEngineNotFound { name: String, location: Location },
+
+    #[snafu(display("Unsupported gRPC request, kind: {}, location: {}", kind, location))]
+    UnsupportedGrpcRequest { kind: String, location: Location },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -559,7 +583,9 @@ impl ErrorExt for Error {
             | MissingInsertBody { .. }
             | ShutdownInstance { .. }
             | CloseTableEngine { .. }
-            | JoinTask { .. } => StatusCode::Internal,
+            | JoinTask { .. }
+            | RegionNotFound { .. }
+            | RegionEngineNotFound { .. } => StatusCode::Internal,
 
             StartServer { source, .. }
             | ShutdownServer { source, .. }
@@ -570,7 +596,9 @@ impl ErrorExt for Error {
             OpenLogStore { source, .. } => source.status_code(),
             RuntimeResource { .. } => StatusCode::RuntimeResourcesExhausted,
             MetaClientInit { source, .. } => source.status_code(),
-            TableIdProviderNotFound { .. } => StatusCode::Unsupported,
+            TableIdProviderNotFound { .. } | UnsupportedGrpcRequest { .. } => {
+                StatusCode::Unsupported
+            }
             BumpTableId { source, .. } => source.status_code(),
             ColumnDefaultValue { source, .. } => source.status_code(),
             UnrecognizedTableOption { .. } => StatusCode::InvalidArguments,
@@ -581,6 +609,7 @@ impl ErrorExt for Error {
             StartProcedureManager { source } | StopProcedureManager { source } => {
                 source.status_code()
             }
+            HandleRegionRequest { source, .. } => source.status_code(),
         }
     }
 

src/datanode/src/instance/grpc.rs

@@ -42,7 +42,7 @@ use table::table::adapter::DfTableProviderAdapter;
 use crate::error::{
     self, CatalogSnafu, DecodeLogicalPlanSnafu, DeleteExprToRequestSnafu, DeleteSnafu,
     ExecuteLogicalPlanSnafu, ExecuteSqlSnafu, InsertDataSnafu, InsertSnafu, JoinTaskSnafu,
-    PlanStatementSnafu, Result, TableNotFoundSnafu,
+    PlanStatementSnafu, Result, TableNotFoundSnafu, UnsupportedGrpcRequestSnafu,
 };
 use crate::instance::Instance;
 
@@ -221,8 +221,10 @@ impl GrpcQueryHandler for Instance {
                 self.handle_query(query, ctx).await
             }
             Request::Ddl(request) => self.handle_ddl(request, ctx).await,
-            Request::RowInserts(_) => unreachable!(),
-            Request::RowDelete(_) => unreachable!(),
+            Request::RowInserts(_) | Request::RowDelete(_) => UnsupportedGrpcRequestSnafu {
+                kind: "row insert/delete",
+            }
+            .fail(),
         }
     }
 }

src/datanode/src/lib.rs

@@ -23,9 +23,9 @@ pub mod instance;
 pub mod metrics;
 #[cfg(any(test, feature = "testing"))]
 mod mock;
+pub mod region_server;
 pub mod server;
 pub mod sql;
 mod store;
-
 #[cfg(test)]
 mod tests;

src/datanode/src/region_server.rs

@@ -0,0 +1,103 @@
+// 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.
+
+use std::collections::HashMap;
+
+use common_query::Output;
+use common_telemetry::info;
+use dashmap::DashMap;
+use snafu::{OptionExt, ResultExt};
+use store_api::region_engine::RegionEngineRef;
+use store_api::region_request::RegionRequest;
+use store_api::storage::RegionId;
+
+use crate::error::{
+    HandleRegionRequestSnafu, RegionEngineNotFoundSnafu, RegionNotFoundSnafu, Result,
+};
+
+#[derive(Default)]
+pub struct RegionServer {
+    engines: HashMap<String, RegionEngineRef>,
+    region_map: DashMap<RegionId, RegionEngineRef>,
+}
+
+impl RegionServer {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn register_engine(&mut self, engine: RegionEngineRef) {
+        let engine_name = engine.name();
+        self.engines.insert(engine_name.to_string(), engine);
+    }
+
+    pub async fn handle_request(
+        &self,
+        region_id: RegionId,
+        request: RegionRequest,
+    ) -> Result<Output> {
+        // TODO(ruihang): add some metrics
+
+        let region_change = match &request {
+            RegionRequest::Create(create) => RegionChange::Register(create.engine.clone()),
+            RegionRequest::Open(open) => RegionChange::Register(open.engine.clone()),
+            RegionRequest::Close(_) | RegionRequest::Drop(_) => RegionChange::Deregisters,
+            RegionRequest::Write(_)
+            | RegionRequest::Read(_)
+            | RegionRequest::Delete(_)
+            | RegionRequest::Alter(_)
+            | RegionRequest::Flush(_)
+            | RegionRequest::Compact(_) => RegionChange::None,
+        };
+
+        let engine = match &region_change {
+            RegionChange::Register(engine_type) => self
+                .engines
+                .get(engine_type)
+                .with_context(|| RegionEngineNotFoundSnafu { name: engine_type })?
+                .clone(),
+            RegionChange::None | RegionChange::Deregisters => self
+                .region_map
+                .get(&region_id)
+                .with_context(|| RegionNotFoundSnafu { region_id })?
+                .clone(),
+        };
+        let engine_type = engine.name();
+
+        let result = engine
+            .handle_request(region_id, request)
+            .await
+            .with_context(|_| HandleRegionRequestSnafu { region_id })?;
+
+        match region_change {
+            RegionChange::None => {}
+            RegionChange::Register(_) => {
+                info!("Region {region_id} is registered to engine {engine_type}");
+                self.region_map.insert(region_id, engine);
+            }
+            RegionChange::Deregisters => {
+                info!("Region {region_id} is deregistered from engine {engine_type}");
+                self.region_map.remove(&region_id);
+            }
+        }
+
+        Ok(result)
+    }
+}
+
+enum RegionChange {
+    None,
+    Register(String),
+    Deregisters,
+}