Bootstrap

Golang 并发编程:Context 包的使用与并发控制

一、简介

在并发编程中,任务管理和资源控制是非常重要的,而 Golang 的 context 为我们提供了一种优雅的方式来传递取消信号超时控制Context 用于在多个 Goroutine 之间传递上下文信息,避免 Goroutine 无法按需停止而导致资源浪费。

本篇博客将详细介绍 context 包的用法,并通过实例讲解如何在超时、取消任务多 Goroutine 协作场景中使用它。


二、Context 的基本概念

Context 是一种携带取消信号、截止时间(超时)和元数据的上下文对象,主要用于父 Goroutine 与子 Goroutine 的协作。它通过层级化的结构来管理多个并发任务。

1. context 包常用函数
  • context.Background():创建根上下文,通常用于程序入口。
  • context.TODO():占位符上下文,表示未来会替换为实际上下文。
  • context.WithCancel(parent Context):创建带取消功能的子上下文。
  • context.WithTimeout(parent Context, timeout time.Duration):创建带超时功能的子上下文。
  • context.WithDeadline(parent Context, deadline time.Time):基于指定的截止时间创建上下文。
  • context.WithValue(parent Context, key, value interface{}):传递携带额外数据的上下文。

三、Context 的基本用法

1. WithCancel:取消任务的上下文

示例:使用 WithCancel 取消 Goroutine

package main

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

func worker(ctx context.Context, id int) {
    for {
        select {
        case <-ctx.Done(): // 接收取消信号
            fmt.Printf("Worker %d stopped\n", id)
            return
        default:
            fmt.Printf("Worker %d is working...\n", id)
            time.Sleep(time.Second)
        }
    }
}

func main() {
    ctx, cancel := context.WithCancel(context.Background()) // 创建可取消的上下文

    for i := 1; i <= 3; i++ {
        go worker(ctx, i)
    }

    time.Sleep(3 * time.Second) // 模拟主 Goroutine 的其他工作
    fmt.Println("Cancelling all workers...")
    cancel() // 发送取消信号

    time.Sleep(1 * time.Second) // 等待所有 Goroutine 退出
    fmt.Println("All workers stopped.")
}

输出:

Worker 1 is working...
Worker 2 is working...
Worker 3 is working...
...
Cancelling all workers...
Worker 1 stopped
Worker 2 stopped
Worker 3 stopped
All workers stopped.

解析:

  • context.WithCancel 创建的上下文可以通过调用 cancel() 发送取消信号,从而优雅地停止所有子 Goroutine。

四、超时控制:WithTimeoutWithDeadline

1. 使用 WithTimeout 控制任务超时

示例:在 2 秒内完成任务,否则超时退出

package main

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

func worker(ctx context.Context) {
    select {
    case <-time.After(3 * time.Second): // 模拟长时间任务
        fmt.Println("Task completed")
    case <-ctx.Done(): // 接收超时信号
        fmt.Println("Task timed out")
    }
}

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) // 设置 2 秒超时
    defer cancel() // 确保资源释放

    go worker(ctx)

    time.Sleep(4 * time.Second) // 等待任务完成或超时
}

输出:

Task timed out

解析:

  • 通过 context.WithTimeout 创建的上下文,2 秒内未完成任务时自动发送取消信号。
  • 超时上下文可避免 Goroutine 无限期运行,帮助更好地管理资源。
2. 使用 WithDeadline 设定截止时间

WithDeadlineWithTimeout 类似,只是使用具体的时间点来控制超时。


五、传递上下文中的数据:WithValue

有时,我们需要在多个 Goroutine 之间传递一些元数据。WithValue 允许我们将键值对存入上下文,并在子 Goroutine 中访问。

示例:传递用户信息

package main

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

func greetUser(ctx context.Context) {
    if user, ok := ctx.Value("user").(string); ok {
        fmt.Printf("Hello, %s!\n", user)
    } else {
        fmt.Println("No user found.")
    }
}

func main() {
    ctx := context.WithValue(context.Background(), "user", "Alice") // 在上下文中存入用户信息
    go greetUser(ctx)

    time.Sleep(1 * time.Second) // 确保 Goroutine 执行完毕
}

输出:

Hello, Alice!

解析:

  • WithValue 允许我们为上下文设置键值对,便于在多 Goroutine 间传递数据。

注意:

  • 不建议用 WithValue 传递重要的控制信息,例如取消信号或超时。

六、Context 的应用场景

  1. API 请求的超时控制:确保 HTTP 请求不会无限期等待。
  2. 任务取消:当用户主动取消某个操作时,通知相关的 Goroutine 停止工作。
  3. 传递元数据:例如在服务链路中传递用户身份、请求 ID 等信息。

七、完整示例:多任务协作控制

示例:启动多个任务,随时可取消所有任务

package main

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

func worker(ctx context.Context, id int, wg *sync.WaitGroup) {
    defer wg.Done()

    for {
        select {
        case <-ctx.Done():
            fmt.Printf("Worker %d stopped\n", id)
            return
        default:
            fmt.Printf("Worker %d is processing...\n", id)
            time.Sleep(500 * time.Millisecond)
        }
    }
}

func main() {
    var wg sync.WaitGroup
    ctx, cancel := context.WithCancel(context.Background()) // 创建上下文

    for i := 1; i <= 3; i++ {
        wg.Add(1)
        go worker(ctx, i, &wg)
    }

    time.Sleep(2 * time.Second)
    fmt.Println("Cancelling all workers...")
    cancel() // 取消所有任务

    wg.Wait() // 等待所有任务完成
    fmt.Println("All workers stopped.")
}

输出:

Worker 1 is processing...
Worker 2 is processing...
Worker 3 is processing...
...
Cancelling all workers...
Worker 1 stopped
Worker 2 stopped
Worker 3 stopped
All workers stopped.

八、小结

  1. Context 用于并发控制:在 Goroutine 中传递取消信号、超时信号或携带元数据。
  2. 超时控制与资源管理:使用 WithTimeoutWithCancel 及时终止任务,避免资源浪费。
  3. 多 Goroutine 协作:通过 Context 实现多个 Goroutine 之间的优雅通信。

;