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
329impl BatchTaskExecution {
330    pub fn new(
331        prost_tid: &PbTaskId,
332        plan: PlanFragment,
333        context: Arc<dyn BatchTaskContext>,
334        runtime: Arc<BackgroundShutdownRuntime>,
335    ) -> Result<Self> {
336        let task_id = TaskId::from(prost_tid);
337
338        let (sender, receivers) = create_output_channel(
339            plan.get_exchange_info()?,
340            context.get_config().developer.output_channel_size,
341        )?;
342
343        let mut rts = Vec::new();
344        rts.extend(receivers.into_iter().map(Some));
345
346        let (shutdown_tx, shutdown_rx) = ShutdownToken::new();
347        Ok(Self {
348            task_id,
349            plan,
350            state: Mutex::new(TaskStatus::Pending),
351            receivers: Mutex::new(rts),
352            failure: Arc::new(Mutex::new(None)),
353            context,
354            runtime,
355            sender,
356            shutdown_tx,
357            shutdown_rx,
358            heartbeat_join_handle: Mutex::new(None),
359        })
360    }
361
362    pub fn get_task_id(&self) -> &TaskId {
363        &self.task_id
364    }
365
366    /// `async_execute` executes the task in background, it spawns a tokio coroutine and returns
367    /// immediately. The result produced by the task will be sent to one or more channels, according
368    /// to a particular shuffling strategy. For example, in hash shuffling, the result will be
369    /// hash partitioned across multiple channels.
370    /// To obtain the result, one must pick one of the channels to consume via [`TaskOutputId`]. As
371    /// such, parallel consumers are able to consume the result independently.
372    pub async fn async_execute(
373        self: Arc<Self>,
374        state_tx: Option<StateReporter>,
375        tracing_context: TracingContext,
376        expr_context: ExprContext,
377    ) -> Result<()> {
378        let mut state_tx = state_tx;
379        trace!(
380            "Prepare executing plan [{:?}]: {}",
381            self.task_id,
382            serde_json::to_string_pretty(self.plan.get_root()?).unwrap()
383        );
384
385        let exec = expr_context_scope(
386            expr_context.clone(),
387            ExecutorBuilder::new(
388                self.plan.root.as_ref().unwrap(),
389                &self.task_id,
390                self.context.clone(),
391                self.shutdown_rx.clone(),
392            )
393            .build(),
394        )
395        .await?;
396
397        let sender = self.sender.clone();
398        let _failure = self.failure.clone();
399        let task_id = self.task_id.clone();
400
401        // After we init the output receivers, it's must safe to schedule next stage -- able to send
402        // TaskStatus::Running here.
403        // Init the state receivers. Swap out later.
404        self.change_state_notify(TaskStatus::Running, state_tx.as_mut(), None)
405            .await?;
406
407        // Clone `self` to make compiler happy because of the move block.
408        let t_1 = self.clone();
409        let this = self.clone();
410        async fn notify_panic(
411            this: &BatchTaskExecution,
412            state_tx: Option<&mut StateReporter>,
413            message: Option<&str>,
414        ) {
415            let err_str = if let Some(message) = message {
416                format!("execution panic: {}", message)
417            } else {
418                "execution panic".into()
419            };
420
421            if let Err(e) = this
422                .change_state_notify(TaskStatus::Failed, state_tx, Some(err_str))
423                .await
424            {
425                warn!(
426                    error = %e.as_report(),
427                    "The status receiver in FE has closed so the status push is failed",
428                );
429            }
430        }
431        // Spawn task for real execution.
432        let fut = async move {
433            trace!("Executing plan [{:?}]", task_id);
434            let sender = sender;
435            let mut state_tx_1 = state_tx.clone();
436
437            let task = |task_id: TaskId| async move {
438                let span = tracing_context.attach(tracing::info_span!(
439                    "batch_execute",
440                    task_id = task_id.task_id,
441                    stage_id = task_id.stage_id,
442                    query_id = task_id.query_id,
443                ));
444
445                // We should only pass a reference of sender to execution because we should only
446                // close it after task error has been set.
447                expr_context_scope(
448                    expr_context,
449                    t_1.run(exec, sender, state_tx_1.as_mut()).instrument(span),
450                )
451                .await;
452            };
453
454            if let Err(error) = AssertUnwindSafe(task(task_id.clone()))
455                .rw_catch_unwind()
456                .await
457            {
458                let message = panic_message::get_panic_message(&error);
459                error!(?task_id, error = message, "Batch task panic");
460                notify_panic(&this, state_tx.as_mut(), message).await;
461            }
462        };
463
464        self.runtime.spawn(fut);
465
466        Ok(())
467    }
468
469    /// Change state and notify frontend for task status via streaming GRPC.
470    pub async fn change_state_notify(
471        &self,
472        task_status: TaskStatus,
473        state_tx: Option<&mut StateReporter>,
474        err_str: Option<String>,
475    ) -> Result<()> {
476        self.change_state(task_status);
477        // Notify frontend the task status.
478        if let Some(reporter) = state_tx {
479            reporter
480                .send(TaskInfoResponse {
481                    task_id: Some(self.task_id.to_prost()),
482                    task_status: task_status.into(),
483                    error_message: err_str.unwrap_or("".to_owned()),
484                })
485                .await
486        } else {
487            Ok(())
488        }
489    }
490
491    pub fn change_state(&self, task_status: TaskStatus) {
492        *self.state.lock() = task_status;
493        tracing::debug!(
494            "Task {:?} state changed to {:?}",
495            &self.task_id,
496            task_status
497        );
498    }
499
500    async fn run(
501        &self,
502        root: BoxedExecutor,
503        mut sender: ChanSenderImpl,
504        state_tx: Option<&mut StateReporter>,
505    ) {
506        self.context
507            .batch_metrics()
508            .as_ref()
509            .inspect(|m| m.batch_manager_metrics().task_num.inc());
510        let mut data_chunk_stream = root.execute();
511        let mut state;
512        let mut error = None;
513
514        let mut shutdown_rx = self.shutdown_rx.clone();
515        loop {
516            select! {
517                biased;
518                // `shutdown_rx` can't be removed here to avoid `sender.send(data_chunk)` blocked whole execution.
519                _ = shutdown_rx.cancelled() => {
520                    match self.shutdown_rx.message() {
521                        ShutdownMsg::Abort(e) => {
522                            error = Some(BatchError::Aborted(e));
523                            state = TaskStatus::Aborted;
524                            break;
525                        }
526                        ShutdownMsg::Cancel => {
527                            state = TaskStatus::Cancelled;
528                            break;
529                        }
530                        ShutdownMsg::Init => {
531                            unreachable!("Init message should not be received here!")
532                        }
533                    }
534                }
535                data_chunk = data_chunk_stream.next()=> {
536                    match data_chunk {
537                        Some(Ok(data_chunk)) => {
538                            if let Err(e) = sender.send(data_chunk).await {
539                                match e {
540                                    BatchError::SenderError => {
541                                        // This is possible since when we have limit executor in parent
542                                        // stage, it may early stop receiving data from downstream, which
543                                        // leads to close of channel.
544                                        warn!("Task receiver closed!");
545                                        state = TaskStatus::Finished;
546                                        break;
547                                    }
548                                    x => {
549                                        error!("Failed to send data!");
550                                        error = Some(x);
551                                        state = TaskStatus::Failed;
552                                        break;
553                                    }
554                                }
555                            }
556                        }
557                        Some(Err(e)) => match self.shutdown_rx.message() {
558                            ShutdownMsg::Init => {
559                                // There is no message received from shutdown channel, which means it caused
560                                // task failed.
561                                error!(error = %e.as_report(), "Batch task failed");
562                                error = Some(e);
563                                state = TaskStatus::Failed;
564                                break;
565                            }
566                            ShutdownMsg::Abort(_) => {
567                                error = Some(e);
568                                state = TaskStatus::Aborted;
569                                break;
570                            }
571                            ShutdownMsg::Cancel => {
572                                state = TaskStatus::Cancelled;
573                                break;
574                            }
575                        },
576                        None => {
577                            debug!("Batch task {:?} finished successfully.", self.task_id);
578                            state = TaskStatus::Finished;
579                            break;
580                        }
581                    }
582                }
583            }
584        }
585
586        let error = error.map(Arc::new);
587        self.failure.lock().clone_from(&error);
588        let err_str = error.as_ref().map(|e| e.to_report_string());
589        if let Err(e) = sender.close(error).await {
590            match e {
591                SenderError => {
592                    // This is possible since when we have limit executor in parent
593                    // stage, it may early stop receiving data from downstream, which
594                    // leads to close of channel.
595                    warn!("Task receiver closed when sending None!");
596                }
597                _x => {
598                    error!("Failed to close task output channel: {:?}", self.task_id);
599                    state = TaskStatus::Failed;
600                }
601            }
602        }
603
604        if let Err(e) = self.change_state_notify(state, state_tx, err_str).await {
605            warn!(
606                error = %e.as_report(),
607                "The status receiver in FE has closed so the status push is failed",
608            );
609        }
610
611        self.context
612            .batch_metrics()
613            .as_ref()
614            .inspect(|m| m.batch_manager_metrics().task_num.dec());
615    }
616
617    pub fn abort(&self, err_msg: String) {
618        // No need to set state to be Aborted here cuz it will be set by shutdown receiver.
619        // Stop task execution.
620        if self.shutdown_tx.abort(err_msg) {
621            info!("Abort task {:?} done", self.task_id);
622        } else {
623            debug!("The task has already died before this request.")
624        }
625    }
626
627    pub fn cancel(&self) {
628        if !self.shutdown_tx.cancel() {
629            debug!("The task has already died before this request.");
630        }
631    }
632
633    pub fn get_task_output(&self, output_id: &PbTaskOutputId) -> Result<TaskOutput> {
634        let task_id = TaskId::from(output_id.get_task_id()?);
635        let receiver = self.receivers.lock()[output_id.get_output_id() as usize]
636            .take()
637            .with_context(|| {
638                format!(
639                    "Task{:?}'s output{} has already been taken.",
640                    task_id,
641                    output_id.get_output_id(),
642                )
643            })?;
644        let task_output = TaskOutput {
645            receiver,
646            output_id: output_id.try_into()?,
647            failure: self.failure.clone(),
648        };
649        Ok(task_output)
650    }
651
652    pub fn check_if_running(&self) -> Result<()> {
653        if *self.state.lock() != TaskStatus::Running {
654            bail!("task {:?} is not running", self.get_task_id());
655        }
656        Ok(())
657    }
658
659    pub fn check_if_aborted(&self) -> Result<bool> {
660        match *self.state.lock() {
661            TaskStatus::Aborted => Ok(true),
662            TaskStatus::Finished => bail!("task {:?} has been finished", self.get_task_id()),
663            _ => Ok(false),
664        }
665    }
666
667    /// Check the task status: whether has ended.
668    pub fn is_end(&self) -> bool {
669        let guard = self.state.lock();
670        !(*guard == TaskStatus::Running || *guard == TaskStatus::Pending)
671    }
672}
673
674impl BatchTaskExecution {
675    pub(crate) fn set_heartbeat_join_handle(&self, join_handle: JoinHandle<()>) {
676        *self.heartbeat_join_handle.lock() = Some(join_handle);
677    }
678
679    pub(crate) fn heartbeat_join_handle(&self) -> Option<JoinHandle<()>> {
680        self.heartbeat_join_handle.lock().take()
681    }
682}
683
684#[cfg(test)]
685mod tests {
686    use super::*;
687
688    #[test]
689    fn test_task_output_id_debug() {
690        let task_id = TaskId {
691            task_id: 1,
692            stage_id: 2,
693            query_id: "abc".to_owned(),
694        };
695        let task_output_id = TaskOutputId {
696            task_id,
697            output_id: 3,
698        };
699        assert_eq!(
700            format!("{:?}", task_output_id),
701            "TaskOutputId { query_id: \"abc\", stage_id: 2, task_id: 1, output_id: 3 }"
702        );
703    }
704}
risingwave_batch/task/task_execution.rs

risingwave_batch/task/
task_execution.rs
