Go context基本数据结构剖析

Go context 概述

在 Go 语言的编程世界里，context（上下文）扮演着至关重要的角色，尤其是在处理并发编程、控制流管理以及资源清理等场景。context 包提供了一种机制，用于在不同的 goroutine 之间传递截止时间、取消信号以及其他请求范围的值。

context 接口剖析

在 Go 语言中，context 是一个接口，定义如下：

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key interface{}) interface{}
}

1. Deadline 方法：
   • Deadline 方法返回截止时间，这个截止时间是当前 context 应该被取消的时间点。如果 ok 为 false，则表示没有设置截止时间。
   • 示例代码：

package main

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

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    deadline, ok := ctx.Deadline()
    if ok {
        fmt.Printf("截止时间: %v\n", deadline)
    } else {
        fmt.Println("未设置截止时间")
    }
}

2. Done 方法：
   • Done 方法返回一个只读通道，当 context 被取消或者超时的时候，这个通道会被关闭。在 goroutine 中可以通过监听这个通道来判断是否需要停止当前的工作。
   • 示例代码：

package main

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

func worker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("收到取消信号，停止工作")
            return
        default:
            fmt.Println("正在工作...")
            time.Sleep(500 * time.Millisecond)
        }
    }
}

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    go worker(ctx)

    time.Sleep(3 * time.Second)
}

3. Err 方法：
   • Err 方法返回 context 被取消的原因。如果 context 还没有被取消，Err 方法会返回 nil。当 context 被取消时，根据不同的取消原因，Err 方法会返回相应的错误。
   • 示例代码：

package main

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

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    time.Sleep(3 * time.Second)

    err := ctx.Err()
    if err != nil {
        fmt.Printf("context 取消原因: %v\n", err)
    }
}

4. Value 方法：
   • Value 方法用于从 context 中获取一个键值对的值。这个方法主要用于在不同的 goroutine 之间传递请求范围的数据，比如请求的用户认证信息等。
   • 示例代码：

package main

import (
    "context"
    "fmt"
)

func main() {
    ctx := context.WithValue(context.Background(), "user", "John")

    value := ctx.Value("user")
    if value != nil {
        fmt.Printf("用户: %v\n", value)
    }
}

context 的实现类型

1. background 与 todo 上下文：
   • context.Background 和 context.TODO 是 context 包中两个特殊的上下文，它们是所有 context 的根上下文。
   • context.Background 通常用于主函数、初始化以及测试代码中，作为上下文树的根。它没有截止时间，也不会被取消。
   • context.TODO 用于暂时不知道该使用什么上下文的情况，通常用于代码需要上下文，但还没有合适的上下文传入的场景。
   • 示例代码：

package main

import (
    "context"
    "fmt"
)

func main() {
    bgCtx := context.Background()
    todoCtx := context.TODO()

    fmt.Printf("background 上下文: %T\n", bgCtx)
    fmt.Printf("TODO 上下文: %T\n", todoCtx)
}

2. WithCancel 上下文：
   • context.WithCancel 函数创建一个可以手动取消的上下文。它接受一个父上下文，并返回一个新的上下文和一个取消函数。调用取消函数可以取消这个新的上下文，同时也会取消所有从这个上下文派生出来的子上下文。
   • 示例代码：

package main

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

func worker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("收到取消信号，停止工作")
            return
        default:
            fmt.Println("正在工作...")
            time.Sleep(500 * time.Millisecond)
        }
    }
}

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

    go worker(ctx)

    time.Sleep(2 * time.Second)
    cancel()

    time.Sleep(1 * time.Second)
}

3. WithTimeout 上下文：
   • context.WithTimeout 函数创建一个带有超时时间的上下文。它接受一个父上下文和一个超时时间，返回一个新的上下文和一个取消函数。当超时时间到达时，上下文会自动取消，同时也会调用取消函数。
   • 示例代码：

package main

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

func worker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("收到取消信号，停止工作")
            return
        default:
            fmt.Println("正在工作...")
            time.Sleep(500 * time.Millisecond)
        }
    }
}

func main() {
    parentCtx := context.Background()
    ctx, cancel := context.WithTimeout(parentCtx, 2*time.Second)
    defer cancel()

    go worker(ctx)

    time.Sleep(3 * time.Second)
}

4. WithDeadline 上下文：
   • context.WithDeadline 函数创建一个带有截止时间的上下文。它接受一个父上下文和一个截止时间，返回一个新的上下文和一个取消函数。当到达截止时间时，上下文会自动取消，同时也会调用取消函数。
   • 示例代码：

package main

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

func worker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("收到取消信号，停止工作")
            return
        default:
            fmt.Println("正在工作...")
            time.Sleep(500 * time.Millisecond)
        }
    }
}

func main() {
    parentCtx := context.Background()
    deadline := time.Now().Add(2 * time.Second)
    ctx, cancel := context.WithDeadline(parentCtx, deadline)
    defer cancel()

    go worker(ctx)

    time.Sleep(3 * time.Second)
}

5. WithValue 上下文：
   • context.WithValue 函数创建一个新的上下文，并携带一个键值对。这个键值对可以在不同的 goroutine 之间传递，通过 Value 方法获取。
   • 示例代码：

package main

import (
    "context"
    "fmt"
)

func process(ctx context.Context) {
    value := ctx.Value("key")
    if value != nil {
        fmt.Printf("获取到的值: %v\n", value)
    }
}

func main() {
    ctx := context.WithValue(context.Background(), "key", "value")
    go process(ctx)

    time.Sleep(1 * time.Second)
}

context 的数据结构分析

1. emptyCtx 结构体：
   • emptyCtx 是一个空结构体，用于实现 context.Background 和 context.TODO。它没有任何字段，因为这两个上下文不需要任何额外的状态。
   • 其定义如下：

type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}

func (*emptyCtx) Done() <-chan struct{} {
    return nil
}

func (*emptyCtx) Err() error {
    return nil
}

func (*emptyCtx) Value(key interface{}) interface{} {
    return nil
}

2. cancelCtx 结构体：
   • cancelCtx 结构体用于实现可以取消的上下文，如 context.WithCancel 创建的上下文。
   • 其定义如下：

type cancelCtx struct {
    Context

    mu       sync.Mutex
    done     atomic.Value
    children map[canceler]struct{}
    err      error
}

- `Context` 嵌入了父上下文，使得 `cancelCtx` 可以继承父上下文的行为。
- `mu` 是一个互斥锁，用于保护 `done`、`children` 和 `err` 字段。
- `done` 字段是一个 `atomic.Value`，用于存储一个通道，当上下文被取消时，这个通道会被关闭。
- `children` 字段是一个 map，存储了所有依赖这个上下文的子上下文。
- `err` 字段存储上下文取消的原因。

3. timerCtx 结构体： - timerCtx 结构体用于实现带有截止时间或超时时间的上下文，如 context.WithTimeout 和 context.WithDeadline 创建的上下文。 - 其定义如下：

type timerCtx struct {
    cancelCtx
    timer *time.Timer
    deadline time.Time
}

- `timerCtx` 嵌入了 `cancelCtx`，继承了其取消功能。
- `timer` 字段是一个 `time.Timer`，用于在到达截止时间或超时时间时触发取消操作。
- `deadline` 字段存储截止时间。

4. valueCtx 结构体： - valueCtx 结构体用于实现携带键值对的上下文，如 context.WithValue 创建的上下文。 - 其定义如下：

type valueCtx struct {
    Context
    key, val interface{}
}

- `Context` 嵌入了父上下文。
- `key` 和 `val` 字段分别存储键值对中的键和值。

context 的使用场景

1. 控制并发请求：
   • 在微服务架构中，一个请求可能会触发多个并发的子请求。通过 context 可以统一控制这些子请求的生命周期，当主请求取消或者超时的时候，所有子请求也会被取消。
   • 示例代码：

package main

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

func subRequest(ctx context.Context, id int) {
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("子请求 %d 收到取消信号，停止工作\n", id)
            return
        default:
            fmt.Printf("子请求 %d 正在工作...\n", id)
            time.Sleep(500 * time.Millisecond)
        }
    }
}

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

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

    time.Sleep(3 * time.Second)
}

2. 传递请求范围数据：
   • 在一个请求的处理过程中，可能需要在不同的函数和 goroutine 之间传递一些请求范围的数据，如用户认证信息、请求 ID 等。通过 context.WithValue 可以方便地实现这一需求。
   • 示例代码：

package main

import (
    "context"
    "fmt"
)

func handleRequest(ctx context.Context) {
    requestID := ctx.Value("requestID")
    if requestID != nil {
        fmt.Printf("处理请求，请求 ID: %v\n", requestID)
    }
}

func main() {
    ctx := context.WithValue(context.Background(), "requestID", "12345")
    go handleRequest(ctx)

    time.Sleep(1 * time.Second)
}

3. 资源清理：
   • 当一个上下文被取消时，可以通过监听 Done 通道来进行资源清理工作，如关闭文件、数据库连接等。
   • 示例代码：

package main

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

func main() {
    file, err := os.Open("test.txt")
    if err != nil {
        fmt.Println("打开文件失败:", err)
        return
    }
    defer file.Close()

    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    go func() {
        select {
        case <-ctx.Done():
            fmt.Println("关闭文件")
            file.Close()
        }
    }()

    time.Sleep(3 * time.Second)
}

context 使用的注意事项

1. 避免在全局变量中使用 context：
   • context 应该作为参数在函数调用链中传递，而不是作为全局变量。因为全局变量的 context 无法根据不同的请求进行动态调整，容易导致逻辑错误。
2. 正确处理取消信号：
   • 在 goroutine 中使用 context 时，要及时处理 Done 通道的关闭信号，确保在上下文取消时能够及时停止工作，避免资源泄漏。
3. 合理设置截止时间：
   • 在使用带有截止时间或超时时间的上下文时，要根据实际业务需求合理设置时间，避免设置过长导致资源浪费，或者设置过短导致业务处理不完整。
4. 注意 context 传递过程中的性能：
   • 虽然 context 的传递相对轻量级，但在高并发场景下，如果频繁创建和传递 context，尤其是带有大量键值对的 valueCtx，可能会对性能产生一定影响。要根据实际情况进行优化，如减少不必要的键值对传递。

通过深入剖析 Go 语言中 context 的基本数据结构，我们可以更好地理解其工作原理，从而在并发编程中更加灵活、高效地使用 context，编写出健壮、可维护的代码。无论是控制并发请求、传递请求范围数据还是进行资源清理，context 都为我们提供了强大而便捷的工具。在实际应用中，要充分理解其特性，并注意使用过程中的各种事项，以发挥 context 的最大价值。