[Refactor] ♻️ Avoid Unlock Prematurely When Apply Wrong at Resource NoEnough Which Can Lead to Concurrency Issues #5

Author: Harry-zklcdc

internal/schedulers/cpuscheduler.go

@@ -80,6 +80,7 @@ func (cs *cpuScheduler) Apply(num int) (string, error) {
 	}
 
 	cs.Lock()
+	defer cs.Unlock()
 
 	keys := make([]int, 0, len(cs.CpuStatusMap))
 	for k := range cs.CpuStatusMap {
@@ -103,29 +104,36 @@ func (cs *cpuScheduler) Apply(num int) (string, error) {
 	}
 
 	if len(applyCpus) < num {
-		cs.Unlock()
-		cs.Restore(applyCpus)
+		cs.restore(applyCpus)
 		return "", xerrors.NewCpuNotEnoughError()
 	}
 
 	cpuSet := strings.Trim(strings.Join(applyCpus, ","), ",")
 
-	cs.Unlock()
 	go cs.putToEtcd()
 
 	return cpuSet, nil
 }
 
 func (cs *cpuScheduler) Restore(cpuSet []string) error {
+
 	cs.Lock()
 	defer cs.Unlock()
 
+	err := cs.restore(cpuSet)
+	if err != nil {
+		return errors.Wrap(err, "restore failed")
+	}
+	go cs.putToEtcd()
+
+	return nil
+}
+
+func (cs *cpuScheduler) restore(cpuSet []string) error {
 	for _, cpu := range cpuSet {
 		cs.CpuStatusMap[cpu] = 0
 	}
 
-	go cs.putToEtcd()
-
 	return nil
 }
 

internal/schedulers/gpuscheduler.go

@@ -86,6 +86,7 @@ func (gs *gpuScheduler) Apply(num int) ([]string, error) {
 	}
 
 	gs.Lock()
+	defer gs.Unlock()
 
 	var availableGpus []string
 	for k, v := range gs.GpuStatusMap {
@@ -99,12 +100,10 @@ func (gs *gpuScheduler) Apply(num int) ([]string, error) {
 	}
 
 	if len(availableGpus) < num {
-		gs.Unlock()
-		gs.Restore(availableGpus)
+		gs.restore(availableGpus)
 		return nil, xerrors.NewGpuNotEnoughError()
 	}
 
-	gs.Unlock()
 	go gs.putToEtcd()
 
 	return availableGpus, nil
@@ -119,11 +118,19 @@ func (gs *gpuScheduler) Restore(gpus []string) {
 	gs.Lock()
 	defer gs.Unlock()
 
+	gs.restore(gpus)
+
+	go gs.putToEtcd()
+}
+
+func (gs *gpuScheduler) restore(gpus []string) {
+	if len(gpus) <= 0 || len(gpus) > gs.AvailableGpuNums {
+		return
+	}
+
 	for _, gpu := range gpus {
 		gs.GpuStatusMap[gpu] = 0
 	}
-
-	go gs.putToEtcd()
 }
 
 func (gs *gpuScheduler) serialize() *string {

internal/schedulers/portscheduler.go

@@ -113,11 +113,15 @@ func (ps *portScheduler) Restore(ports []string) {
 	ps.Lock()
 	defer ps.Unlock()
 
+	ps.restore(ports)
+
+	go ps.putToEtcd()
+}
+
+func (ps *portScheduler) restore(ports []string) {
 	for _, port := range ports {
 		delete(ps.UsedPortSet, port)
 	}
-
-	go ps.putToEtcd()
 }
 
 func (ps *portScheduler) serialize() *string {

internal/schedulers/scheduler.go

@@ -5,4 +5,5 @@ type Scheduler interface {
 	Restore([]string)
 	serialize() *string
 	putToEtcd()
+	restore([]string)
 }