go - Golang 中 context 执行了Cancel，但是<- ctx.Done() 没有执行？

Question

package main

import (
    "context"
    "fmt"
)

func gen(ctx context.Context) <-chan interface{} {
    ch := make(chan interface{})
    go func() {
        n := 0
        for {
            select {
            case <-ctx.Done():
                fmt.Println("done")
            default:
                n += 1
                ch <- n
            }
        }

    }()
    return ch
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())

    for n := range gen(ctx) {
        fmt.Println(n)
        if n == 5 {
            break
        }
    }
    defer cancel()
}

期望 能打印出 "done"

Answer 1

应该是阻塞在 ch<-n 了。

cancel() 如果执行的比较晚，在 cancel 之前第六个 ch<-n 就已经执行了的话，就阻塞了。（main 逻辑相对比较长，比如还有一个 Println ，所以大概率 cancel 会比较晚。

可以试试这个：

package main

import (
    "context"
    "fmt"
)

func gen(ctx context.Context) <-chan interface{} {
    ch := make(chan interface{})
    go func() {
        n := 0
        for {
            select {
            case <-ctx.Done():
                fmt.Println("done")
                close(ch) // 关闭 channel ，range 结束
                return
            default:
                n += 1
                ch <- n
            }
        }

    }()
    return ch
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())

    for n := range gen(ctx) {
        fmt.Println(n)
        if n == 5 {
            cancel()
            // break 不能 break，否则没有从 ch 中的读取，就会一直阻塞在 ch <- n
        }
    }
}

这个依然会打印 6 个，因为 cancel 的时候第六个 ch<-n 已经执行了。

Answer 2

在你的代码中，<-ctx.Done() 没有执行并打印 "done" 的原因是 cancel() 函数调用发生在 for n := range gen(ctx) 循环结束后，而循环结束的条件是 n == 5。由于 cancel() 是在循环结束后被调用，此时 gen 函数中的 goroutine 已经因为 main 函数中的循环结束而不再从 ch 发送数据，因此 select 语句中的 case <-ctx.Done(): 分支没有机会被执行。

为了解决这个问题，你需要确保在 main 函数中，在循环结束之前调用 cancel() 函数，并且让 gen 函数中的 goroutine 有机会响应 ctx.Done() 信号。你可以通过修改 main 函数来实现这一点，如下所示：

package main

import (
    "context"
    "fmt"
    "time"
)

func gen(ctx context.Context) <-chan interface{} {
    ch := make(chan interface{})
    go func() {
        n := 0
        for {
            select {
            case <-ctx.Done():
                fmt.Println("done")
                return // 确保goroutine在接收到取消信号后退出
            default:
                n += 1
                ch <- n
                // 为了避免goroutine过快地发送数据，可以添加一点延迟
                time.Sleep(100 * time.Millisecond)
            }
        }
    }()
    return ch
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())

    done := make(chan bool)
    go func() {
        for n := range gen(ctx) {
            fmt.Println(n)
            if n == 5 {
                cancel() // 在n==5时取消context
                done <- true
                break
            }
        }
    }()

    <-done // 等待goroutine完成
}

在这个修改后的版本中，我添加了一个 done 通道来确保 main 函数等待 gen 函数中的 goroutine 完成其工作。当 n 等于 5 时，cancel() 被调用，并且 done 通道被发送一个值以通知 main 函数可以安全地继续执行。同时，我在 gen 函数的 goroutine 中添加了一个 return 语句，以确保在接收到取消信号后 goroutine 能够退出。此外，我还添加了一个小的延迟来模拟数据生成的过程，以避免在测试时过快地达到 n == 5 的条件。