risingwave_batch/task/
env.rs

1// Copyright 2022 RisingWave Labs
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 risingwave_common::config::{BatchConfig, MetricLevel};
18use risingwave_common::util::addr::HostAddr;
19use risingwave_common::util::worker_util::WorkerNodeId;
20use risingwave_connector::source::iceberg::IcebergScanMetrics;
21use risingwave_connector::source::monitor::SourceMetrics;
22use risingwave_dml::dml_manager::DmlManagerRef;
23use risingwave_rpc_client::ComputeClientPoolRef;
24use risingwave_storage::StateStoreImpl;
25
26use crate::monitor::{BatchExecutorMetrics, BatchManagerMetrics, BatchSpillMetrics};
27use crate::task::BatchManager;
28
29/// The global environment for task execution.
30/// The instance will be shared by every task.
31#[derive(Clone)]
32pub struct BatchEnvironment {
33    /// Endpoint the batch task manager listens on.
34    server_addr: HostAddr,
35
36    /// Reference to the task manager.
37    task_manager: Arc<BatchManager>,
38
39    /// Batch related configurations.
40    config: Arc<BatchConfig>,
41
42    /// Current worker node id.
43    worker_id: WorkerNodeId,
44
45    /// State store for table scanning.
46    state_store: StateStoreImpl,
47
48    /// Executor level metrics.
49    executor_metrics: Arc<BatchExecutorMetrics>,
50
51    /// Compute client pool for batch gRPC exchange.
52    client_pool: ComputeClientPoolRef,
53
54    /// Manages dml information.
55    dml_manager: DmlManagerRef,
56
57    /// Metrics for source.
58    source_metrics: Arc<SourceMetrics>,
59
60    /// Batch spill metrics
61    spill_metrics: Arc<BatchSpillMetrics>,
62
63    /// Metrics for iceberg scan.
64    iceberg_scan_metrics: Arc<IcebergScanMetrics>,
65
66    metric_level: MetricLevel,
67}
68
69impl BatchEnvironment {
70    #[expect(clippy::too_many_arguments)]
71    pub fn new(
72        task_manager: Arc<BatchManager>,
73        server_addr: HostAddr,
74        config: Arc<BatchConfig>,
75        worker_id: WorkerNodeId,
76        state_store: StateStoreImpl,
77        executor_metrics: Arc<BatchExecutorMetrics>,
78        client_pool: ComputeClientPoolRef,
79        dml_manager: DmlManagerRef,
80        source_metrics: Arc<SourceMetrics>,
81        spill_metrics: Arc<BatchSpillMetrics>,
82        iceberg_scan_metrics: Arc<IcebergScanMetrics>,
83        metric_level: MetricLevel,
84    ) -> Self {
85        BatchEnvironment {
86            server_addr,
87            task_manager,
88            config,
89            worker_id,
90            state_store,
91            executor_metrics,
92            client_pool,
93            dml_manager,
94            source_metrics,
95            spill_metrics,
96            iceberg_scan_metrics,
97            metric_level,
98        }
99    }
100
101    // Create an instance for testing purpose.
102    pub fn for_test() -> Self {
103        use risingwave_dml::dml_manager::DmlManager;
104        use risingwave_rpc_client::ComputeClientPool;
105        use risingwave_storage::monitor::MonitoredStorageMetrics;
106
107        BatchEnvironment {
108            task_manager: Arc::new(BatchManager::new(
109                BatchConfig::default(),
110                BatchManagerMetrics::for_test(),
111                u64::MAX,
112                None,
113            )),
114            server_addr: "127.0.0.1:2333".parse().unwrap(),
115            config: Arc::new(BatchConfig::default()),
116            worker_id: WorkerNodeId::default(),
117            state_store: StateStoreImpl::shared_in_memory_store(Arc::new(
118                MonitoredStorageMetrics::unused(),
119            )),
120            client_pool: Arc::new(ComputeClientPool::for_test()),
121            dml_manager: Arc::new(DmlManager::for_test()),
122            source_metrics: Arc::new(SourceMetrics::default()),
123            executor_metrics: BatchExecutorMetrics::for_test(),
124            spill_metrics: BatchSpillMetrics::for_test(),
125            iceberg_scan_metrics: IcebergScanMetrics::for_test(),
126            metric_level: MetricLevel::Debug,
127        }
128    }
129
130    pub fn server_address(&self) -> &HostAddr {
131        &self.server_addr
132    }
133
134    pub fn task_manager(&self) -> Arc<BatchManager> {
135        self.task_manager.clone()
136    }
137
138    pub fn config(&self) -> &BatchConfig {
139        self.config.as_ref()
140    }
141
142    pub fn worker_id(&self) -> WorkerNodeId {
143        self.worker_id
144    }
145
146    pub fn state_store(&self) -> StateStoreImpl {
147        self.state_store.clone()
148    }
149
150    pub fn manager_metrics(&self) -> Arc<BatchManagerMetrics> {
151        self.task_manager.metrics()
152    }
153
154    pub fn executor_metrics(&self) -> Arc<BatchExecutorMetrics> {
155        self.executor_metrics.clone()
156    }
157
158    pub fn client_pool(&self) -> ComputeClientPoolRef {
159        self.client_pool.clone()
160    }
161
162    pub fn dml_manager_ref(&self) -> DmlManagerRef {
163        self.dml_manager.clone()
164    }
165
166    pub fn source_metrics(&self) -> Arc<SourceMetrics> {
167        self.source_metrics.clone()
168    }
169
170    pub fn spill_metrics(&self) -> Arc<BatchSpillMetrics> {
171        self.spill_metrics.clone()
172    }
173
174    pub fn metric_level(&self) -> MetricLevel {
175        self.metric_level
176    }
177
178    pub fn iceberg_scan_metrics(&self) -> Arc<IcebergScanMetrics> {
179        self.iceberg_scan_metrics.clone()
180    }
181}
risingwave_batch/task/env.rs

risingwave_batch/task/
env.rs
