Go Context使用的实战技巧

Go Context 基础概念

在 Go 语言中，Context（上下文）是一个用于在 API 边界之间传递截止日期、取消信号及其他请求范围的值的对象。Context 通常用于控制多个 goroutine 的生命周期，尤其是在处理 HTTP 请求时，它能有效地管理资源并确保程序按预期响应。

Context 接口定义如下：

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

• Deadline：返回当前 Context 的截止时间。ok 为 true 时表示设置了截止时间，deadline 为截止时间点。
• Done：返回一个只读通道 <-chan struct{}，当此 Context 被取消或超时时，该通道会被关闭。
• Err：返回 Context 被取消或超时的原因。如果 Context 尚未取消或超时，返回 nil。
• Value：根据传入的 key 获取对应的值，常用于在不同的 goroutine 间传递请求范围的数据。

上下文的创建

1. Background 与 TODO
   • context.Background 是所有 Context 的根，通常用于程序的主入口，如 main 函数中启动的顶级 goroutine。

   func main() {
       ctx := context.Background()
       // 使用 ctx 启动 goroutine
   }
   

   • context.TODO 用于暂时不知道使用什么 Context 的情况，它主要作为一个占位符，提醒开发者后续需要替换为合适的 Context。

   func someFunction() {
       ctx := context.TODO()
       // 后续代码，这里应该被替换为合适的 Context
   }
   

2. WithCancel context.WithCancel 用于创建一个可以手动取消的 Context。它返回一个新的 Context 和一个取消函数 cancel。调用 cancel 函数会关闭 Context 的 Done 通道，所有基于此 Context 创建的子 Context 也会被取消。

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel() // 确保函数结束时取消，避免资源泄漏

    go func(ctx context.Context) {
        for {
            select {
            case <-ctx.Done():
                fmt.Println("goroutine cancelled")
                return
            default:
                fmt.Println("goroutine working")
                time.Sleep(100 * time.Millisecond)
            }
        }
    }(ctx)

    time.Sleep(500 * time.Millisecond)
    cancel()
    time.Sleep(200 * time.Millisecond)
}

在上述代码中，我们创建了一个可取消的 Context，并在一个 goroutine 中监听其取消信号。主线程等待一段时间后调用 cancel 函数，取消 goroutine 的执行。

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

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

    go func(ctx context.Context) {
        for {
            select {
            case <-ctx.Done():
                fmt.Println("goroutine cancelled due to deadline")
                return
            default:
                fmt.Println("goroutine working")
                time.Sleep(100 * time.Millisecond)
            }
        }
    }(ctx)

    time.Sleep(800 * time.Millisecond)
}

在此示例中，我们设置了一个 500 毫秒后的截止时间。当超过这个时间，即使主线程没有手动调用 cancel，goroutine 也会因为 Context 超时而被取消。

4. WithTimeout context.WithTimeout 是 context.WithDeadline 的便捷函数，它通过传入一个超时时间来创建一个带有截止时间的 Context。

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

    go func(ctx context.Context) {
        for {
            select {
            case <-ctx.Done():
                fmt.Println("goroutine cancelled due to timeout")
                return
            default:
                fmt.Println("goroutine working")
                time.Sleep(100 * time.Millisecond)
            }
        }
    }(ctx)

    time.Sleep(800 * time.Millisecond)
}

这段代码与 context.WithDeadline 的效果类似，只是使用 WithTimeout 更简洁，直接传入超时时间即可。

在 HTTP 服务中使用 Context

1. 处理请求取消 在 HTTP 服务中，Context 常用于处理客户端取消请求的情况。当客户端关闭连接时，对应的 Context 会被取消，从而可以中断正在进行的处理。

package main

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

func longRunningHandler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    fmt.Println("Handler started")
    select {
    case <-ctx.Done():
        fmt.Println("Request cancelled by client")
        http.Error(w, "Request cancelled", http.StatusRequestTimeout)
        return
    case <-time.After(3 * time.Second):
        fmt.Println("Handler completed")
        fmt.Fprintf(w, "Long running task completed")
    }
}

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

在这个示例中，longRunningHandler 函数从请求中获取 Context。如果在任务完成前 Context 被取消（客户端关闭连接），则会返回相应的错误信息。

2. 设置请求超时 可以通过 Context 为 HTTP 请求设置超时，避免处理过长时间的请求占用资源。

package main

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

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

    select {
    case <-ctx.Done():
        fmt.Println("Request timed out")
        http.Error(w, "Request timed out", http.StatusRequestTimeout)
        return
    case <-time.After(3 * time.Second):
        fmt.Println("Handler completed")
        fmt.Fprintf(w, "Task completed")
    }
}

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

在 timeoutHandler 中，我们创建了一个 2 秒超时的 Context。如果任务在 2 秒内未完成，Context 会被取消，返回超时错误。

在数据库操作中使用 Context

1. 控制数据库查询超时 当进行数据库查询时，使用 Context 可以避免长时间运行的查询占用资源。以 database/sql 包为例：

package main

import (
    "context"
    "database/sql"
    "fmt"
    _ "github.com/lib/pq"
    "time"
)

func main() {
    db, err := sql.Open("postgres", "user=postgres dbname=mydb sslmode=disable")
    if err != nil {
        panic(err)
    }
    defer db.Close()

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

    var result string
    err = db.QueryRowContext(ctx, "SELECT some_column FROM some_table WHERE some_condition").Scan(&result)
    if err != nil {
        if err == context.DeadlineExceeded {
            fmt.Println("Query timed out")
        } else {
            fmt.Println("Query error:", err)
        }
        return
    }
    fmt.Println("Query result:", result)
}

在此代码中，我们使用 QueryRowContext 方法并传入带有超时的 Context。如果查询在 3 秒内未完成，会返回 context.DeadlineExceeded 错误。

2. 事务与 Context 在数据库事务中，Context 也非常有用。它可以确保在事务处理过程中，如果 Context 被取消，事务也能正确回滚。

package main

import (
    "context"
    "database/sql"
    "fmt"
    _ "github.com/lib/pq"
    "time"
)

func main() {
    db, err := sql.Open("postgres", "user=postgres dbname=mydb sslmode=disable")
    if err != nil {
        panic(err)
    }
    defer db.Close()

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

    tx, err := db.BeginTx(ctx, nil)
    if err != nil {
        fmt.Println("Failed to start transaction:", err)
        return
    }

    _, err = tx.ExecContext(ctx, "INSERT INTO some_table (column1, column2) VALUES ($1, $2)", "value1", "value2")
    if err != nil {
        fmt.Println("Failed to execute query in transaction:", err)
        tx.Rollback()
        return
    }

    err = tx.Commit()
    if err != nil {
        fmt.Println("Failed to commit transaction:", err)
        return
    }
    fmt.Println("Transaction completed successfully")
}

在这个事务处理的例子中，我们使用带有超时的 Context 来启动事务，并在执行 SQL 语句时也传入 Context。如果在事务处理过程中 Context 超时而取消，我们会回滚事务，确保数据的一致性。

Context 的嵌套与传递

1. Context 嵌套 在实际应用中，经常需要创建嵌套的 Context。例如，在一个 HTTP 处理函数中启动多个子 goroutine，每个子 goroutine 可能需要不同的截止时间或取消逻辑，但它们又共享一些请求范围的数据。

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

    subCtx, subCancel := context.WithTimeout(ctx, 2*time.Second)
    defer subCancel()

    go func(ctx context.Context) {
        select {
        case <-ctx.Done():
            fmt.Println("Sub - goroutine cancelled")
            return
        case <-time.After(3 * time.Second):
            fmt.Println("Sub - goroutine completed")
        }
    }(subCtx)

    time.Sleep(4 * time.Second)
}

在上述代码中，我们创建了一个主 Context ctx，并基于它创建了一个子 Context subCtx，子 Context 的超时时间更短。当主 Context 被取消或超时时，子 Context 也会相应地被取消。但子 Context 可以在主 Context 之前因为自身的超时设置而被取消。

2. Context 传递 Context 需要在函数调用链中正确传递，以确保所有相关的 goroutine 都能接收到取消或超时信号。

package main

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

func worker1(ctx context.Context) {
    fmt.Println("Worker1 started")
    worker2(ctx)
    fmt.Println("Worker1 ended")
}

func worker2(ctx context.Context) {
    fmt.Println("Worker2 started")
    select {
    case <-ctx.Done():
        fmt.Println("Worker2 cancelled")
        return
    case <-time.After(2 * time.Second):
        fmt.Println("Worker2 completed")
    }
    fmt.Println("Worker2 ended")
}

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

    worker1(ctx)
}

在这个例子中，main 函数创建了一个带有超时的 Context，并将其传递给 worker1 函数，worker1 又将其传递给 worker2。这样，当 Context 超时时，worker2 能够接收到取消信号并正确处理。

使用 Context 传递请求范围数据

1. 自定义键类型 为了在 Context 中安全地传递数据，通常会定义自定义的键类型。

type requestIDKey struct{}

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

    value := ctx.Value(requestIDKey{})
    if value != nil {
        fmt.Println("Request ID:", value)
    }
}

在上述代码中，我们定义了一个 requestIDKey 类型，并使用它作为键在 Context 中存储和获取请求 ID。这种方式可以避免键冲突，因为不同的包可以定义自己的唯一键类型。

2. 在函数调用链中传递数据

type userIDKey struct{}

func handleRequest(ctx context.Context) {
    userID := "user - 1"
    ctx = context.WithValue(ctx, userIDKey{}, userID)
    processRequest(ctx)
}

func processRequest(ctx context.Context) {
    value := ctx.Value(userIDKey{})
    if value != nil {
        fmt.Println("Processing request for user:", value)
    }
}

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

在此示例中，handleRequest 函数将用户 ID 存储在 Context 中，并传递给 processRequest 函数。processRequest 函数从 Context 中获取用户 ID 并进行相应处理，展示了如何在函数调用链中传递请求范围的数据。

Context 使用的常见问题与注意事项

1. 避免泄漏 Context 在使用 context.WithCancel、context.WithDeadline 和 context.WithTimeout 创建 Context 时，一定要记得在函数结束时调用取消函数 cancel，以避免资源泄漏。

func someFunction() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    // 函数主体逻辑
}

如果忘记调用 cancel，可能会导致 goroutine 无法正确取消，占用不必要的资源。

2. 正确选择 Context 要根据实际需求选择合适的 Context 创建函数。例如，如果需要手动取消，使用 context.WithCancel；如果需要设置截止时间，使用 context.WithDeadline 或 context.WithTimeout。在 HTTP 处理中，要根据客户端请求的特点选择合适的超时设置。

3. 不要将 Context 放入结构体 Context 应该作为函数参数传递，而不是嵌入到结构体中。因为 Context 是用于传递请求范围的数据和控制信号，它的生命周期与函数调用相关，而不是与结构体实例的生命周期绑定。

// 错误示例
type MyStruct struct {
    ctx context.Context
}

func (ms *MyStruct) someMethod() {
    // 使用 ms.ctx
}

// 正确示例
func someFunction(ctx context.Context) {
    // 使用 ctx
}

4. 注意 Context 的继承关系 子 Context 会继承父 Context 的取消和超时信号。在创建嵌套的 Context 时，要注意设置合理的超时时间，避免子 Context 因为父 Context 的过早取消而无法完成预期任务，或者子 Context 的超时长于父 Context 导致资源浪费。

通过深入理解和正确运用 Go 语言的 Context，我们可以更好地控制 goroutine 的生命周期，处理 HTTP 请求、数据库操作等场景中的超时和取消逻辑，以及在不同的 goroutine 和函数之间传递请求范围的数据，从而编写出更健壮、高效的并发程序。在实际开发中，不断地实践和总结经验，能够更加熟练地掌握 Context 的各种实战技巧，提升程序的质量和性能。