risingwave_batch/task/
task_execution.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::fmt::{Debug, Formatter};
16use std::panic::AssertUnwindSafe;
17use std::sync::Arc;
18
19use anyhow::Context;
20use futures::StreamExt;
21use parking_lot::Mutex;
22use risingwave_common::array::DataChunk;
23use risingwave_common::util::panic::FutureCatchUnwindExt;
24use risingwave_common::util::runtime::BackgroundShutdownRuntime;
25use risingwave_common::util::tracing::TracingContext;
26use risingwave_expr::expr_context::expr_context_scope;
27use risingwave_pb::PbFieldNotFound;
28use risingwave_pb::batch_plan::{PbTaskId, PbTaskOutputId, PlanFragment};
29use risingwave_pb::plan_common::ExprContext;
30use risingwave_pb::task_service::task_info_response::TaskStatus;
31use risingwave_pb::task_service::{GetDataResponse, TaskInfoResponse};
32use thiserror_ext::AsReport;
33use tokio::select;
34use tokio::task::JoinHandle;
35use tracing::Instrument;
36
37use crate::error::BatchError::SenderError;
38use crate::error::{BatchError, Result, SharedResult};
39use crate::executor::{BoxedExecutor, ExecutorBuilder};
40use crate::rpc::service::exchange::ExchangeWriter;
41use crate::rpc::service::task_service::TaskInfoResponseResult;
42use crate::task::BatchTaskContext;
43use crate::task::channel::{ChanReceiverImpl, ChanSenderImpl, create_output_channel};
44
45// Now we will only at most have 2 status for each status channel. Running -> Failed or Finished.
46pub const TASK_STATUS_BUFFER_SIZE: usize = 2;
47
48/// Send batch task status (local/distributed) to frontend.
49///
50///
51/// Local mode use `StateReporter::Local`, Distributed mode use `StateReporter::Distributed` to send
52/// status (Failed/Finished) update. `StateReporter::Mock` is only used in test and do not takes any
53/// effect. Local sender only report Failed update, Distributed sender will also report
54/// Finished/Pending/Starting/Aborted etc.
55#[derive(Clone)]
56pub enum StateReporter {
57    Distributed(tokio::sync::mpsc::Sender<TaskInfoResponseResult>),
58    Mock(),
59}
60
61impl StateReporter {
62    pub async fn send(&mut self, val: TaskInfoResponse) -> Result<()> {
63        match self {
64            Self::Distributed(s) => s.send(Ok(val)).await.map_err(|_| SenderError),
65            Self::Mock() => Ok(()),
66        }
67    }
68
69    pub fn new_with_dist_sender(s: tokio::sync::mpsc::Sender<TaskInfoResponseResult>) -> Self {
70        Self::Distributed(s)
71    }
72
73    pub fn new_with_test() -> Self {
74        Self::Mock()
75    }
76}
77
78#[derive(PartialEq, Eq, Hash, Clone, Debug, Default)]
79pub struct TaskId {
80    pub task_id: u64,
81    pub stage_id: u32,
82    pub query_id: String,
83}
84
85#[derive(PartialEq, Eq, Hash, Clone, Default)]
86pub struct TaskOutputId {
87    pub task_id: TaskId,
88    pub output_id: u64,
89}
90
91/// More compact formatter compared to derived `fmt::Debug`.
92impl Debug for TaskOutputId {
93    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
94        f.write_fmt(format_args!(
95            "TaskOutputId {{ query_id: \"{}\", stage_id: {}, task_id: {}, output_id: {} }}",
96            self.task_id.query_id, self.task_id.stage_id, self.task_id.task_id, self.output_id
97        ))
98    }
99}
100
101impl From<&PbTaskId> for TaskId {
102    fn from(prost: &PbTaskId) -> Self {
103        TaskId {
104            task_id: prost.task_id,
105            stage_id: prost.stage_id,
106            query_id: prost.query_id.clone(),
107        }
108    }
109}
110
111impl TaskId {
112    pub fn to_prost(&self) -> PbTaskId {
113        PbTaskId {
114            task_id: self.task_id,
115            stage_id: self.stage_id,
116            query_id: self.query_id.clone(),
117        }
118    }
119}
120
121impl TryFrom<&PbTaskOutputId> for TaskOutputId {
122    type Error = PbFieldNotFound;
123
124    fn try_from(prost: &PbTaskOutputId) -> std::result::Result<Self, PbFieldNotFound> {
125        Ok(TaskOutputId {
126            task_id: TaskId::from(prost.get_task_id()?),
127            output_id: prost.get_output_id(),
128        })
129    }
130}
131
132impl TaskOutputId {
133    pub fn to_prost(&self) -> PbTaskOutputId {
134        PbTaskOutputId {
135            task_id: Some(self.task_id.to_prost()),
136            output_id: self.output_id,
137        }
138    }
139}
140
141pub struct TaskOutput {
142    receiver: ChanReceiverImpl,
143    output_id: TaskOutputId,
144    failure: Arc<Mutex<Option<Arc<BatchError>>>>,
145}
146
147impl std::fmt::Debug for TaskOutput {
148    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
149        f.debug_struct("TaskOutput")
150            .field("output_id", &self.output_id)
151            .field("failure", &self.failure)
152            .finish_non_exhaustive()
153    }
154}
155
156impl TaskOutput {
157    /// Write the data in serialized format to `ExchangeWriter`.
158    /// Return whether the data stream is finished.
159    async fn take_data_inner(
160        &mut self,
161        writer: &mut impl ExchangeWriter,
162        at_most_num: Option<usize>,
163    ) -> Result<bool> {
164        let mut cnt: usize = 0;
165        let limited = at_most_num.is_some();
166        let at_most_num = at_most_num.unwrap_or(usize::MAX);
167        loop {
168            if limited && cnt >= at_most_num {
169                return Ok(false);
170            }
171            match self.receiver.recv().await {
172                // Received some data
173                Ok(Some(chunk)) => {
174                    trace!(
175                        "Task output id: {:?}, data len: {:?}",
176                        self.output_id,
177                        chunk.cardinality()
178                    );
179                    let pb = chunk.to_protobuf().await;
180                    let resp = GetDataResponse {
181                        record_batch: Some(pb),
182                    };
183                    writer.write(Ok(resp)).await?;
184                }
185                // Reached EOF
186                Ok(None) => {
187                    break;
188                }
189                // Error happened
190                Err(e) => {
191                    writer.write(Err(tonic::Status::from(&*e))).await?;
192                    break;
193                }
194            }
195            cnt += 1;
196        }
197        Ok(true)
198    }
199
200    /// Take at most num data and write the data in serialized format to `ExchangeWriter`.
201    /// Return whether the data stream is finished.
202    pub async fn take_data_with_num(
203        &mut self,
204        writer: &mut impl ExchangeWriter,
205        num: usize,
206    ) -> Result<bool> {
207        self.take_data_inner(writer, Some(num)).await
208    }
209
210    /// Take all data and write the data in serialized format to `ExchangeWriter`.
211    pub async fn take_data(&mut self, writer: &mut impl ExchangeWriter) -> Result<()> {
212        let finish = self.take_data_inner(writer, None).await?;
213        assert!(finish);
214        Ok(())
215    }
216
217    /// Directly takes data without serialization.
218    pub async fn direct_take_data(&mut self) -> SharedResult<Option<DataChunk>> {
219        Ok(self.receiver.recv().await?.map(|c| c.into_data_chunk()))
220    }
221
222    pub fn id(&self) -> &TaskOutputId {
223        &self.output_id
224    }
225}
226
227#[derive(Clone, Debug)]
228pub enum ShutdownMsg {
229    /// Used in init, it never occur in receiver later.
230    Init,
231    Abort(String),
232    Cancel,
233}
234
235/// A token which can be used to signal a shutdown request.
236#[derive(Clone)]
237pub struct ShutdownSender(tokio::sync::watch::Sender<ShutdownMsg>);
238
239impl ShutdownSender {
240    /// Send a cancel message. Return true if the message is sent successfully.
241    pub fn cancel(&self) -> bool {
242        self.0.send(ShutdownMsg::Cancel).is_ok()
243    }
244
245    /// Send an abort message. Return true if the message is sent successfully.
246    pub fn abort(&self, msg: impl Into<String>) -> bool {
247        self.0.send(ShutdownMsg::Abort(msg.into())).is_ok()
248    }
249}
250
251/// A token which can be used to receive a shutdown signal.
252#[derive(Clone)]
253pub struct ShutdownToken(tokio::sync::watch::Receiver<ShutdownMsg>);
254
255impl ShutdownToken {
256    /// Create an empty token.
257    pub fn empty() -> Self {
258        Self::new().1
259    }
260
261    /// Create a new token.
262    pub fn new() -> (ShutdownSender, Self) {
263        let (tx, rx) = tokio::sync::watch::channel(ShutdownMsg::Init);
264        (ShutdownSender(tx), ShutdownToken(rx))
265    }
266
267    /// Return error if the shutdown token has been triggered.
268    pub fn check(&self) -> Result<()> {
269        match &*self.0.borrow() {
270            ShutdownMsg::Init => Ok(()),
271            msg => bail!("Receive shutdown msg: {msg:?}"),
272        }
273    }
274
275    /// Wait until cancellation is requested.
276    ///
277    /// # Cancel safety
278    /// This method is cancel safe.
279    pub async fn cancelled(&mut self) {
280        if matches!(*self.0.borrow(), ShutdownMsg::Init)
281            && let Err(_err) = self.0.changed().await
282        {
283            std::future::pending::<()>().await;
284        }
285    }
286
287    /// Return true if the shutdown token has been triggered.
288    pub fn is_cancelled(&self) -> bool {
289        !matches!(*self.0.borrow(), ShutdownMsg::Init)
290    }
291
292    /// Return the current shutdown message.
293    pub fn message(&self) -> ShutdownMsg {
294        self.0.borrow().clone()
295    }
296}
297
298/// `BatchTaskExecution` represents a single task execution.
299pub struct BatchTaskExecution {
300    /// Task id.
301    task_id: TaskId,
302
303    /// Inner plan to execute.
304    plan: PlanFragment,
305
306    /// Task state.
307    state: Mutex<TaskStatus>,
308
309    /// Receivers data of the task.
310    receivers: Mutex<Vec<Option<ChanReceiverImpl>>>,
311
312    /// Sender for sending chunks between different executors.
313    sender: ChanSenderImpl,
314
315    /// Context for task execution
316    context: Arc<dyn BatchTaskContext>,
317
318    /// The execution failure.
319    failure: Arc<Mutex<Option<Arc<BatchError>>>>,
320
321    /// Runtime for the batch tasks.
322    runtime: Arc<BackgroundShutdownRuntime>,
323
324    shutdown_tx: ShutdownSender,
325    shutdown_rx: ShutdownToken,
326    heartbeat_join_handle: Mutex<Option<JoinHandle<()>>>,
327
328    /// Registry for await-tree. The task will be registered as a root when spawned,
329    /// so that it shows up in the await-tree dump.
330    await_tree_reg: Option<await_tree::Registry>,
331}
332
333impl BatchTaskExecution {
334    pub fn new(
335        prost_tid: &PbTaskId,
336        plan: PlanFragment,
337        context: Arc<dyn BatchTaskContext>,
338        runtime: Arc<BackgroundShutdownRuntime>,
339        await_tree_reg: Option<await_tree::Registry>,
340    ) -> Result<Self> {
341        let task_id = TaskId::from(prost_tid);
342
343        let (sender, receivers) = create_output_channel(
344            plan.get_exchange_info()?,
345            context.get_config().developer.output_channel_size,
346        )?;
347
348        let mut rts = Vec::new();
349        rts.extend(receivers.into_iter().map(Some));
350
351        let (shutdown_tx, shutdown_rx) = ShutdownToken::new();
352        Ok(Self {
353            task_id,
354            plan,
355            state: Mutex::new(TaskStatus::Pending),
356            receivers: Mutex::new(rts),
357            failure: Arc::new(Mutex::new(None)),
358            context,
359            runtime,
360            sender,
361            shutdown_tx,
362            shutdown_rx,
363            heartbeat_join_handle: Mutex::new(None),
364            await_tree_reg,
365        })
366    }
367
368    pub fn get_task_id(&self) -> &TaskId {
369        &self.task_id
370    }
371
372    /// `async_execute` executes the task in background, it spawns a tokio coroutine and returns
373    /// immediately. The result produced by the task will be sent to one or more channels, according
374    /// to a particular shuffling strategy. For example, in hash shuffling, the result will be
375    /// hash partitioned across multiple channels.
376    /// To obtain the result, one must pick one of the channels to consume via [`TaskOutputId`]. As
377    /// such, parallel consumers are able to consume the result independently.
378    pub async fn async_execute(
379        self: Arc<Self>,
380        state_tx: Option<StateReporter>,
381        tracing_context: TracingContext,
382        expr_context: ExprContext,
383    ) -> Result<()> {
384        let mut state_tx = state_tx;
385        trace!(
386            "Prepare executing plan [{:?}]: {}",
387            self.task_id,
388            serde_json::to_string_pretty(self.plan.get_root()?).unwrap()
389        );
390
391        let exec = expr_context_scope(
392            expr_context.clone(),
393            ExecutorBuilder::new(
394                self.plan.root.as_ref().unwrap(),
395                &self.task_id,
396                self.context.clone(),
397                self.shutdown_rx.clone(),
398            )
399            .build(),
400        )
401        .await?;
402
403        let sender = self.sender.clone();
404        let _failure = self.failure.clone();
405        let task_id = self.task_id.clone();
406
407        // After we init the output receivers, it's must safe to schedule next stage -- able to send
408        // TaskStatus::Running here.
409        // Init the state receivers. Swap out later.
410        self.change_state_notify(TaskStatus::Running, state_tx.as_mut(), None)
411            .await?;
412
413        // Clone `self` to make compiler happy because of the move block.
414        let t_1 = self.clone();
415        let this = self.clone();
416        async fn notify_panic(
417            this: &BatchTaskExecution,
418            state_tx: Option<&mut StateReporter>,
419            message: Option<&str>,
420        ) {
421            let err_str = if let Some(message) = message {
422                format!("execution panic: {}", message)
423            } else {
424                "execution panic".into()
425            };
426
427            if let Err(e) = this
428                .change_state_notify(TaskStatus::Failed, state_tx, Some(err_str))
429                .await
430            {
431                warn!(
432                    error = %e.as_report(),
433                    "The status receiver in FE has closed so the status push is failed",
434                );
435            }
436        }
437        // Spawn task for real execution.
438        let fut = async move {
439            trace!("Executing plan [{:?}]", task_id);
440            let sender = sender;
441            let mut state_tx_1 = state_tx.clone();
442
443            let task = |task_id: TaskId| async move {
444                let span = tracing_context.attach(tracing::info_span!(
445                    "batch_execute",
446                    task_id = task_id.task_id,
447                    stage_id = task_id.stage_id,
448                    query_id = task_id.query_id,
449                ));
450
451                // We should only pass a reference of sender to execution because we should only
452                // close it after task error has been set.
453                expr_context_scope(
454                    expr_context,
455                    t_1.run(exec, sender, state_tx_1.as_mut()).instrument(span),
456                )
457                .await;
458            };
459
460            if let Err(error) = AssertUnwindSafe(task(task_id.clone()))
461                .rw_catch_unwind()
462                .await
463            {
464                let message = panic_message::get_panic_message(&error);
465                error!(?task_id, error = message, "Batch task panic");
466                notify_panic(&this, state_tx.as_mut(), message).await;
467            }
468        };
469
470        if let Some(reg) = self.await_tree_reg.clone() {
471            let key = crate::task::await_tree_key::BatchTask(self.task_id.clone());
472            let span = await_tree::span!(
473                "Batch Task (query {} stage {} task {})",
474                self.task_id.query_id,
475                self.task_id.stage_id,
476                self.task_id.task_id
477            );
478            self.runtime.spawn(reg.register(key, span).instrument(fut));
479        } else {
480            self.runtime.spawn(fut);
481        }
482
483        Ok(())
484    }
485
486    /// Change state and notify frontend for task status via streaming GRPC.
487    pub async fn change_state_notify(
488        &self,
489        task_status: TaskStatus,
490        state_tx: Option<&mut StateReporter>,
491        err_str: Option<String>,
492    ) -> Result<()> {
493        self.change_state(task_status);
494        // Notify frontend the task status.
495        if let Some(reporter) = state_tx {
496            reporter
497                .send(TaskInfoResponse {
498                    task_id: Some(self.task_id.to_prost()),
499                    task_status: task_status.into(),
500                    error_message: err_str.unwrap_or("".to_owned()),
501                })
502                .await
503        } else {
504            Ok(())
505        }
506    }
507
508    pub fn change_state(&self, task_status: TaskStatus) {
509        *self.state.lock() = task_status;
510        tracing::debug!(
511            "Task {:?} state changed to {:?}",
512            &self.task_id,
513            task_status
514        );
515    }
516
517    async fn run(
518        &self,
519        root: BoxedExecutor,
520        mut sender: ChanSenderImpl,
521        state_tx: Option<&mut StateReporter>,
522    ) {
523        self.context
524            .batch_metrics()
525            .as_ref()
526            .inspect(|m| m.batch_manager_metrics().task_num.inc());
527        let mut data_chunk_stream = root.execute();
528        let mut state;
529        let mut error = None;
530
531        let mut shutdown_rx = self.shutdown_rx.clone();
532        loop {
533            select! {
534                biased;
535                // `shutdown_rx` can't be removed here to avoid `sender.send(data_chunk)` blocked whole execution.
536                _ = shutdown_rx.cancelled() => {
537                    match self.shutdown_rx.message() {
538                        ShutdownMsg::Abort(e) => {
539                            error = Some(BatchError::Aborted(e));
540                            state = TaskStatus::Aborted;
541                            break;
542                        }
543                        ShutdownMsg::Cancel => {
544                            state = TaskStatus::Cancelled;
545                            break;
546                        }
547                        ShutdownMsg::Init => {
548                            unreachable!("Init message should not be received here!")
549                        }
550                    }
551                }
552                data_chunk = data_chunk_stream.next()=> {
553                    match data_chunk {
554                        Some(Ok(data_chunk)) => {
555                            if let Err(e) = sender.send(data_chunk).await {
556                                match e {
557                                    BatchError::SenderError => {
558                                        // This is possible since when we have limit executor in parent
559                                        // stage, it may early stop receiving data from downstream, which
560                                        // leads to close of channel.
561                                        warn!("Task receiver closed!");
562                                        state = TaskStatus::Finished;
563                                        break;
564                                    }
565                                    x => {
566                                        error!("Failed to send data!");
567                                        error = Some(x);
568                                        state = TaskStatus::Failed;
569                                        break;
570                                    }
571                                }
572                            }
573                        }
574                        Some(Err(e)) => match self.shutdown_rx.message() {
575                            ShutdownMsg::Init => {
576                                // There is no message received from shutdown channel, which means it caused
577                                // task failed.
578                                error!(error = %e.as_report(), "Batch task failed");
579                                error = Some(e);
580                                state = TaskStatus::Failed;
581                                break;
582                            }
583                            ShutdownMsg::Abort(_) => {
584                                error = Some(e);
585                                state = TaskStatus::Aborted;
586                                break;
587                            }
588                            ShutdownMsg::Cancel => {
589                                state = TaskStatus::Cancelled;
590                                break;
591                            }
592                        },
593                        None => {
594                            debug!("Batch task {:?} finished successfully.", self.task_id);
595                            state = TaskStatus::Finished;
596                            break;
597                        }
598                    }
599                }
600            }
601        }
602
603        let error = error.map(Arc::new);
604        self.failure.lock().clone_from(&error);
605        let err_str = error.as_ref().map(|e| e.to_report_string());
606        if let Err(e) = sender.close(error).await {
607            match e {
608                SenderError => {
609                    // This is possible since when we have limit executor in parent
610                    // stage, it may early stop receiving data from downstream, which
611                    // leads to close of channel.
612                    warn!("Task receiver closed when sending None!");
613                }
614                _x => {
615                    error!("Failed to close task output channel: {:?}", self.task_id);
616                    state = TaskStatus::Failed;
617                }
618            }
619        }
620
621        if let Err(e) = self.change_state_notify(state, state_tx, err_str).await {
622            warn!(
623                error = %e.as_report(),
624                "The status receiver in FE has closed so the status push is failed",
625            );
626        }
627
628        self.context
629            .batch_metrics()
630            .as_ref()
631            .inspect(|m| m.batch_manager_metrics().task_num.dec());
632    }
633
634    pub fn abort(&self, err_msg: String) {
635        // No need to set state to be Aborted here cuz it will be set by shutdown receiver.
636        // Stop task execution.
637        if self.shutdown_tx.abort(err_msg) {
638            info!("Abort task {:?} done", self.task_id);
639        } else {
640            debug!("The task has already died before this request.")
641        }
642    }
643
644    pub fn cancel(&self) {
645        if !self.shutdown_tx.cancel() {
646            debug!("The task has already died before this request.");
647        }
648    }
649
650    pub fn get_task_output(&self, output_id: &PbTaskOutputId) -> Result<TaskOutput> {
651        let task_id = TaskId::from(output_id.get_task_id()?);
652        let receiver = self.receivers.lock()[output_id.get_output_id() as usize]
653            .take()
654            .with_context(|| {
655                format!(
656                    "Task{:?}'s output{} has already been taken.",
657                    task_id,
658                    output_id.get_output_id(),
659                )
660            })?;
661        let task_output = TaskOutput {
662            receiver,
663            output_id: output_id.try_into()?,
664            failure: self.failure.clone(),
665        };
666        Ok(task_output)
667    }
668
669    pub fn check_if_running(&self) -> Result<()> {
670        if *self.state.lock() != TaskStatus::Running {
671            bail!("task {:?} is not running", self.get_task_id());
672        }
673        Ok(())
674    }
675
676    pub fn check_if_aborted(&self) -> Result<bool> {
677        match *self.state.lock() {
678            TaskStatus::Aborted => Ok(true),
679            TaskStatus::Finished => bail!("task {:?} has been finished", self.get_task_id()),
680            _ => Ok(false),
681        }
682    }
683
684    /// Check the task status: whether has ended.
685    pub fn is_end(&self) -> bool {
686        let guard = self.state.lock();
687        !(*guard == TaskStatus::Running || *guard == TaskStatus::Pending)
688    }
689}
690
691impl BatchTaskExecution {
692    pub(crate) fn set_heartbeat_join_handle(&self, join_handle: JoinHandle<()>) {
693        *self.heartbeat_join_handle.lock() = Some(join_handle);
694    }
695
696    pub(crate) fn heartbeat_join_handle(&self) -> Option<JoinHandle<()>> {
697        self.heartbeat_join_handle.lock().take()
698    }
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704
705    #[test]
706    fn test_task_output_id_debug() {
707        let task_id = TaskId {
708            task_id: 1,
709            stage_id: 2,
710            query_id: "abc".to_owned(),
711        };
712        let task_output_id = TaskOutputId {
713            task_id,
714            output_id: 3,
715        };
716        assert_eq!(
717            format!("{:?}", task_output_id),
718            "TaskOutputId { query_id: \"abc\", stage_id: 2, task_id: 1, output_id: 3 }"
719        );
720    }
721}
risingwave_batch/task/task_execution.rs

risingwave_batch/task/
task_execution.rs
