fix bug that nil in the taskBoxCh

Signed-off-by: Weizhen Wang <[email protected]>
pingcap · Jan 3, 2023 · da506e4 · da506e4
1 parent fe1e900
commit da506e4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/util/gpool/spmc/spmcpool.go b/util/gpool/spmc/spmcpool.go
@@ -216,7 +216,8 @@ func (p *Pool[T, U, C, CT, TF]) ReleaseAndWait() {
 	close(p.stopCh)
 	p.release()
 	for {
-		if p.Running() == 0 && p.heartbeatDone.Load() {
+		// Wait for all workers to exit and all task to be completed.
+		if p.Running() == 0 && p.heartbeatDone.Load() && p.waitingTask.Load() == 0 {
 			return
 		}
 	}

diff --git a/util/gpool/spmc/worker.go b/util/gpool/spmc/worker.go
@@ -55,10 +55,10 @@ func (w *goWorker[T, U, C, CT, TF]) run() {
 		}()
 
 		for f := range w.taskBoxCh {
-			w.pool.subWaitingTask()
 			if f == nil {
 				return
 			}
+			w.pool.subWaitingTask()
 			ctx := f.GetContextFunc().GetContext()
 			if f.GetResultCh() != nil {
 				for t := range f.GetTaskCh() {