From 0061e53f5afd12e599b269285d74653eb9eb49f2 Mon Sep 17 00:00:00 2001 From: Dylan Chen Date: Mon, 23 Sep 2024 18:10:17 +0800 Subject: [PATCH] fix --- src/batch/src/monitor/stats.rs | 13 +++++++++++-- src/batch/src/task/context.rs | 2 +- src/batch/src/task/env.rs | 2 +- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/batch/src/monitor/stats.rs b/src/batch/src/monitor/stats.rs index b84779708a821..a96e4f1756f66 100644 --- a/src/batch/src/monitor/stats.rs +++ b/src/batch/src/monitor/stats.rs @@ -58,8 +58,9 @@ impl BatchExecutorMetrics { } /// Create a new `BatchTaskMetrics` instance used in tests or other places. - pub fn for_test() -> Self { - GLOBAL_BATCH_EXECUTOR_METRICS.clone() + #[cfg(test)] + pub fn for_test() -> Arc { + Arc::new(GLOBAL_BATCH_EXECUTOR_METRICS.clone()) } } @@ -89,6 +90,14 @@ impl BatchMetricsInner { pub fn batch_manager_metrics(&self) -> &BatchManagerMetrics { &self.batch_manager_metrics } + + #[cfg(test)] + pub fn for_test() -> BatchMetrics { + Arc::new(Self { + batch_manager_metrics: BatchManagerMetrics::for_test(), + executor_metrics: BatchExecutorMetrics::for_test(), + }) + } } #[derive(Clone)] diff --git a/src/batch/src/task/context.rs b/src/batch/src/task/context.rs index cc563cfae8a07..9fd82e8bbd2eb 100644 --- a/src/batch/src/task/context.rs +++ b/src/batch/src/task/context.rs @@ -136,7 +136,7 @@ impl ComputeNodeContext { pub fn for_test() -> Self { Self { env: BatchEnvironment::for_test(), - batch_metrics: BatchMetrics::for_test().into(), + batch_metrics: BatchMetricsInner::for_test(), mem_context: MemoryContext::none(), } } diff --git a/src/batch/src/task/env.rs b/src/batch/src/task/env.rs index 45b509201a478..85cbd75360ca8 100644 --- a/src/batch/src/task/env.rs +++ b/src/batch/src/task/env.rs @@ -114,7 +114,7 @@ impl BatchEnvironment { client_pool: Arc::new(ComputeClientPool::for_test()), dml_manager: Arc::new(DmlManager::for_test()), source_metrics: Arc::new(SourceMetrics::default()), - executor_metrics: Arc::new(BatchExecutorMetrics::for_test()), + executor_metrics: BatchExecutorMetrics::for_test(), spill_metrics: BatchSpillMetrics::for_test(), metric_level: MetricLevel::Debug, }