golang 多个worker正常关闭的示例

代码如下，如有问题请联系 baibaibai_000@163.com

package work_test

import (
    "math/rand"
    "runtime"
    "sync"
    "testing"
    "time"
)

// TestWorker
// test目的：展示当多个worker同时处理一个通道的任务，被关闭时，通道剩余的任务需要被全部处理
func TestWorker(t *testing.T) {

    rand.Seed(time.Now().UnixNano())

    // 虚拟的处理task的函数
    work := func(t int) {

    }

    // 任务通道
    taskChan := make(chan int, 1024)

    // 关闭
    closeChan := make(chan struct{})
    wg := sync.WaitGroup{}

    // worker个数
    workerNum := runtime.NumCPU()

    // 起多个worker来处理taskChan的事务
    for i := 0; i < workerNum; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()

        LOOPING:
            for {
                select {
                case t := <-taskChan: //处理task
                    work(t)
                case <-closeChan: //收到外界关闭的事件
                    // 继续处理taskChan剩余的task，直到通道为空
                    for hasTask := true; hasTask; {
                        select {
                        case t := <-taskChan:
                            work(t)
                        default:
                            hasTask = false
                        }
                    }
                    break LOOPING
                }
            }

        }()

    }

    // 填充任务
    go func() {
        ticker := time.NewTicker(time.Millisecond * time.Duration(1))
        defer ticker.Stop()

        for {
            select {
            case <-ticker.C:
                taskChan <- 1
            case <-closeChan:
                return
            }
        }

    }()

    // 随机个事件后关闭
    time.Sleep(time.Duration(2+rand.Int()%4) * time.Second)
    close(closeChan)

    wg.Wait()

    // 如果taskChan还有没处理的任务就失败了
    if len(taskChan) > 0 {
        t.FailNow()
    }

}