aft rebase fix

Signed-off-by: discord9 <discord9@163.com>

Author: discord9

src/common/meta/src/instruction.rs

@@ -565,6 +565,20 @@ impl Instruction {
             _ => None,
         }
     }
+
+    pub fn into_get_file_refs(self) -> Option<GetFileRefs> {
+        match self {
+            Self::GetFileRefs(get_file_refs) => Some(get_file_refs),
+            _ => None,
+        }
+    }
+
+    pub fn into_gc_regions(self) -> Option<GcRegions> {
+        match self {
+            Self::GcRegions(gc_regions) => Some(gc_regions),
+            _ => None,
+        }
+    }
 }
 
 /// The reply of [UpgradeRegion].

src/datanode/src/heartbeat.rs

@@ -321,7 +321,7 @@ impl HeartbeatTask {
                         let gc_stat = gc_limiter.gc_stat();
                         gc_stat.into_extensions(&mut extensions);
 
-                        let req = HeartbeatRequest {
+                        let mut req = HeartbeatRequest {
                             region_stats,
                             topic_stats,
                             duration_since_epoch,

src/datanode/src/heartbeat/handler.rs

@@ -20,7 +20,7 @@ use common_meta::heartbeat::handler::{
 use common_meta::instruction::{Instruction, InstructionReply};
 use common_telemetry::error;
 use snafu::OptionExt;
-use store_api::storage::{GcReport, RegionId};
+use store_api::storage::GcReport;
 
 mod close_region;
 mod downgrade_region;
@@ -32,7 +32,9 @@ mod upgrade_region;
 
 use crate::heartbeat::handler::close_region::CloseRegionsHandler;
 use crate::heartbeat::handler::downgrade_region::DowngradeRegionsHandler;
+use crate::heartbeat::handler::file_ref::GetFileRefsHandler;
 use crate::heartbeat::handler::flush_region::FlushRegionsHandler;
+use crate::heartbeat::handler::gc_worker::GcRegionsHandler;
 use crate::heartbeat::handler::open_region::OpenRegionsHandler;
 use crate::heartbeat::handler::upgrade_region::UpgradeRegionsHandler;
 use crate::heartbeat::task_tracker::TaskTracker;
@@ -110,12 +112,8 @@ impl RegionHeartbeatResponseHandler {
             Instruction::DowngradeRegions(_) => Ok(Box::new(DowngradeRegionsHandler)),
             Instruction::UpgradeRegion(_) => Ok(Box::new(UpgradeRegionsHandler)),
             Instruction::InvalidateCaches(_) => InvalidHeartbeatResponseSnafu.fail(),
-            Instruction::GetFileRefs(get_file_refs) => Ok(Box::new(move |handler_context| {
-                Box::pin(handler_context.handle_get_file_refs_instruction(get_file_refs))
-            })),
-            Instruction::GcRegions(gc_regions) => Ok(Box::new(move |handler_context| {
-                Box::pin(handler_context.handle_gc_regions_instruction(gc_regions))
-            })),
+            Instruction::GetFileRefs(_) => Ok(Box::new(GetFileRefsHandler)),
+            Instruction::GcRegions(_) => Ok(Box::new(GcRegionsHandler)),
         }
     }
 }
@@ -149,14 +147,18 @@ impl HeartbeatResponseHandler for RegionHeartbeatResponseHandler {
         let gc_tasks = self.gc_tasks.clone();
         let handler = self.build_handler(&instruction)?;
         let _handle = common_runtime::spawn_global(async move {
-            let reply = handler(HandlerContext {
-                region_server,
-                catchup_tasks,
-                downgrade_tasks,
-                flush_tasks,
-                gc_tasks,
-            })
-            .await;
+            let reply = handler
+                .handle(
+                    &HandlerContext {
+                        region_server,
+                        catchup_tasks,
+                        downgrade_tasks,
+                        flush_tasks,
+                        gc_tasks,
+                    },
+                    instruction,
+                )
+                .await;
 
             if let Some(reply) = reply
                 && let Err(e) = mailbox.send((meta, reply)).await

src/datanode/src/heartbeat/handler/file_ref.rs

@@ -13,17 +13,22 @@
 // limitations under the License.
 
 use common_error::ext::ErrorExt;
-use common_meta::instruction::{GetFileRefs, GetFileRefsReply, InstructionReply};
+use common_meta::instruction::{GetFileRefsReply, Instruction, InstructionReply};
 use store_api::storage::FileRefsManifest;
 
-use crate::heartbeat::handler::HandlerContext;
-impl HandlerContext {
-    /// Handles GetFileRefs instruction by getting file references from MitoEngine.
-    pub(crate) async fn handle_get_file_refs_instruction(
-        self,
-        get_file_refs: GetFileRefs,
+use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
+
+pub struct GetFileRefsHandler;
+
+#[async_trait::async_trait]
+impl InstructionHandler for GetFileRefsHandler {
+    async fn handle(
+        &self,
+        ctx: &HandlerContext,
+        instruction: Instruction,
     ) -> Option<InstructionReply> {
-        let region_server = self.region_server;
+        let get_file_refs = instruction.into_get_file_refs().unwrap();
+        let region_server = &ctx.region_server;
 
         // Get the MitoEngine
         let Some(mito_engine) = region_server.mito_engine() else {

src/datanode/src/heartbeat/handler/gc_worker.rs

@@ -12,20 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use common_meta::instruction::{GcRegions, GcRegionsReply, InstructionReply};
+use common_meta::instruction::{GcRegionsReply, Instruction, InstructionReply};
 use common_telemetry::{debug, warn};
 use mito2::gc::LocalGcWorker;
 use snafu::{OptionExt, ResultExt};
 use store_api::storage::{FileRefsManifest, RegionId};
 
 use crate::error::{GcMitoEngineSnafu, InvalidGcArgsSnafu, Result, UnexpectedSnafu};
-use crate::heartbeat::handler::HandlerContext;
+use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
 
-impl HandlerContext {
-    pub(crate) async fn handle_gc_regions_instruction(
-        self,
-        gc_regions: GcRegions,
+pub struct GcRegionsHandler;
+
+#[async_trait::async_trait]
+impl InstructionHandler for GcRegionsHandler {
+    async fn handle(
+        &self,
+        ctx: &HandlerContext,
+        instruction: Instruction,
     ) -> Option<InstructionReply> {
+        let gc_regions = instruction.into_gc_regions().unwrap();
         let region_ids = gc_regions.regions.clone();
         debug!("Received gc regions instruction: {:?}", region_ids);
 
@@ -45,6 +50,7 @@ impl HandlerContext {
 
         let (region_id, gc_worker) = match self
             .create_gc_worker(
+                ctx,
                 region_ids,
                 &gc_regions.file_refs_manifest,
                 gc_regions.full_file_listing,
@@ -59,7 +65,7 @@ impl HandlerContext {
             }
         };
 
-        let register_result = self
+        let register_result = ctx
             .gc_tasks
             .try_register(
                 region_id,
@@ -83,7 +89,7 @@ impl HandlerContext {
             }));
         }
         let mut watcher = register_result.into_watcher();
-        let result = self.gc_tasks.wait_until_finish(&mut watcher).await;
+        let result = ctx.gc_tasks.wait_until_finish(&mut watcher).await;
         match result {
             Ok(report) => Some(InstructionReply::GcRegions(GcRegionsReply {
                 result: Ok(report),
@@ -93,16 +99,19 @@ impl HandlerContext {
             })),
         }
     }
+}
 
+impl GcRegionsHandler {
     async fn create_gc_worker(
         &self,
+        ctx: &HandlerContext,
         mut region_ids: Vec<RegionId>,
         file_ref_manifest: &FileRefsManifest,
         full_file_listing: bool,
     ) -> Result<(RegionId, LocalGcWorker)> {
         // always use the smallest region id on datanode as the target region id
         region_ids.sort_by_key(|r| r.region_number());
-        let mito_engine = self
+        let mito_engine = ctx
             .region_server
             .mito_engine()
             .with_context(|| UnexpectedSnafu {