golang每日一题：context、goroutine相关

24/11/28

题目描述

整体的超时控制
A子协程发送数字 0-9
B子协程计算a发来数字的平方
主线程打印输出最后的平方数

go 复制代码

package main

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

func TestGoroutine(t *testing.T) {
    // 设置超时
    ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
    defer cancel() // 在函数结束时取消 context

    sendNum := make(chan int)
    resMul := make(chan int)

    // 启动第一个 goroutine，用于发送数据
    go func(ctx context.Context) {
        defer close(sendNum) // 在发送完数据后关闭 sendNum 通道
        for i := 0; i < 200; i++ {
            select {
            case sendNum <- i:
            case <-ctx.Done(): // 如果上下文被取消或超时，退出
                return
            }
        }
    }(ctx)

    // 启动第二个 goroutine，用于接收并处理数据
    go func(ctx context.Context) {
        defer close(resMul) // 在处理完数据后关闭 resMul 通道
        for {
            select {
            case num := <-sendNum:
                resMul <- num * num
            case <-ctx.Done(): // 如果上下文被取消或超时，退出
                return
            }
        }
    }(ctx)

    // 主 goroutine 处理结果
    for {
        select {
        case tem := <-resMul:
            fmt.Printf("resMul = %v\n", tem)
        case <-ctx.Done():
            fmt.Println(ctx.Err()) // 打印超时或取消的错误信息
            fmt.Println("time out...")
            return
        }
    }
}

改进：

defer cancel()：在函数结束时调用 cancel() 来释放上下文资源。这样做可以确保 context 在不再需要时被清理。
关闭通道：在发送数据完毕后，我们显式地关闭了 sendNum 通道。这有助于避免死锁，并且通知接收方数据已发送完毕。类似地，resMul 通道也应该在最后关闭。
避免无限循环：通过在 select 中监听 ctx.Done()，在 context 超时或取消时退出 goroutine，而不需要使用额外的 flag 控制循环。
更优雅的退出机制：通过 ctx.Done() 监听上下文取消，确保所有 goroutine 在超时或取消时能够正确退出。