Go Context的跨协程传递技巧

Go Context 的跨协程传递基础

理解 Go Context

在 Go 语言的并发编程中，context 包提供了一种强大的机制，用于在多个 goroutine 之间传递截止日期、取消信号以及其他请求范围的值。Context 类型是一个接口，它定义了四个方法：Deadline、Done、Err 和 Value。

• Deadline 方法返回 Context 被取消的时间点，或者是否没有截止时间。

func (c Context) Deadline() (deadline time.Time, ok bool)

• Done 方法返回一个只读的 <-chan struct{}，当 Context 被取消或超时时，这个通道会被关闭。

func (c Context) Done() <-chan struct{}

• Err 方法返回 Context 被取消的原因。如果 Context 还没有被取消，Err 返回 nil。

func (c Context) Err() error

• Value 方法返回与 Context 关联的键对应的值，如果没有关联的值，则返回 nil。

func (c Context) Value(key interface{}) interface{}

Context 的创建

Go 提供了几个函数来创建不同类型的 Context。最常用的是 context.Background 和 context.TODO。

context.Background 通常用于应用程序的根 Context，它永远不会被取消，没有截止日期，也没有关联的值。

func Background() Context

context.TODO 用于暂时不确定使用哪个 Context 的情况，通常在函数还没有接收到 Context 参数，但需要创建一个 Context 的时候使用。

func TODO() Context

此外，context.WithCancel、context.WithDeadline 和 context.WithTimeout 函数用于创建可取消的 Context。

context.WithCancel 创建一个可取消的 Context，返回一个取消函数 cancel，调用 cancel 函数可以取消这个 Context。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

context.WithDeadline 创建一个带有截止日期的 Context，当截止日期到达时，Context 会自动取消。

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)

context.WithTimeout 是 context.WithDeadline 的便捷函数，用于创建一个在指定时间段后自动取消的 Context。

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

简单的跨协程传递示例

下面是一个简单的示例，展示如何在两个 goroutine 之间传递 Context 并取消操作。

package main

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

func worker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("worker: context cancelled")
            return
        default:
            fmt.Println("worker: working...")
            time.Sleep(1 * time.Second)
        }
    }
}

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

    go worker(ctx)

    time.Sleep(5 * time.Second)
    fmt.Println("main: exiting")
}

在这个示例中，main 函数创建了一个带有 3 秒超时的 Context，并将其传递给 worker goroutine。worker goroutine 在循环中不断检查 ctx.Done() 通道是否关闭。当 Context 超时时，ctx.Done() 通道会被关闭，worker goroutine 会收到取消信号并退出。

复杂场景下的 Context 跨协程传递

多级协程传递

在实际应用中，可能会有多层 goroutine 调用，需要将 Context 一直传递下去。

package main

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

func subWorker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("subWorker: context cancelled")
            return
        default:
            fmt.Println("subWorker: working...")
            time.Sleep(1 * time.Second)
        }
    }
}

func worker(ctx context.Context) {
    go subWorker(ctx)
    for {
        select {
        case <-ctx.Done():
            fmt.Println("worker: context cancelled")
            return
        default:
            fmt.Println("worker: working...")
            time.Sleep(1 * time.Second)
        }
    }
}

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

    go worker(ctx)

    time.Sleep(5 * time.Second)
    fmt.Println("main: exiting")
}

在这个例子中，main 函数创建的 Context 被传递给 worker goroutine，worker goroutine 又将其传递给 subWorker goroutine。当 Context 超时时，worker 和 subWorker 都会收到取消信号并退出。

传递值的 Context

Context 还可以用于在多个 goroutine 之间传递请求范围的值。通过 context.WithValue 函数可以创建一个携带值的 Context。

package main

import (
    "context"
    "fmt"
)

type requestIDKey struct{}

func worker(ctx context.Context) {
    requestID := ctx.Value(requestIDKey{}).(string)
    fmt.Printf("worker: received request ID %s\n", requestID)
}

func main() {
    ctx := context.WithValue(context.Background(), requestIDKey{}, "12345")
    go worker(ctx)

    fmt.Println("main: started worker")
    select {}
}

在这个示例中，context.WithValue 创建了一个带有 requestID 值的 Context，并将其传递给 worker goroutine。worker goroutine 通过 ctx.Value 方法获取到这个值。

Context 与 HTTP 处理

在 HTTP 服务器编程中，Context 也起着至关重要的作用。net/http 包中的 http.Request 结构体包含一个 Context 字段，用于传递请求相关的信息。

package main

import (
    "context"
    "fmt"
    "net/http"
    "time"
)

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
    defer cancel()

    done := make(chan struct{})
    go func() {
        time.Sleep(3 * time.Second)
        close(done)
    }()

    select {
    case <-ctx.Done():
        fmt.Println("handler: context cancelled")
        http.Error(w, "request timed out", http.StatusGatewayTimeout)
    case <-done:
        fmt.Println("handler: operation completed")
        fmt.Fprintf(w, "operation completed")
    }
}

func main() {
    http.HandleFunc("/", handler)
    fmt.Println("Server listening on :8080")
    http.ListenAndServe(":8080", nil)
}

在这个示例中，http.Request 的 Context 被获取，并创建了一个带有 2 秒超时的新 Context。在一个 goroutine 中模拟一个耗时 3 秒的操作，当 Context 超时或操作完成时，相应地处理响应。

Context 跨协程传递的注意事项

避免在全局变量中使用 Context

虽然在某些情况下在全局变量中使用 Context 看起来很方便，但这会导致代码难以理解和测试。Context 应该作为参数在函数和 goroutine 之间传递，这样可以明确每个操作的上下文。

例如，避免这样的代码：

var globalCtx context.Context

func init() {
    globalCtx = context.Background()
}

func worker() {
    // 使用 globalCtx
}

而应该采用传递参数的方式：

func worker(ctx context.Context) {
    // 使用传入的 ctx
}

func main() {
    ctx := context.Background()
    go worker(ctx)
}

正确处理取消信号

在使用 Context 的 Done 通道时，要确保在 goroutine 中正确处理取消信号。如果 goroutine 没有及时响应取消信号，可能会导致资源泄漏或程序无法正常退出。

例如，在执行一些清理操作时，要先检查 Context 是否被取消：

func worker(ctx context.Context) {
    // 模拟一些操作
    for i := 0; i < 10; i++ {
        select {
        case <-ctx.Done():
            // 执行清理操作
            fmt.Println("worker: cleaning up")
            return
        default:
            fmt.Println("worker: working...", i)
            time.Sleep(1 * time.Second)
        }
    }
}

不要在 Context 中传递敏感信息

虽然 Context 可以用于传递值，但不应该用于传递敏感信息，如密码、密钥等。因为 Context 可能会被记录、打印或传递给不可信的代码，从而导致敏感信息泄露。

注意 Context 的生命周期

要清楚 Context 的生命周期，特别是在使用 WithCancel、WithDeadline 和 WithTimeout 创建的 Context 时。确保在不需要 Context 时及时调用取消函数，以避免资源浪费。

例如，在使用 WithCancel 创建 Context 时：

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

    // 一段时间后取消
    time.Sleep(3 * time.Second)
    cancel()

    time.Sleep(1 * time.Second)
    fmt.Println("main: exiting")
}

在这个示例中，main 函数在 3 秒后调用 cancel 函数取消 Context，以确保 worker goroutine 能及时收到取消信号并退出。

Context 与并发安全

并发访问 Context

Context 本身是并发安全的，可以在多个 goroutine 中同时访问。例如，多个 goroutine 可以同时读取 Context 的值或监听 Done 通道。

package main

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

func worker1(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("worker1: context cancelled")
            return
        default:
            fmt.Println("worker1: working...")
            time.Sleep(1 * time.Second)
        }
    }
}

func worker2(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("worker2: context cancelled")
            return
        default:
            fmt.Println("worker2: working...")
            time.Sleep(1 * time.Second)
        }
    }
}

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

    go worker1(ctx)
    go worker2(ctx)

    time.Sleep(5 * time.Second)
    fmt.Println("main: exiting")
}

在这个示例中，worker1 和 worker2 两个 goroutine 同时监听同一个 Context 的 Done 通道，这是安全的。

避免数据竞争

虽然 Context 本身是并发安全的，但如果在 Context 中传递的是可变的数据结构，需要注意避免数据竞争。例如，如果传递的是一个共享的 map，在多个 goroutine 中读写这个 map 可能会导致数据竞争。

package main

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

type dataKey struct{}

func worker(ctx context.Context, wg *sync.WaitGroup) {
    data := ctx.Value(dataKey{}).(map[string]int)
    // 这里对 data 进行读写操作需要加锁，否则可能出现数据竞争
    data["count"]++
    fmt.Printf("worker: data count is %d\n", data["count"])
    wg.Done()
}

func main() {
    data := map[string]int{"count": 0}
    ctx := context.WithValue(context.Background(), dataKey{}, data)

    var wg sync.WaitGroup
    for i := 0; i < 10; i++ {
        wg.Add(1)
        go worker(ctx, &wg)
    }

    wg.Wait()
    fmt.Println("main: all workers done")
}

在这个示例中，如果不对 data 进行同步，多个 worker goroutine 同时读写 data 中的 count 字段会导致数据竞争。可以使用 sync.Mutex 或其他同步机制来解决这个问题。

package main

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

type dataKey struct{}

func worker(ctx context.Context, wg *sync.WaitGroup, mu *sync.Mutex) {
    data := ctx.Value(dataKey{}).(map[string]int)
    mu.Lock()
    data["count"]++
    fmt.Printf("worker: data count is %d\n", data["count"])
    mu.Unlock()
    wg.Done()
}

func main() {
    data := map[string]int{"count": 0}
    ctx := context.WithValue(context.Background(), dataKey{}, data)

    var wg sync.WaitGroup
    var mu sync.Mutex
    for i := 0; i < 10; i++ {
        wg.Add(1)
        go worker(ctx, &wg, &mu)
    }

    wg.Wait()
    fmt.Println("main: all workers done")
}

通过添加 sync.Mutex，确保了对共享数据的安全访问。

总结 Context 跨协程传递的最佳实践

始终传递 Context

在编写函数和 goroutine 时，始终将 Context 作为第一个参数传递，即使当前函数暂时不需要使用 Context。这样可以确保在后续的代码演进中，Context 能够顺利地在调用链中传递。

尽早检查取消信号

在 goroutine 中，尽早检查 Context 的取消信号，特别是在执行长时间运行的操作之前。这样可以确保在 Context 被取消时，goroutine 能够及时响应并清理资源。

合理使用 Context 的方法

根据具体需求，合理使用 Context 的 Deadline、Done、Err 和 Value 方法。例如，在需要设置超时的情况下使用 Deadline 方法，在监听取消信号时使用 Done 通道等。

结合同步机制

当在 Context 中传递共享数据时，结合同步机制（如 sync.Mutex、sync.RWMutex 等）来确保数据的并发安全。

避免滥用 Context

不要滥用 Context 来传递不相关的数据，保持 Context 的简洁性和专注性。只在 Context 中传递与请求范围相关且在多个 goroutine 间需要共享的数据。

通过遵循这些最佳实践，可以更有效地在 Go 语言的并发编程中使用 Context 进行跨协程传递，提高程序的健壮性和可维护性。

在实际的项目开发中，要根据具体的业务场景和需求，灵活运用 Context 的跨协程传递技巧。无论是简单的并发任务还是复杂的分布式系统，正确使用 Context 都能帮助我们更好地管理资源、处理超时和取消操作，从而构建出高效、稳定的 Go 程序。