From 6f5332753c79adbcfecd384c2f296ae43b556ebc Mon Sep 17 00:00:00 2001
From: "Da K. Ma" <klaus1982.cn@gmail.com>
Date: Mon, 9 Jul 2018 17:17:28 +0800
Subject: [PATCH] Deleted Job when PDB deleted.

Signed-off-by: Da K. Ma <klaus1982.cn@gmail.com>
---
 pkg/scheduler/api/helpers.go          |  4 +++-
 pkg/scheduler/api/job_info.go         |  4 ++++
 pkg/scheduler/cache/event_handlers.go | 19 +++++++++++++++----
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/pkg/scheduler/api/helpers.go b/pkg/scheduler/api/helpers.go
index 8fa8279ac6..2f18c276bd 100644
--- a/pkg/scheduler/api/helpers.go
+++ b/pkg/scheduler/api/helpers.go
@@ -98,5 +98,7 @@ func MergeErrors(errs ...error) error {
 
 // JobTerminated checkes whether job was terminated.
 func JobTerminated(job *JobInfo) bool {
-	return job.SchedSpec == nil && len(job.Tasks) == 0
+	return job.SchedSpec == nil &&
+		job.PDB == nil &&
+		len(job.Tasks) == 0
 }
diff --git a/pkg/scheduler/api/job_info.go b/pkg/scheduler/api/job_info.go
index cd8f321aa8..e3e764ec5c 100644
--- a/pkg/scheduler/api/job_info.go
+++ b/pkg/scheduler/api/job_info.go
@@ -162,6 +162,10 @@ func (ps *JobInfo) SetPDB(pbd *policyv1.PodDisruptionBudget) {
 	ps.PDB = pbd
 }
 
+func (ps *JobInfo) UnsetPDB() {
+	ps.PDB = nil
+}
+
 func (ps *JobInfo) GetTasks(statuses ...TaskStatus) []*TaskInfo {
 	var res []*TaskInfo
 
diff --git a/pkg/scheduler/cache/event_handlers.go b/pkg/scheduler/cache/event_handlers.go
index 7e431878b1..4fb45838fb 100644
--- a/pkg/scheduler/cache/event_handlers.go
+++ b/pkg/scheduler/cache/event_handlers.go
@@ -362,7 +362,6 @@ func (sc *SchedulerCache) deleteSchedulingSpec(ss *arbv1.SchedulingSpec) error {
 	// Unset SchedulingSpec
 	job.UnsetSchedulingSpec()
 
-	// TODO (k82cn): find another way to clean up Job.
 	sc.deleteJob(job)
 
 	return nil
@@ -456,12 +455,24 @@ func (sc *SchedulerCache) setPDB(pdb *policyv1.PodDisruptionBudget) error {
 }
 
 // Assumes that lock is already acquired.
-func (sc *SchedulerCache) updatePDB(oldQueue, newQueue *policyv1.PodDisruptionBudget) error {
-	return sc.setPDB(newQueue)
+func (sc *SchedulerCache) updatePDB(oldPDB, newPDB *policyv1.PodDisruptionBudget) error {
+	return sc.setPDB(newPDB)
 }
 
 // Assumes that lock is already acquired.
-func (sc *SchedulerCache) deletePDB(queue *policyv1.PodDisruptionBudget) error {
+func (sc *SchedulerCache) deletePDB(pdb *policyv1.PodDisruptionBudget) error {
+	jobID := arbapi.JobID(utils.GetController(pdb))
+
+	job, found := sc.Jobs[jobID]
+	if !found {
+		return fmt.Errorf("can not found job %v:%v/%v", jobID, pdb.Namespace, pdb.Name)
+	}
+
+	// Unset SchedulingSpec
+	job.UnsetPDB()
+
+	sc.deleteJob(job)
+
 	return nil
 }