risingwave_batch/task/
task_manager.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::collections::{HashMap, hash_map};
16use std::net::SocketAddr;
17use std::sync::Arc;
18
19use anyhow::Context;
20use parking_lot::Mutex;
21use risingwave_common::config::BatchConfig;
22use risingwave_common::memory::MemoryContext;
23use risingwave_common::util::runtime::BackgroundShutdownRuntime;
24use risingwave_common::util::tracing::TracingContext;
25use risingwave_pb::batch_plan::{PbTaskId, PbTaskOutputId, PlanFragment};
26use risingwave_pb::plan_common::ExprContext;
27use risingwave_pb::task_service::task_info_response::TaskStatus;
28use risingwave_pb::task_service::{GetDataResponse, TaskInfoResponse};
29use tokio::sync::mpsc::Sender;
30use tonic::Status;
31
32use super::BatchTaskContext;
33use crate::error::Result;
34use crate::monitor::BatchManagerMetrics;
35use crate::rpc::service::exchange::GrpcExchangeWriter;
36use crate::task::{BatchTaskExecution, StateReporter, TaskId, TaskOutput, TaskOutputId};
37
38pub mod await_tree_key {
39    /// Await-tree key type for batch tasks.
40    #[derive(Debug, Clone, PartialEq, Eq, Hash)]
41    pub struct BatchTask(pub crate::task::TaskId);
42}
43
44/// `BatchManager` is responsible for managing all batch tasks.
45#[derive(Clone)]
46pub struct BatchManager {
47    /// Every task id has a corresponding task execution.
48    tasks: Arc<Mutex<HashMap<TaskId, Arc<BatchTaskExecution>>>>,
49
50    /// Runtime for the batch manager.
51    runtime: Arc<BackgroundShutdownRuntime>,
52
53    /// Batch configuration
54    config: BatchConfig,
55
56    /// Memory context used for batch tasks in cn.
57    mem_context: MemoryContext,
58
59    /// Metrics for batch manager.
60    metrics: Arc<BatchManagerMetrics>,
61
62    /// Registry for await-tree.
63    await_tree_reg: Option<await_tree::Registry>,
64}
65
66impl BatchManager {
67    pub fn new(
68        config: BatchConfig,
69        metrics: Arc<BatchManagerMetrics>,
70        mem_limit: u64,
71        await_tree_config: Option<await_tree::Config>,
72    ) -> Self {
73        let runtime = {
74            let mut builder = tokio::runtime::Builder::new_multi_thread();
75            if let Some(worker_threads_num) = config.worker_threads_num {
76                builder.worker_threads(worker_threads_num);
77            }
78            builder
79                .thread_name("rw-batch")
80                .enable_all()
81                .build()
82                .unwrap()
83        };
84
85        let mem_context = MemoryContext::root(metrics.batch_total_mem.clone(), mem_limit);
86        BatchManager {
87            tasks: Arc::new(Mutex::new(HashMap::new())),
88            runtime: Arc::new(runtime.into()),
89            config,
90            metrics,
91            mem_context,
92            await_tree_reg: await_tree_config.map(await_tree::Registry::new),
93        }
94    }
95
96    /// Get the registry of await-trees.
97    pub fn await_tree_reg(&self) -> Option<&await_tree::Registry> {
98        self.await_tree_reg.as_ref()
99    }
100
101    pub(crate) fn metrics(&self) -> Arc<BatchManagerMetrics> {
102        self.metrics.clone()
103    }
104
105    pub fn memory_context_ref(&self) -> MemoryContext {
106        self.mem_context.clone()
107    }
108
109    pub async fn fire_task(
110        self: &Arc<Self>,
111        tid: &PbTaskId,
112        plan: PlanFragment,
113        context: Arc<dyn BatchTaskContext>, // ComputeNodeContext
114        state_reporter: StateReporter,
115        tracing_context: TracingContext,
116        expr_context: ExprContext,
117    ) -> Result<()> {
118        trace!("Received task id: {:?}, plan: {:?}", tid, plan);
119        let task = BatchTaskExecution::new(
120            tid,
121            plan,
122            context,
123            self.runtime(),
124            self.await_tree_reg.clone(),
125        )?;
126        let task_id = task.get_task_id().clone();
127        let task = Arc::new(task);
128        // Here the task id insert into self.tasks is put in front of `.async_execute`, cuz when
129        // send `TaskStatus::Running` in `.async_execute`, the query runner may schedule next stage,
130        // it's possible do not found parent task id in theory.
131        let ret = if let hash_map::Entry::Vacant(e) = self.tasks.lock().entry(task_id.clone()) {
132            e.insert(task.clone());
133
134            let this = self.clone();
135            let task_id = task_id.clone();
136            let state_reporter = state_reporter.clone();
137            let heartbeat_join_handle = self.runtime.spawn(async move {
138                this.start_task_heartbeat(state_reporter, task_id).await;
139            });
140            task.set_heartbeat_join_handle(heartbeat_join_handle);
141
142            Ok(())
143        } else {
144            bail!(
145                "can not create duplicate task with the same id: {:?}",
146                task_id,
147            );
148        };
149        task.async_execute(Some(state_reporter), tracing_context, expr_context)
150            .await
151            .inspect_err(|_| {
152                self.cancel_task(&task_id.to_prost());
153            })?;
154        ret
155    }
156
157    #[cfg(test)]
158    async fn fire_task_for_test(
159        self: &Arc<Self>,
160        tid: &PbTaskId,
161        plan: PlanFragment,
162    ) -> Result<()> {
163        use crate::task::ComputeNodeContext;
164
165        self.fire_task(
166            tid,
167            plan,
168            ComputeNodeContext::for_test(),
169            StateReporter::new_with_test(),
170            TracingContext::none(),
171            ExprContext {
172                time_zone: "UTC".to_owned(),
173                strict_mode: false,
174            },
175        )
176        .await
177    }
178
179    async fn start_task_heartbeat(&self, mut state_reporter: StateReporter, task_id: TaskId) {
180        let _metric_guard = scopeguard::guard((), |_| {
181            tracing::debug!("heartbeat worker for task {:?} stopped", task_id);
182            self.metrics.batch_heartbeat_worker_num.dec();
183        });
184        tracing::debug!("heartbeat worker for task {:?} started", task_id);
185        self.metrics.batch_heartbeat_worker_num.inc();
186        // The heartbeat is to ensure task cancellation when frontend's cancellation request fails
187        // to reach compute node (for any reason like RPC fails, frontend crashes).
188        let mut heartbeat_interval = tokio::time::interval(core::time::Duration::from_secs(60));
189        heartbeat_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
190        heartbeat_interval.reset();
191        loop {
192            heartbeat_interval.tick().await;
193            if !self.tasks.lock().contains_key(&task_id) {
194                break;
195            }
196            if state_reporter
197                .send(TaskInfoResponse {
198                    task_id: Some(task_id.to_prost()),
199                    task_status: TaskStatus::Ping.into(),
200                    error_message: "".to_owned(),
201                })
202                .await
203                .is_err()
204            {
205                tracing::warn!("try to cancel task {:?} due to heartbeat", task_id);
206                // Task may have been cancelled, but it's fine to `cancel_task` again.
207                self.cancel_task(&task_id.to_prost());
208                break;
209            }
210        }
211    }
212
213    pub fn get_data(
214        &self,
215        tx: Sender<std::result::Result<GetDataResponse, Status>>,
216        peer_addr: SocketAddr,
217        pb_task_output_id: &PbTaskOutputId,
218    ) -> Result<()> {
219        let task_id = TaskOutputId::try_from(pb_task_output_id)?;
220        tracing::debug!(target: "events::compute::exchange", peer_addr = %peer_addr, from = ?task_id, "serve exchange RPC");
221        let mut task_output = self.take_output(pb_task_output_id)?;
222        self.runtime.spawn(async move {
223            let mut writer = GrpcExchangeWriter::new(tx.clone());
224            tokio::select! {
225                result = task_output.take_data(&mut writer) => match result {
226                    Ok(_) => {
227                        tracing::trace!(
228                            from = ?task_id,
229                            "exchanged {} chunks",
230                            writer.written_chunks(),
231                        );
232                        Ok(())
233                    }
234                    Err(e) => tx.send(Err(e.into())).await,
235                },
236                _ = tx.closed() => Ok(()),
237            }
238        });
239        Ok(())
240    }
241
242    pub fn take_output(&self, output_id: &PbTaskOutputId) -> Result<TaskOutput> {
243        let task_id = TaskId::from(output_id.get_task_id()?);
244        self.tasks
245            .lock()
246            .get(&task_id)
247            .with_context(|| format!("task {:?} not found", task_id))?
248            .get_task_output(output_id)
249    }
250
251    pub fn cancel_task(&self, sid: &PbTaskId) {
252        let sid = TaskId::from(sid);
253        match self.tasks.lock().remove(&sid) {
254            Some(task) => {
255                tracing::trace!("Removed task: {:?}", task.get_task_id());
256                // Use `cancel` rather than `abort` here since this is not an error which should be
257                // propagated to upstream.
258                task.cancel();
259                if let Some(heartbeat_join_handle) = task.heartbeat_join_handle() {
260                    heartbeat_join_handle.abort();
261                }
262            }
263            None => {
264                warn!("Task {:?} not found for cancel", sid)
265            }
266        };
267    }
268
269    /// Returns error if task is not running.
270    pub fn check_if_task_running(&self, task_id: &TaskId) -> Result<()> {
271        match self.tasks.lock().get(task_id) {
272            Some(task) => task.check_if_running(),
273            None => bail!("task {:?} not found", task_id),
274        }
275    }
276
277    pub fn check_if_task_aborted(&self, task_id: &TaskId) -> Result<bool> {
278        match self.tasks.lock().get(task_id) {
279            Some(task) => task.check_if_aborted(),
280            None => bail!("task {:?} not found", task_id),
281        }
282    }
283
284    #[cfg(test)]
285    async fn wait_until_task_aborted(&self, task_id: &TaskId) -> Result<()> {
286        use std::time::Duration;
287        loop {
288            match self.tasks.lock().get(task_id) {
289                Some(task) => {
290                    let ret = task.check_if_aborted();
291                    match ret {
292                        Ok(true) => return Ok(()),
293                        Ok(false) => {}
294                        Err(err) => return Err(err),
295                    }
296                }
297                None => bail!("task {:?} not found", task_id),
298            }
299            tokio::time::sleep(Duration::from_millis(100)).await
300        }
301    }
302
303    pub fn runtime(&self) -> Arc<BackgroundShutdownRuntime> {
304        self.runtime.clone()
305    }
306
307    pub fn config(&self) -> &BatchConfig {
308        &self.config
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use std::sync::Arc;
315
316    use risingwave_common::config::BatchConfig;
317    use risingwave_pb::batch_plan::exchange_info::DistributionMode;
318    use risingwave_pb::batch_plan::plan_node::NodeBody;
319    use risingwave_pb::batch_plan::{
320        ExchangeInfo, PbTaskId, PbTaskOutputId, PlanFragment, PlanNode,
321    };
322
323    use crate::monitor::BatchManagerMetrics;
324    use crate::task::{BatchManager, TaskId};
325
326    #[tokio::test]
327    async fn test_task_not_found() {
328        let manager = Arc::new(BatchManager::new(
329            BatchConfig::default(),
330            BatchManagerMetrics::for_test(),
331            u64::MAX,
332            None,
333        ));
334        let task_id = TaskId {
335            task_id: 0,
336            stage_id: 0,
337            query_id: "abc".to_owned(),
338        };
339
340        let error = manager.check_if_task_running(&task_id).unwrap_err();
341        assert!(error.to_string().contains("not found"), "{:?}", error);
342
343        let output_id = PbTaskOutputId {
344            task_id: Some(risingwave_pb::batch_plan::TaskId {
345                stage_id: 0,
346                task_id: 0,
347                query_id: "".to_owned(),
348            }),
349            output_id: 0,
350        };
351        let error = manager.take_output(&output_id).unwrap_err();
352        assert!(error.to_string().contains("not found"), "{:?}", error);
353    }
354
355    #[tokio::test]
356    // see https://github.com/risingwavelabs/risingwave/issues/11979
357    #[ignore]
358    async fn test_task_cancel_for_busy_loop() {
359        let manager = Arc::new(BatchManager::new(
360            BatchConfig::default(),
361            BatchManagerMetrics::for_test(),
362            u64::MAX,
363            None,
364        ));
365        let plan = PlanFragment {
366            root: Some(PlanNode {
367                children: vec![],
368                identity: "".to_owned(),
369                node_body: Some(NodeBody::BusyLoopExecutor(true)),
370            }),
371            exchange_info: Some(ExchangeInfo {
372                mode: DistributionMode::Single as i32,
373                distribution: None,
374            }),
375        };
376        let task_id = PbTaskId {
377            query_id: "".to_owned(),
378            stage_id: 0,
379            task_id: 0,
380        };
381        manager.fire_task_for_test(&task_id, plan).await.unwrap();
382        manager.cancel_task(&task_id);
383        let task_id = TaskId::from(&task_id);
384        assert!(!manager.tasks.lock().contains_key(&task_id));
385    }
386
387    #[tokio::test]
388    // see https://github.com/risingwavelabs/risingwave/issues/11979
389    #[ignore]
390    async fn test_task_abort_for_busy_loop() {
391        let manager = Arc::new(BatchManager::new(
392            BatchConfig::default(),
393            BatchManagerMetrics::for_test(),
394            u64::MAX,
395            None,
396        ));
397        let plan = PlanFragment {
398            root: Some(PlanNode {
399                children: vec![],
400                identity: "".to_owned(),
401                node_body: Some(NodeBody::BusyLoopExecutor(true)),
402            }),
403            exchange_info: Some(ExchangeInfo {
404                mode: DistributionMode::Single as i32,
405                distribution: None,
406            }),
407        };
408        let task_id = PbTaskId {
409            query_id: "".to_owned(),
410            stage_id: 0,
411            task_id: 0,
412        };
413        manager.fire_task_for_test(&task_id, plan).await.unwrap();
414        let task_id = TaskId::from(&task_id);
415        manager
416            .tasks
417            .lock()
418            .get(&task_id)
419            .unwrap()
420            .abort("Abort Test".to_owned());
421        assert!(manager.wait_until_task_aborted(&task_id).await.is_ok());
422    }
423}
risingwave_batch/task/task_manager.rs

risingwave_batch/task/
task_manager.rs
