译 | 如何优雅地关闭 Go 中的工作 goroutine

func workerFunc() {
    for event := range jobsChan { // 阻塞直到接收到一个事件，或者该 channel 被关闭。
        // handle the event...
    }
}

func callbackFunc(event int) {
    jobsChan<-event
}

func startConsumer(ctx context.Context) {
    // Loop until a ctx.Done() is received. Note that select{} blocks until either case happens
    for {
            select {
            case event := <-intermediateChan:
                jobsChan <- event
            case _ <- ctx.Done():
                close(jobsChan)
                return             // exit this function so we don't consume anything more from the intermediate chan
            }
    }
}

for event := range jobsChan { // <- on the close(jobsChan), all goroutines waiting for jobs here will exit the for-loop
    // handle the event...
}

func main() {
    ... rest of program ...
    
    termChan := make(chan os.Signal)
    signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)
    <-termChan // Blocks here until either SIGINT or SIGTERM is received.
    // 接下来呢？
}

func main() {
    ctx, cancelFunc := context.WithCancel(ctx.Background())
    // ... some omitted code ...
    
    go startConsumer(ctx) // pass the cancellable context to the consumer function
    
    // ... some more omitted code ...
    <-termChan
    
    cancelFunc() // call the cancelfunc to notify the consumer it's time to shut stuff down.
}

const numberOfWorkers = 4

func main() {
    // ... omitted ...
    wg := &sync.WaitGroup{}
    wg.Add(numberOfWorkers)
    
    // Start [workerPoolSize] workers
    for i := 0; i < workerPoolSize; i++ {
        go workerFunc(wg)
    }
    
    // ... more omitted stuff ...
    
    <-termChan    // wait for SIGINT / SIGTERM
    cancelFunc()  // send the shutdown signal through the context.Context
    wg.Wait()     // program will wait here until all worker goroutines have reported that they're done
    fmt.Println("Workers done, shutting down!")
}

func workerFunc(wg *sync.WaitGroup) {
    defer wg.Done() // Mark this goroutine as done! once the function exits
    for event := range jobsChan {
        // handle the event...
    }
}

$ go run main.go 
  Worker 3 starting
  Worker 2 starting
  Worker 1 starting
  Worker 0 starting
  Worker 3 finished processing job 0
  Worker 0 finished processing job 3
  ^C*********************************     <-- HERE I PRESS CTRL+C
  Shutdown signal received
  *********************************
  Worker 3 finished processing job 4
  Worker 2 finished processing job 1
  Worker 1 finished processing job 2
  Consumer received cancellation signal, closing jobsChan!   <-- Here, the consumer receives the <-ctx.Done()
  Worker 3 finished processing job 6
  Worker 0 finished processing job 5
  Worker 1 finished processing job 8
  Worker 2 finished processing job 7
  Worker 0 finished processing job 10
  Worker 0 interrupted                    <-- Worker 0 has finished job #10, 3 left
  Worker 2 finished processing job 12
  Worker 2 interrupted                    <-- Worker 2 has finished job #12, 2 left
  Worker 3 finished processing job 9
  Worker 3 interrupted                    <-- Worker 3 has finished job #9, 1 left
  Worker 1 finished processing job 11
  Worker 1 interrupted                    <-- Worker 1 has finished job #11, all done
  All workers done, shutting down!

“如果可以处理一个或多个通信，那么选择进行处理的那个 chanel 是通过统一的伪随机选择的。（If one or more of the communications can proceed, a single one that can proceed is chosen via a uniform pseudo-random selection）”。

jobsChan := make(chan int)

$ go run main.go
.... omitted for brevity ....
^C*********************************
Shutdown signal received
*********************************
Worker 3 finished processing job 3
Worker 3 started job 5
Worker 0 finished processing job 4
Worker 0 started job 6
Consumer received cancellation signal, closing jobsChan! <-- again, it may take some time until the consumer is handed <-ctx.Done()
Consumer closed jobsChan
Worker 1 finished processing job 1     <-- From here on, we see that each worker finishes exactly one job before being interrupted.
Worker 1 interrupted
Worker 2 finished processing job 2
Worker 2 interrupted
Worker 0 finished processing job 6
Worker 0 interrupted
Worker 3 finished processing job 5
Worker 3 interrupted
All workers done, shutting down!

package main

import (
    "context"
    "fmt"
    "math/rand"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

const workerPoolSize = 4

func main() {
    // 创建消费者
    consumer := Consumer{
        ingestChan: make(chan int, 1),
        jobsChan:   make(chan int, workerPoolSize),
    }

    // 模拟外部库：每秒发送 10 个事件
    producer := Producer{callbackFunc: consumer.callbackFunc}
    go producer.start()

    // 设置取消 context 和 waitgroup
    ctx, cancelFunc := context.WithCancel(context.Background())
    wg := &sync.WaitGroup{}

    // 传递取消 context 以启动消费者
    go consumer.startConsumer(ctx)

    // 启动 worker，并添加 [workerPoolSize] 到 WaitGroup
    wg.Add(workerPoolSize)
    for i := 0; i < workerPoolSize; i++ {
        go consumer.workerFunc(wg, i)
    }

    // 处理终止信号，并等待 termChan 信号
    termChan := make(chan os.Signal)
    signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)
    
    <-termChan         // 这里阻塞直到接收到信号

    // 处理关闭
    fmt.Println("*********************************\nShutdown signal received\n*********************************")
    cancelFunc()       // 向 context.Context 发送取消信号
    wg.Wait()          // 这里阻塞直至所有 worker 完成
    
    fmt.Println("All workers done, shutting down!")
}

// -- 从这里起，下面是 Consumer！
type Consumer struct {
    ingestChan chan int
    jobsChan   chan int
}

// 每当外部库传递一个事件给我们，就会调用 callbackFunc。
func (c Consumer) callbackFunc(event int) {
    c.ingestChan <- event
}

// workerFunc 启动一个 worker 函数，它会遍历 jobsChan，直到该 channel 关闭。
func (c Consumer) workerFunc(wg *sync.WaitGroup, index int) {
    defer wg.Done()

    fmt.Printf("Worker %d starting\n", index)
    for eventIndex := range c.jobsChan {
        // 模拟工作执行 1 ~ 3 秒
        fmt.Printf("Worker %d started job %d\n", index, eventIndex)
        time.Sleep(time.Millisecond * time.Duration(1000+rand.Intn(2000)))
        fmt.Printf("Worker %d finished processing job %d\n", index, eventIndex)
    }
    fmt.Printf("Worker %d interrupted\n", index)
}

// startConsumer 作为 ingestChan 和 jobsChan 之间的代理，使用 select 语句以支持优雅关闭。
func (c Consumer) startConsumer(ctx context.Context) {
    for {
        select {
        case job := <-c.ingestChan:
            c.jobsChan <- job
        case <-ctx.Done():
            fmt.Println("Consumer received cancellation signal, closing jobsChan!")
            close(c.jobsChan)
            fmt.Println("Consumer closed jobsChan")
            return
        }
    }
}

// -- Producer 模拟一个外部库，每 100ms 有新数据时simulates an external library that invokes the 
// 它会调用注册的回调函数。
type Producer struct {
    callbackFunc func(event int)
}
func (p Producer) start() {
    eventIndex := 0
    for {
        p.callbackFunc(eventIndex)
        eventIndex++
        time.Sleep(time.Millisecond * 100)
    }
}