wip

Signed-off-by: Zhenchi <[email protected]>

Author: zhongzc

Cargo.lock

@@ -147,7 +147,8 @@ etcd-client = { git = "https://github.com/GreptimeTeam/etcd-client", rev = "f62d
 fst = "0.4.7"
 futures = "0.3"
 futures-util = "0.3"
-greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "14b9dc40bdc8288742b0cefc7bb024303b7429ef" }
+# greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "14b9dc40bdc8288742b0cefc7bb024303b7429ef" }
+greptime-proto = { path = "../greptime-proto" }
 hex = "0.4"
 http = "1"
 humantime = "2.1"

Cargo.toml

@@ -53,6 +53,7 @@ metric-engine.workspace = true
 mito2.workspace = true
 num_cpus.workspace = true
 object-store.workspace = true
+partition.workspace = true
 prometheus.workspace = true
 prost.workspace = true
 query.workspace = true

src/datanode/Cargo.toml

@@ -19,6 +19,7 @@ use common_error::define_into_tonic_status;
 use common_error::ext::{BoxedError, ErrorExt};
 use common_error::status_code::StatusCode;
 use common_macro::stack_trace_debug;
+use mito2::remap_manifest::Error as RemapManifestError;
 use snafu::{Location, Snafu};
 use store_api::storage::RegionId;
 use table::error::Error as TableError;
@@ -396,6 +397,14 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Failed to remap manifests: {}", source))]
+    RemapManifest {
+        #[snafu(implicit)]
+        location: Location,
+        #[snafu(source)]
+        source: RemapManifestError,
+    },
+
     #[snafu(display("Not yet implemented: {what}"))]
     NotYetImplemented { what: String },
 }
@@ -469,6 +478,7 @@ impl ErrorExt for Error {
 
             ObjectStore { source, .. } => source.status_code(),
             BuildCacheStore { .. } => StatusCode::StorageUnavailable,
+            RemapManifest { .. } => StatusCode::Unexpected,
         }
     }
 

src/datanode/src/error.rs

@@ -24,7 +24,9 @@ use api::region::RegionResponse;
 use api::v1::meta::TopicStat;
 use api::v1::region::sync_request::ManifestInfo;
 use api::v1::region::{
-    ListMetadataRequest, RegionResponse as RegionResponseV1, SyncRequest, region_request,
+    ApplyStagedManifestRequest, ListMetadataRequest, PauseRequest, PublishRegionRuleRequest,
+    RegionResponse as RegionResponseV1, RemapManifestRequest, ResumeRequest,
+    StageRegionRuleRequest, SyncRequest, region_request,
 };
 use api::v1::{ResponseHeader, Status};
 use arrow_flight::{FlightData, Ticket};
@@ -84,6 +86,8 @@ use crate::error::{
 use crate::event_listener::RegionServerEventListenerRef;
 use crate::region_server::catalog::{NameAwareCatalogList, NameAwareDataSourceInjectorBuilder};
 
+const REMAP_STATS_EXTENSION_KEY: &str = "repartition.manifest.stats";
+
 #[derive(Clone)]
 pub struct RegionServer {
     inner: Arc<RegionServerInner>,
@@ -370,6 +374,24 @@ impl RegionServer {
         }
     }
 
+    /// Temporarily pauses compaction and snapshot related activities for the region.
+    ///
+    /// Currently a stub; real implementation will coordinate with region worker.
+    pub async fn pause_compaction_and_snapshot(&self, region_id: RegionId) -> Result<()> {
+        info!("pause_compaction_and_snapshot stub invoked for region {region_id}");
+        let _ = region_id;
+        Ok(())
+    }
+
+    /// Resumes compaction and snapshot related activities for the region.
+    ///
+    /// Currently a stub; real implementation will coordinate with region worker.
+    pub async fn resume_compaction_and_snapshot(&self, region_id: RegionId) -> Result<()> {
+        info!("resume_compaction_and_snapshot stub invoked for region {region_id}");
+        let _ = region_id;
+        Ok(())
+    }
+
     /// Stop the region server.
     pub async fn stop(&self) -> Result<()> {
         self.inner.stop().await
@@ -538,6 +560,124 @@ impl RegionServer {
         Ok(response)
     }
 
+    async fn handle_pause_region_request(&self, request: &PauseRequest) -> Result<RegionResponse> {
+        let region_id = RegionId::from_u64(request.region_id);
+        let tracing_context = TracingContext::from_current_span();
+        let span = tracing_context.attach(info_span!(
+            "RegionServer::handle_pause_region_request",
+            region_id = region_id.to_string()
+        ));
+
+        self.pause_compaction_and_snapshot(region_id)
+            .trace(span)
+            .await
+            .map(|_| RegionResponse::new(AffectedRows::default()))
+    }
+
+    async fn handle_resume_region_request(
+        &self,
+        request: &ResumeRequest,
+    ) -> Result<RegionResponse> {
+        let region_id = RegionId::from_u64(request.region_id);
+        let tracing_context = TracingContext::from_current_span();
+        let span = tracing_context.attach(info_span!(
+            "RegionServer::handle_resume_region_request",
+            region_id = region_id.to_string()
+        ));
+
+        self.resume_compaction_and_snapshot(region_id)
+            .trace(span)
+            .await
+            .map(|_| RegionResponse::new(AffectedRows::default()))
+    }
+
+    async fn handle_stage_region_rule_request(
+        &self,
+        request: &StageRegionRuleRequest,
+    ) -> Result<RegionResponse> {
+        let region_id = RegionId::from_u64(request.region_id);
+        info!(
+            "Stage region rule for region {region_id} with version {}",
+            request.rule_version
+        );
+        match self
+            .set_region_role_state_gracefully(region_id, SettableRegionRoleState::StagingLeader)
+            .await?
+        {
+            SetRegionRoleStateResponse::Success(_) | SetRegionRoleStateResponse::NotFound => {
+                Ok(RegionResponse::new(AffectedRows::default()))
+            }
+            SetRegionRoleStateResponse::InvalidTransition(err) => {
+                Err(err).with_context(|_| HandleRegionRequestSnafu { region_id })
+            }
+        }
+    }
+
+    async fn handle_publish_region_rule_request(
+        &self,
+        request: &PublishRegionRuleRequest,
+    ) -> Result<RegionResponse> {
+        let region_id = RegionId::from_u64(request.region_id);
+        info!(
+            "Publish region rule for region {region_id} with version {}",
+            request.rule_version
+        );
+        match self
+            .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader)
+            .await?
+        {
+            SetRegionRoleStateResponse::Success(_) | SetRegionRoleStateResponse::NotFound => {
+                Ok(RegionResponse::new(AffectedRows::default()))
+            }
+            SetRegionRoleStateResponse::InvalidTransition(err) => {
+                Err(err).with_context(|_| HandleRegionRequestSnafu { region_id })
+            }
+        }
+    }
+
+    async fn handle_remap_manifest_request(
+        &self,
+        request: &RemapManifestRequest,
+    ) -> Result<RegionResponse> {
+        info!(
+            "received remap manifest request for table {} group {}",
+            request.table_id, request.group_id
+        );
+
+        let stats_json = serde_json::to_vec(&serde_json::json!({
+            "files_per_region": HashMap::<u64, usize>::new(),
+            "total_file_refs": 0u64,
+            "empty_regions": Vec::<u64>::new(),
+            "group_id": &request.group_id,
+        }))
+        .context(SerializeJsonSnafu)?;
+
+        let mut extensions = HashMap::new();
+        extensions.insert(REMAP_STATS_EXTENSION_KEY.to_string(), stats_json);
+
+        Ok(RegionResponse {
+            affected_rows: 0,
+            extensions,
+            metadata: Vec::new(),
+        })
+    }
+
+    async fn handle_apply_staged_manifest_request(
+        &self,
+        request: &ApplyStagedManifestRequest,
+    ) -> Result<RegionResponse> {
+        info!(
+            "received manifest apply request for table {} group {} publish={} regions {:?}",
+            request.table_id, request.group_id, request.publish, request.region_ids
+        );
+
+        Ok(RegionResponse {
+            affected_rows: 0,
+            extensions: HashMap::new(),
+            metadata: Vec::new(),
+        })
+    }
+
     /// Sync region manifest and registers new opened logical regions.
     pub async fn sync_region(
         &self,
@@ -569,6 +709,26 @@ impl RegionServerHandler for RegionServer {
             region_request::Body::Sync(sync_request) => {
                 self.handle_sync_region_request(sync_request).await
             }
+            region_request::Body::Pause(pause_request) => {
+                self.handle_pause_region_request(pause_request).await
+            }
+            region_request::Body::Resume(resume_request) => {
+                self.handle_resume_region_request(resume_request).await
+            }
+            region_request::Body::StageRegionRule(stage_request) => {
+                self.handle_stage_region_rule_request(stage_request).await
+            }
+            region_request::Body::PublishRegionRule(publish_request) => {
+                self.handle_publish_region_rule_request(publish_request)
+                    .await
+            }
+            region_request::Body::RemapManifest(remap_request) => {
+                self.handle_remap_manifest_request(remap_request).await
+            }
+            region_request::Body::ApplyStagedManifest(apply_request) => {
+                self.handle_apply_staged_manifest_request(apply_request)
+                    .await
+            }
             region_request::Body::ListMetadata(list_metadata_request) => {
                 self.handle_list_metadata_request(list_metadata_request)
                     .await