Merge pull request #17 from coralogix/more-logging

Catch panics in task execution so jobs do not hang

Author: thinkharderdev

ballista/rust/executor/src/executor_server.rs

@@ -16,9 +16,11 @@
 // under the License.
 
 use ballista_core::BALLISTA_VERSION;
+use futures::FutureExt;
 use std::collections::HashMap;
 use std::convert::TryInto;
 use std::ops::Deref;
+use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tokio::sync::{mpsc, RwLock};
@@ -52,6 +54,7 @@ use tokio::task::JoinHandle;
 
 use crate::as_task_status;
 use crate::cpu_bound_executor::DedicatedExecutor;
+use crate::execution_loop::any_to_string;
 use crate::executor::Executor;
 use crate::shutdown::ShutdownNotifier;
 
@@ -197,8 +200,6 @@ struct ExecutorEnv {
     tx_stop: mpsc::Sender<bool>,
 }
 
-unsafe impl Sync for ExecutorEnv {}
-
 impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T, U> {
     fn new(
         scheduler_to_register: SchedulerGrpcClient<Channel>,
@@ -350,17 +351,26 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
             plan,
         )?;
 
-        let execution_result = self
-            .executor
-            .execute_shuffle_write(
+        let execution_result =
+            match AssertUnwindSafe(self.executor.execute_shuffle_write(
                 task_id.job_id.clone(),
                 task_id.stage_id as usize,
                 task_id.partition_id as usize,
                 shuffle_writer_plan.clone(),
                 task_context,
                 shuffle_output_partitioning,
-            )
-            .await;
+            ))
+            .catch_unwind()
+            .await
+            {
+                Ok(Ok(r)) => Ok(r),
+                Ok(Err(r)) => Err(r),
+                Err(r) => {
+                    error!("Error executing task: {:?}", any_to_string(&r));
+                    Err(BallistaError::Internal(format!("{:#?}", any_to_string(&r))))
+                }
+            };
+
         info!("Done with task {}", task_id_log);
         debug!("Statistics: {:?}", execution_result);
 
@@ -606,13 +616,20 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorGrpc
         } = request.into_inner();
         let task_sender = self.executor_env.tx_task.clone();
         for task in tasks {
-            task_sender
+            if let Err(e) = task_sender
                 .send(CuratorTaskDefinition {
                     scheduler_id: scheduler_id.clone(),
                     task,
                 })
                 .await
-                .unwrap();
+            {
+                let msg = format!(
+                    "Error launching tasks, failed to send to task channel: {:?}",
+                    e
+                );
+                error!("{}", msg);
+                return Err(Status::internal(msg));
+            }
         }
         Ok(Response::new(LaunchTaskResult { success: true }))
     }

ballista/rust/scheduler/src/state/execution_graph.rs

@@ -647,24 +647,30 @@ impl ExecutionGraph {
         for event in events {
             match event {
                 StageEvent::StageResolved(stage_id) => {
+                    info!("stage {}/{} resolved", self.job_id, stage_id);
                     self.resolve_stage(stage_id)?;
                     has_resolved = true;
                 }
                 StageEvent::StageCompleted(stage_id) => {
+                    info!("stage {}/{} completed", self.job_id, stage_id);
                     self.complete_stage(stage_id);
                 }
                 StageEvent::StageFailed(stage_id, err_msg, failed_at) => {
+                    error!("stage {}/{} failed: {}", self.job_id, stage_id, err_msg);
                     job_err_msg = format!("{}{}\n", job_err_msg, &err_msg);
                     failed_at_time = failed_at;
                     self.fail_stage(stage_id, err_msg);
                 }
                 StageEvent::RollBackRunningStage(stage_id) => {
+                    warn!("running stage {}/{} rolled back", self.job_id, stage_id);
                     self.rollback_running_stage(stage_id)?;
                 }
                 StageEvent::RollBackResolvedStage(stage_id) => {
+                    warn!("running stage {}/{} rolled back", self.job_id, stage_id);
                     self.rollback_resolved_stage(stage_id)?;
                 }
                 StageEvent::ReRunCompletedStage(stage_id) => {
+                    info!("re-running completed stage {}/{}", self.job_id, stage_id);
                     self.rerun_completed_stage(stage_id);
                 }
             }

ballista/rust/scheduler/src/state/task_manager.rs

@@ -151,7 +151,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
         let mut events: Vec<QueryStageSchedulerEvent> = vec![];
         for (job_id, statuses) in job_updates {
             let num_tasks = statuses.len();
-            debug!("Updating {} tasks in job {}", num_tasks, job_id);
+            info!("Updating {} tasks in job {}", num_tasks, job_id);
 
             let graph = self.get_active_execution_graph(&job_id).await;
             let job_event = if let Some(graph) = graph {