Go context用法的最佳实践 - 摩柯技术社区

1. 理解Go context的本质

在Go语言中，context（上下文）是一种用于在不同的Go协程之间传递截止时间、取消信号和其他请求范围的值的机制。它是Go 1.7引入的重要特性，对于构建可取消、可控制生命周期和可传递元数据的并发程序至关重要。

本质上，context是一个接口类型，定义如下：

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

Deadline方法返回一个截止时间。如果ok为true，表示截止时间已设置，程序应在截止时间前完成相关操作。
Done方法返回一个只读通道。当context被取消或超时，这个通道会被关闭。
Err方法返回context被取消的原因。如果Done通道未关闭，Err返回nil；如果context被取消，Err返回Canceled错误；如果context超时，Err返回DeadlineExceeded错误。
Value方法用于获取与context关联的键值对中的值。

2. 基本使用场景 - 取消协程

在并发编程中，经常需要能够在某个条件满足时取消正在运行的协程。context为此提供了方便的机制。

首先，来看一个简单的示例，展示如何使用context取消一个长时间运行的协程：

package main

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

func worker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("Worker received cancel signal, exiting...")
            return
        default:
            fmt.Println("Worker is working...")
            time.Sleep(1 * time.Second)
        }
    }
}

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

    time.Sleep(3 * time.Second)
    cancel()
    time.Sleep(1 * time.Second)
}

在上述代码中：

context.WithCancel(context.Background())创建了一个可取消的context和对应的取消函数cancel。context.Background()是所有context的根，通常用于主函数、初始化和测试代码。
worker函数内部通过select语句监听ctx.Done()通道。当该通道接收到信号（即context被取消），worker函数会退出。
在main函数中，启动worker协程后，等待3秒，然后调用cancel函数取消context，从而通知worker协程退出。

3. 超时控制

除了手动取消，context还能设置操作的超时时间。这在处理网络请求、数据库查询等可能长时间阻塞的操作时非常有用。

package main

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

func task(ctx context.Context) {
    select {
    case <-ctx.Done():
        fmt.Println("Task timed out or was canceled, exiting...")
        return
    case <-time.After(5 * time.Second):
        fmt.Println("Task completed successfully")
    }
}

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

    go task(ctx)

    time.Sleep(5 * time.Second)
}

在这个例子中：

context.WithTimeout(context.Background(), 3*time.Second)创建了一个带有3秒超时的context。defer cancel()确保无论函数如何结束，context都会被正确取消，避免资源泄漏。
task函数内部使用select语句，同时监听ctx.Done()通道和一个5秒的定时器通道。由于context设置了3秒超时，在3秒后，ctx.Done()通道会被关闭，task函数会收到超时信号并退出。即使任务本身在5秒内可以完成，由于超时设置，它也会在3秒后被终止。

4. 传递截止时间

context可以携带截止时间信息，以便各个子协程知道整个操作的截止期限。这有助于协调不同层次的协程，确保它们在截止时间前完成任务。

package main

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

func subtask(ctx context.Context) {
    deadline, ok := ctx.Deadline()
    if ok {
        fmt.Printf("Subtask has a deadline: %v\n", deadline)
    } else {
        fmt.Println("Subtask has no specific deadline")
    }

    select {
    case <-ctx.Done():
        fmt.Println("Subtask was canceled or timed out")
        return
    case <-time.After(2 * time.Second):
        fmt.Println("Subtask completed")
    }
}

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

    go subtask(ctx)

    time.Sleep(5 * time.Second)
}

在这段代码里：

context.WithDeadline(context.Background(), deadline)创建了一个带有特定截止时间的context。
subtask函数通过ctx.Deadline()方法获取截止时间信息，并打印出来。如果有截止时间，会打印截止时间；否则，提示没有特定截止时间。
同样，subtask函数通过监听ctx.Done()通道来处理取消或超时情况。如果在截止时间前ctx.Done()通道被关闭，subtask函数会退出。

5. 传递请求范围的值

context还可以用于在不同的协程之间传递请求范围的值，比如用户认证信息、请求ID等。这使得在整个请求处理链中共享这些信息变得非常方便。

package main

import (
    "context"
    "fmt"
)

type User struct {
    Name string
}

func processRequest(ctx context.Context) {
    user, ok := ctx.Value("user").(*User)
    if ok {
        fmt.Printf("Processing request for user: %s\n", user.Name)
    } else {
        fmt.Println("No user information in context")
    }
}

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

    processRequest(ctx)
}

在这个示例中：

context.WithValue(context.Background(), "user", user)创建了一个新的context，并将一个User对象与键"user"关联起来。
processRequest函数通过ctx.Value("user")获取与键"user"关联的值，并断言其类型为*User。如果获取成功，打印用户信息；否则，提示没有用户信息。

6. 在HTTP服务器中使用context

在Go的HTTP服务器编程中，context起着至关重要的作用。它可以用于管理请求的生命周期、传递请求特定的数据以及处理超时。

package main

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

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()

    // 模拟一个长时间运行的任务
    select {
    case <-ctx.Done():
        err := ctx.Err()
        fmt.Fprintf(w, "Request canceled: %v\n", err)
        return
    case <-time.After(5 * time.Second):
        fmt.Fprintf(w, "Request processed successfully")
    }
}

func main() {
    server := &http.Server{
        Addr:    ":8080",
        Handler: http.HandlerFunc(handler),
    }

    go func() {
        if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
            fmt.Printf("Server error: %v\n", err)
        }
    }()

    // 等待一段时间后关闭服务器
    time.Sleep(3 * time.Second)

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

    if err := server.Shutdown(ctx); err != nil {
        fmt.Printf("Server shutdown error: %v\n", err)
    }
}

在上述代码中：

r.Context()获取每个HTTP请求的context。在handler函数中，通过监听ctx.Done()通道来处理请求取消或超时。如果请求在5秒内被取消（比如客户端关闭连接），ctx.Done()通道会被关闭，返回相应的错误信息。
在main函数中，创建了一个HTTP服务器，并启动它。3秒后，使用context.WithTimeout创建一个带有5秒超时的context，并调用server.Shutdown(ctx)关闭服务器。如果在5秒内服务器未能正常关闭，会返回相应的错误。

7. 在数据库操作中使用context

当进行数据库查询等操作时，context同样非常有用，可以控制操作的超时和取消。以下以database/sql包为例。

package main

import (
    "context"
    "database/sql"
    "fmt"
    _ "github.com/lib/pq" // 以PostgreSQL为例，需要安装相应驱动
    "time"
)

func queryData(ctx context.Context, db *sql.DB) {
    var result string
    err := db.QueryRowContext(ctx, "SELECT 'Hello, World!'").Scan(&result)
    if err != nil {
        if err == context.Canceled {
            fmt.Println("Query was canceled")
        } else if err == context.DeadlineExceeded {
            fmt.Println("Query timed out")
        } else {
            fmt.Printf("Query error: %v\n", err)
        }
        return
    }
    fmt.Printf("Query result: %s\n", result)
}

func main() {
    db, err := sql.Open("postgres", "user=postgres dbname=mydb sslmode=disable")
    if err != nil {
        fmt.Printf("Failed to connect to database: %v\n", err)
        return
    }
    defer db.Close()

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

    queryData(ctx, db)
}

在这个示例中：

db.QueryRowContext(ctx, "SELECT 'Hello, World!'").Scan(&result)使用context进行数据库查询。ctx可以控制查询的超时和取消。
如果context被取消（比如外部调用cancel函数），err会等于context.Canceled；如果context超时，err会等于context.DeadlineExceeded。根据不同的错误类型，程序可以进行相应的处理。

8. 嵌套context的正确使用

在实际应用中，经常会遇到需要在多个层次的协程中传递context的情况，这就涉及到嵌套context。正确使用嵌套context可以确保整个系统的取消和超时机制正常工作。

package main

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

func innerTask(ctx context.Context) {
    select {
    case <-ctx.Done():
        fmt.Println("Inner task was canceled")
        return
    case <-time.After(2 * time.Second):
        fmt.Println("Inner task completed")
    }
}

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

    go innerTask(ctx)

    select {
    case <-ctx.Done():
        fmt.Println("Outer task was canceled or timed out")
        return
    case <-time.After(5 * time.Second):
        fmt.Println("Outer task completed")
    }
}

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

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

在上述代码中：

outerTask函数接收一个context，并创建了一个带有3秒超时的子context。这个子context被传递给innerTask函数。
innerTask函数监听接收到的context的取消信号。如果outerTask的context被取消或超时，innerTask也会收到相应的取消信号并退出。
在main函数中，创建了一个可取消的context并传递给outerTask。1秒后，调用cancel函数取消context，从而导致outerTask和innerTask都收到取消信号并退出。

9. 避免常见错误

忘记取消context：在使用context.WithCancel、context.WithTimeout或context.WithDeadline创建context时，一定要确保在合适的时机调用取消函数。忘记调用取消函数可能会导致资源泄漏，例如协程无法正常退出，数据库连接无法释放等。
错误传递context：在传递context时，要确保正确地将父context传递给子协程。如果传递了错误的context，可能会导致取消和超时机制无法按预期工作。例如，在HTTP处理函数中，应该使用r.Context()作为父context进行传递，而不是创建一个新的独立context。
滥用context.Value：虽然context.Value可以方便地在协程间传递数据，但不应该滥用。因为它缺乏类型安全，过多使用可能会使代码难以理解和维护。应该只在必要时使用，并且尽量保持传递的数据简单和明确。

通过正确理解和使用Go语言的context，可以构建出更加健壮、可控制和高效的并发程序，无论是在网络编程、数据库操作还是其他需要并发处理的场景中。