Go 语言中的 context 机制详解

在 Go 语言中，处理并发任务的一个常见需求是在任务之间传递取消信号和截止时间。此外，还可能需要在任务之间传递一些上下文相关的数据。context 包正是为此而设计的，它提供了一种在并发操作中传递这种信息的方法。

context 是 Go 标准库的一部分，它为请求提供了生命周期和取消信号。context 对象可以携带任意的数据，比如身份认证信息、截止时间和取消信号等。

• 管理生命周期：当一个请求或一个工作流开始时，创建一个 context 对象，并在结束时取消它。
• 传递截止时间：如果一个操作需要在一定时间内完成，可以将这个时间限制作为 context 的一部分传递下去。
• 传递取消信号：当一个操作不再需要继续执行时，可以通过取消 context 来通知所有依赖于它的协程。
• 携带额外数据：context 可以用来传递请求特有的信息，如身份验证 token、跟踪 ID 等。

context 包提供了几种创建 context 的方式：

• context.Background()：返回一个空背景 context，它永远不会被取消。
• context.TODO()：用于初始化阶段，表示“我还没有想好用什么 context。
• 创建一个默认的请求上下文。
• 使用 context.Background() 创建一个默认的上下文，作为其他派生上下文的根上下文。
• 返回值: ctx 是创建的上下文对象。ctx := context.Background()。
• 带有截止时间的 context，可以使用 WithDeadline 或 WithTimeout 来创建带有截止时间的 context，它永远不会被取消。
• 创建一个带有5秒超时限制的上下文 context：

ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(5*time.Second))
defer cancel() // 在函数结束前取消此上下文，避免资源泄漏

// 使用 select 语句来监听两个事件
select {
case <-ctx.Done(): // 如果上下文因超时而完成
    fmt.Println("请求超时") // 打印超时信息
case res := <-doWork(ctx): // 如果 doWork 函数返回了结果
    fmt.Println("结果:", res) // 打印结果
}

如果某个操作不再需要执行，可以通过调用cancel函数来取消 context。

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

//调用 cancel 函数来取消context。
go func() {
    time.Sleep(2 * time.Second)
    cancel()
}()

select {
case <-ctx.Done():
    fmt.Println("请求被取消")
case res := <-doWork(ctx):
    fmt.Println("结果:", res)
}

通过 WithValue 函数可以在 context 中存储键值对。

ctx := context.WithValue(context.Background(), "key", "value")

val := ctx.Value("key")
fmt.Println(val) // 输出: value

• 示例：HTTP 服务端使用 context，在 HTTP 服务器中，每个请求都可以与一个 context 相关联。当请求到达时，可以创建一个带有截止时间的 context ，并将其传递给处理函数。

func handler(w http.ResponseWriter, r *http.Request) {
    ctx, cancel := context.WithTimeout(r.context(), 30*time.Second)
    defer cancel()
    // 处理请求逻辑
    doSomething(ctx)
}

func doSomething(ctx context.context) {
    select {
    case <-ctx.Done():
        fmt.Println("请求超时")
    default:
        // 执行业务逻辑
    }
}

• 示例：数据库查询使用 context，在执行数据库查询时，可以使用 context 来控制查询的超时时间。

// 打开与 PostgreSQL 数据库的连接
// 使用 sql.Open 函数根据提供的配置打开数据库连接。

`db := sql.Open("postgres", "user=postgres dbname=test sslmode=disable")`

// 创建带有超时设置的请求上下文
// 使用 context.WithTimeout 根据背景上下文创建一个带有 5 秒超时设置的上下文。
// 此上下文用于控制数据库操作的最大持续时间。
// 返回值: ctx 是创建的上下文，cancel 是用于取消上下文的函数。

`ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)`

// 延迟取消上下文以避免资源泄漏。

`defer cancel()`

// 执行数据库查询
// 使用 db.QueryContext 方法执行 SQL 查询："SELECT * FROM users"。
// 传递上下文以控制查询操作的持续时间。
// 返回值: rows 是查询的结果集，err 是执行期间的错误信息（如果有的话）。

`rows, err := db.QueryContext(ctx, "SELECT * FROM users")`

// 处理执行错误，如果发生错误则记录致命错误。
if err != nil {
	log.Fatal(err)
}
// 延迟关闭结果集以避免资源泄漏。
defer rows.Close()

// 遍历查询结果并打印用户 ID 和名称
// 在结果集中迭代每一行数据，并通过 rows.Scan 将每行的数据读取到 id 和 name 变量中。
// 如果读取过程中出现错误，则记录致命错误。
for rows.Next() {
	var id int
	var name string
	if err := rows.Scan(&id, &name); err != nil {
		log.Fatal(err)
	}
	// 打印当前行的用户 ID 和名称。
	fmt.Println(id, name)
}

context 机制是 Go 语言中处理并发请求时不可或缺的一部分。通过正确地使用 context，可以有效地管理请求的生命周期、设置超时和传递请求特有的数据。掌握了 context 的使用之后，可以编写出更加健壮和响应迅速的并发程序。