实现异步并发 worker 队列

 1type LimitedWaitGroup struct {
 2    wg *sync.WaitGroup
 3    ch chan int
 4}
 5
 6func NewLimitedWaitGroup(size int) *LimitedWaitGroup {
 7    return &LimitedWaitGroup{
 8        wg : new(sync.WaitGroup),
 9        ch : make(chan int, size)
10    }
11}
12
13func (w *LimitedWaitGroup) Add(f func()) {
14    // wait if channel is full
15    w.ch <- 1
16    w.wg.Add(1)
17    go func() {
18        defer w.done()
19        f()
20    }()
21}
22
23func (w *LimitedWaitGroup) done() {
24    <-w.ch
25    w.wg.Done()
26}
27
28func (w *LimitedWaitGroup) Wait() {
29    w.wg.Wait()
30}

 1type TaskFunc func() error
 2
 3type task struct {
 4    tf          TaskFunc        // task function
 5    concurrence int             // concurrence of task
 6    ws          *workerSet      // assign value after task distribute to worker.
 7    status      TaskStatus      // store task status.
 8}
 9
10// TaskStatus is the status of task.
11type TaskStatus int

1type worker struct {
2    ws  *workerSet
3    err error
4}

 1// TaskResult is a manager of submitted task.
 2type TaskResult interface {
 3    // get error if task failed.
 4    Err() error
 5    // wait for task done.
 6    Wait()
 7    // get task status.
 8    Status() TaskStatus
 9    // kill task.
10    Kill()
11}

1// workerSet represent a group of task handle workers
2type workerSet struct {
3    task          *task
4    runningWorker atomic.Int32
5    workers       []*worker
6    ctx           context.Context
7    cancel        context.CancelFunc
8    wg            *sync.WaitGroup
9}

 1// Pool is a buffered worker pool
 2type Pool struct {
 3    // TODO: taskQueue should be a linked list, so that we can get the task from the head of the list and put it back to the head.
 4    // If we use a channel as taskQueue, we can't get the task from the head of the list and put it back to the head.
 5    // But make sure that before change it to linked list, we should have the ability run the task in min(taskQueue length, concurrence) goroutines.
 6    taskQueue         chan *task
 7    enqueuedTaskCount atomic.Int32 // count of unhandled tasks
 8    bufferSize        int          // size of taskQueue buffer, means can count of bufferSize task can wait to be handled
 9    maxWorker         int          // count of how many worker run in concurrence
10    workerSets        []*workerSet
11    lock              *sync.Mutex
12    stopFlag          atomic.Bool
13}

 1func (w *worker) run() {
 2    defer w.ws.done()
 3
 4    var ec = make(chan error, 1)
 5    defer close(ec)
 6    go func() {
 7        ec <- w.ws.task.tf()
 8    }()
 9
10    select {
11    case e := <-ec:
12        w.err = e
13    case <-w.ws.ctx.Done():
14    }
15}

 1func (t TaskStatus) Error() string {
 2    return t.String()
 3}
 4
 5func (t TaskStatus) Err() error {
 6    switch t {
 7    case TaskStatusError, TaskStatusQueueFull, TaskStatusTooManyWorker, TaskStatusPoolClosed, TaskStatusKilled:
 8        return t
 9    }
10
11    return nil
12}
13
14func (t TaskStatus) Wait() {
15    // do nothing.
16}
17
18func (t TaskStatus) Status() TaskStatus {
19    return t
20}
21
22func (t TaskStatus) Kill() {
23    // do nothing.
24}

 1func TestPool_SubmitOrEnqueue(t *testing.T) {
 2    p := NewPool(5, 1)
 3    var (
 4        cnt         int
 5        concurrence = 5
 6    )
 7
 8    tf := func() error {
 9        time.Sleep(time.Second)
10        log.Println("hello world")
11        cnt++
12        return nil
13    }
14
15    got := p.Submit(context.Background(), tf, concurrence)
16    if got.Status() != TaskStatusRunning {
17        t.Errorf("SubmitOrEnqueue() = %v, want %v", got.Status(), TaskStatusRunning)
18        return
19    }
20    got.Wait()
21    if cnt != concurrence {
22        t.Errorf("cnt = %v, want %v", cnt, concurrence)
23    }
24    if got := p.Submit(context.Background(), tf, concurrence); got.Status() != TaskStatusRunning {
25        t.Errorf("SubmitOrEnqueue() = %v, want %v", got, TaskStatusRunning)
26        return
27    }
28    if got := p.Submit(context.Background(), tf, concurrence); got.Status() != TaskStatusEnqueue {
29        t.Errorf("SubmitOrEnqueue() = %v, want %v", got.Status(), TaskStatusEnqueue)
30        return
31    }
32    if got := p.Submit(context.Background(), tf, 6); got.Status() != TaskStatusTooManyWorker {
33        t.Errorf("SubmitOrEnqueue() = %v, want %v", got.Status(), TaskStatusTooManyWorker)
34        return
35    }
36    if got := p.Submit(context.Background(), tf, concurrence); got.Status() != TaskStatusQueueFull {
37        t.Errorf("SubmitOrEnqueue() = %v, want %v", got.Status(), TaskStatusQueueFull)
38        return
39    }
40    p.Stop()
41    if got := p.Submit(context.Background(), tf, 1); got.Status() != TaskStatusPoolClosed {
42        t.Errorf("SubmitOrEnqueue() = %v, want %v", got.Status(), TaskStatusPoolClosed)
43        return
44    }
45}