resourcemanager: fix TaskController.Stop() can't make producer exit in spmcpool

resourcemanager/pooltask/BUILD.bazel

@@ -10,7 +10,10 @@ go_library(
  ],
  importpath = "github.com/pingcap/tidb/resourcemanager/pooltask",
  visibility = ["//visibility:public"],
- deps = ["@org_uber_go_atomic//:atomic"],
+ deps = [
+ "//util/channel",
+ "@org_uber_go_atomic//:atomic",
+ ],
 )
 
 go_test(

resourcemanager/pooltask/task.go

@@ -17,6 +17,8 @@ package pooltask
 import (
  "sync"
  "sync/atomic"
+
+ "github.com/pingcap/tidb/util/channel"
 )
 
 // Context is a interface that can be used to create a context.
@@ -127,35 +129,44 @@ type GPool[T any, U any, C any, CT any, TF Context[CT]] interface {
 
 // TaskController is a controller that can control or watch the pool.
 type TaskController[T any, U any, C any, CT any, TF Context[CT]] struct {
- pool GPool[T, U, C, CT, TF]
- close chan struct{}
- wg *sync.WaitGroup
- taskID uint64
- resultCh chan U
+ pool GPool[T, U, C, CT, TF]
+ productExitCh chan struct{}
+ wg *sync.WaitGroup
+ taskID uint64
+ resultCh chan U
+ inputCh chan Task[T]
 }
 
 // NewTaskController create a controller to deal with pooltask's status.
-func NewTaskController[T any, U any, C any, CT any, TF Context[CT]](p GPool[T, U, C, CT, TF], taskID uint64, closeCh chan struct{}, wg *sync.WaitGroup, resultCh chan U) TaskController[T, U, C, CT, TF] {
+func NewTaskController[T any, U any, C any, CT any, TF Context[CT]](p GPool[T, U, C, CT, TF], taskID uint64, productExitCh chan struct{}, wg *sync.WaitGroup, inputCh chan Task[T], resultCh chan U) TaskController[T, U, C, CT, TF] {
  return TaskController[T, U, C, CT, TF]{
- pool: p,
- taskID: taskID,
- close: closeCh,
- wg: wg,
- resultCh: resultCh,
+ pool: p,
+ taskID: taskID,
+ productExitCh: productExitCh,
+ wg: wg,
+ resultCh: resultCh,
+ inputCh: inputCh,
  }
 }
 
 // Wait is to wait the pool task to stop.
 func (t *TaskController[T, U, C, CT, TF]) Wait() {
- <-t.close
  t.wg.Wait()
  close(t.resultCh)
  t.pool.DeleteTask(t.taskID)
 }
 
 // Stop is to send stop command to the task. But you still need to wait the task to stop.
 func (t *TaskController[T, U, C, CT, TF]) Stop() {
+ close(t.productExitCh)
+ // Clear all the task in the task queue and mark all task complete.
+ // so that ```t.Wait``` is able to close resultCh
+ for range t.inputCh {
+ t.wg.Done()
+ }
  t.pool.StopTask(t.TaskID())
+ // Clear the resultCh to avoid blocking the consumer put result into the channel and cannot exit.
+ channel.Clear(t.resultCh)
 }
 
 // TaskID is to get the task id.

util/gpool/spmc/option.go

@@ -18,6 +18,8 @@ import (
  "time"
 )
 
+const defaultTaskChanLen = 1
+
 // Option represents the optional function.
 type Option func(opts *Options)
 
@@ -103,8 +105,8 @@ func loadTaskOptions(options ...TaskOption) *TaskOptions {
  if opts.ResultChanLen == 0 {
  opts.ResultChanLen = uint64(opts.Concurrency)
  }
- if opts.ResultChanLen == 0 {
- opts.ResultChanLen = uint64(opts.Concurrency)
+ if opts.TaskChanLen == 0 {
+ opts.TaskChanLen = defaultTaskChanLen
  }
  return opts
 }

util/gpool/spmc/spmcpool.go

@@ -219,7 +219,6 @@ func (p *Pool[T, U, C, CT, TF]) release() {
  // There might be some callers waiting in retrieveWorker(), so we need to wake them up to prevent
  // those callers blocking infinitely.
  p.cond.Broadcast()
- close(p.taskCh)
 }
 
 func isClose(exitCh chan struct{}) bool {
@@ -260,9 +259,9 @@ func (p *Pool[T, U, C, CT, TF]) AddProduceBySlice(producer func() ([]T, error),
  taskID := p.NewTaskID()
  var wg sync.WaitGroup
  result := make(chan U, opt.ResultChanLen)
- closeCh := make(chan struct{})
+ productExitCh := make(chan struct{})
  inputCh := make(chan pooltask.Task[T], opt.TaskChanLen)
- tc := pooltask.NewTaskController[T, U, C, CT, TF](p, taskID, closeCh, &wg, result)
+ tc := pooltask.NewTaskController[T, U, C, CT, TF](p, taskID, productExitCh, &wg, inputCh, result)
  p.taskManager.RegisterTask(taskID, int32(opt.Concurrency))
  for i := 0; i < opt.Concurrency; i++ {
  err := p.run()
@@ -274,15 +273,19 @@ func (p *Pool[T, U, C, CT, TF]) AddProduceBySlice(producer func() ([]T, error),
  p.taskManager.AddSubTask(taskID, &taskBox)
  p.taskCh <- &taskBox
  }
+ wg.Add(1)
  go func() {
  defer func() {
  if r := recover(); r != nil {
  logutil.BgLogger().Error("producer panic", zap.Any("recover", r), zap.Stack("stack"))
  }
- close(closeCh)
  close(inputCh)
+ wg.Done()
  }()
  for {
+ if isClose(productExitCh) {
+ return
+ }
  tasks, err := producer()
  if err != nil {
  if errors.Is(err, gpool.ErrProducerClosed) {
@@ -310,10 +313,10 @@ func (p *Pool[T, U, C, CT, TF]) AddProducer(producer func() (T, error), constArg
  taskID := p.NewTaskID()
  var wg sync.WaitGroup
  result := make(chan U, opt.ResultChanLen)
- closeCh := make(chan struct{})
+ productExitCh := make(chan struct{})
  inputCh := make(chan pooltask.Task[T], opt.TaskChanLen)
  p.taskManager.RegisterTask(taskID, int32(opt.Concurrency))
- tc := pooltask.NewTaskController[T, U, C, CT, TF](p, taskID, closeCh, &wg, result)
+ tc := pooltask.NewTaskController[T, U, C, CT, TF](p, taskID, productExitCh, &wg, inputCh, result)
  for i := 0; i < opt.Concurrency; i++ {
  err := p.run()
  if err == gpool.ErrPoolClosed {
@@ -324,15 +327,19 @@ func (p *Pool[T, U, C, CT, TF]) AddProducer(producer func() (T, error), constArg
  p.taskManager.AddSubTask(taskID, &taskBox)
  p.taskCh <- &taskBox
  }
+ wg.Add(1)
  go func() {
  defer func() {
  if r := recover(); r != nil {
  logutil.BgLogger().Error("producer panic", zap.Any("recover", r), zap.Stack("stack"))
  }
- close(closeCh)
  close(inputCh)
+ wg.Done()
  }()
  for {
+ if isClose(productExitCh) {
+ return
+ }
  task, err := producer()
  if err != nil {
  if errors.Is(err, gpool.ErrProducerClosed) {

util/gpool/spmc/spmcpool_test.go

@@ -53,7 +53,7 @@ func TestPool(t *testing.T) {
  }
  }
  // add new task
- resultCh, control := pool.AddProducer(pfunc, myArgs, pooltask.NilContext{}, WithConcurrency(4))
+ resultCh, control := pool.AddProducer(pfunc, myArgs, pooltask.NilContext{}, WithConcurrency(5))
 
  var count atomic.Uint32
  var wg sync.WaitGroup
@@ -112,12 +112,55 @@ func TestStopPool(t *testing.T) {
  require.Greater(t, result, 10)
  }
  }()
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ control.Stop()
+ }()
+ // Waiting task finishing
+ control.Wait()
+ wg.Wait()
+ // close pool
+ pool.ReleaseAndWait()
+}
+
+func TestStopPoolWithSlice(t *testing.T) {
+ type ConstArgs struct {
+ a int
+ }
+ myArgs := ConstArgs{a: 10}
+ // init the pool
+ // input type， output type, constArgs type
+ pool, err := NewSPMCPool[int, int, ConstArgs, any, pooltask.NilContext]("TestStopPoolWithSlice", 3, rmutil.UNKNOWN)
+ require.NoError(t, err)
+ pool.SetConsumerFunc(func(task int, constArgs ConstArgs, ctx any) int {
+ return task + constArgs.a
+ })
+
+ exit := make(chan struct{})
+
+ pfunc := func() ([]int, error) {
+ select {
+ case <-exit:
+ return nil, gpool.ErrProducerClosed
+ default:
+ return []int{1, 2, 3}, nil
+ }
+ }
+ // add new task
+ resultCh, control := pool.AddProduceBySlice(pfunc, myArgs, pooltask.NilContext{}, WithConcurrency(4))
+
+ var wg sync.WaitGroup
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ for result := range resultCh {
+ require.Greater(t, result, 10)
+ control.Stop()
+ }
+ }()
  // Waiting task finishing
- control.Stop()
- close(exit)
  control.Wait()
- // it should pass. Stop can be used after the pool is closed. we should prevent it from panic.
- control.Stop()
  wg.Wait()
  // close pool
  pool.ReleaseAndWait()
@@ -191,9 +234,12 @@ func testTunePool(t *testing.T, name string) {
  for n := pool.Cap(); n > 1; n-- {
  downclockPool(t, pool, tid)
  }
-
- // exit test
- close(exit)
+ wg.Add(1)
+ go func() {
+ // exit test
+ control.Stop()
+ wg.Done()
+ }()
  control.Wait()
  wg.Wait()
  // close pool