go 使用 context 退出子协程 goroutine

一、案例

先看下面一段代码，子协程要在 100个小时候退出，导致主协程一直卡在那边。

func TestCancel1(t *testing.T) {
    wg := sync.WaitGroup{}
    wg.Add(1)
    go func() {
        t.Log("goroutine 1")
        time.Sleep(100 * time.Hour)
        wg.Done()
    }()
    wg.Wait()
    t.Log("goroutine main")
}

输出：

可见主协程一直 等待状态，没发退出。

二、使用context.WithTimeout，给协程一个退出时间

我们在下面的代码里，写了一个context.WithTimeout给3秒，3秒后这个程序没有执行完就退出。

func TestTimeout(t *testing.T) {
    wg := sync.WaitGroup{}
    wg.Add(1)

    parentContext := context.Background()
    cancelCtx, cancelFunc := context.WithTimeout(parentContext, time.Second*3)
    go func(ctx context.Context) {
        for {
            t.Log("goroutine 1")
            select {
            case <-cancelCtx.Done():
                wg.Done()
                return
            case <-time.After(time.Hour * 100):
                wg.Done()
            }
        }

    }(cancelCtx)
    wg.Wait()
    cancelFunc()
    t.Log("goroutine main")
}

三、使用context.WithCancel，在子协程退出

func TestCancel2(t *testing.T) {
    wg := sync.WaitGroup{}
    wg.Add(1)

    parentContext := context.Background()
    cancelCtx, cancelFunc := context.WithCancel(parentContext)
    go func(ctx context.Context) {
        for {
            select {
            case <-cancelCtx.Done():
                return
            case <-time.After(time.Hour * 100):
                cancelFunc()
                wg.Done()
            case <-time.After(time.Second * 3):
                cancelFunc()
                wg.Done()
            }
            t.Log("goroutine 1")
        }

    }(cancelCtx)
    wg.Wait()
    t.Log("goroutine main")
}

输出：

goroutine 1
goroutine main