Add grpc service for the scheduler to make it able for the job shuffle data be cleaned up to be triggered by client explicitly

Author: kyotoYaho

ballista/core/proto/ballista.proto

@@ -948,6 +948,13 @@ message CancelJobResult {
   bool cancelled = 1;
 }
 
+message CleanJobDataParams {
+  string job_id = 1;
+}
+
+message CleanJobDataResult {
+}
+
 message LaunchTaskParams {
   // Allow to launch a task set to an executor at once
   repeated TaskDefinition tasks = 1;
@@ -1014,6 +1021,8 @@ service SchedulerGrpc {
   rpc ExecutorStopped (ExecutorStoppedParams) returns (ExecutorStoppedResult) {}
 
   rpc CancelJob (CancelJobParams) returns (CancelJobResult) {}
+
+  rpc CleanJobData (CleanJobDataParams) returns (CleanJobDataResult) {}
 }
 
 service ExecutorGrpc {

ballista/scheduler/src/scheduler_server/event.rs

@@ -40,6 +40,7 @@ pub enum QueryStageSchedulerEvent {
     JobRunningFailed(String, String),
     JobUpdated(String),
     JobCancel(String),
+    JobDataClean(String),
     TaskUpdating(String, Vec<TaskStatus>),
     ReservationOffering(Vec<ExecutorReservation>),
     ExecutorLost(String, Option<String>),

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -24,12 +24,12 @@ use std::convert::TryInto;
 use ballista_core::serde::protobuf::executor_registration::OptionalHost;
 use ballista_core::serde::protobuf::scheduler_grpc_server::SchedulerGrpc;
 use ballista_core::serde::protobuf::{
-    executor_status, CancelJobParams, CancelJobResult, ExecuteQueryParams,
-    ExecuteQueryResult, ExecutorHeartbeat, ExecutorStatus, ExecutorStoppedParams,
-    ExecutorStoppedResult, GetFileMetadataParams, GetFileMetadataResult,
-    GetJobStatusParams, GetJobStatusResult, HeartBeatParams, HeartBeatResult,
-    PollWorkParams, PollWorkResult, RegisterExecutorParams, RegisterExecutorResult,
-    UpdateTaskStatusParams, UpdateTaskStatusResult,
+    executor_status, CancelJobParams, CancelJobResult, CleanJobDataParams,
+    CleanJobDataResult, ExecuteQueryParams, ExecuteQueryResult, ExecutorHeartbeat,
+    ExecutorStatus, ExecutorStoppedParams, ExecutorStoppedResult, GetFileMetadataParams,
+    GetFileMetadataResult, GetJobStatusParams, GetJobStatusResult, HeartBeatParams,
+    HeartBeatResult, PollWorkParams, PollWorkResult, RegisterExecutorParams,
+    RegisterExecutorResult, UpdateTaskStatusParams, UpdateTaskStatusResult,
 };
 use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata};
 use ballista_core::serde::AsExecutionPlan;
@@ -545,6 +545,30 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
             })?;
         Ok(Response::new(CancelJobResult { cancelled: true }))
     }
+
+    async fn clean_job_data(
+        &self,
+        request: Request<CleanJobDataParams>,
+    ) -> Result<Response<CleanJobDataResult>, Status> {
+        let job_id = request.into_inner().job_id;
+        info!("Received clean data request for job {}", job_id);
+
+        self.query_stage_event_loop
+            .get_sender()
+            .map_err(|e| {
+                let msg = format!("Get query stage event loop error due to {:?}", e);
+                error!("{}", msg);
+                Status::internal(msg)
+            })?
+            .post_event(QueryStageSchedulerEvent::JobDataClean(job_id))
+            .await
+            .map_err(|e| {
+                let msg = format!("Post to query stage event loop error due to {:?}", e);
+                error!("{}", msg);
+                Status::internal(msg)
+            })?;
+        Ok(Response::new(CleanJobDataResult {}))
+    }
 }
 
 #[cfg(all(test, feature = "sled"))]

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -224,6 +224,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                     .cancel_running_tasks(tasks)
                     .await?
             }
+            QueryStageSchedulerEvent::JobDataClean(job_id) => {
+                let executor_manager = self.state.executor_manager.clone();
+                tokio::spawn(async move {
+                    executor_manager.clean_up_job_data(job_id).await;
+                });
+            }
         }
 
         Ok(())