Go使用context管理不同阶段的上下文切换

1. Go 语言中的 context 包概述

在 Go 语言的并发编程中，context包是一个极为重要的工具，用于在不同的 goroutine 之间传递截止时间、取消信号和其他请求范围的值。这对于管理复杂的并发任务，特别是那些需要在不同阶段进行上下文切换的任务，至关重要。

context包定义了Context接口，它有四个方法：

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

• Deadline方法返回当前Context的截止时间。如果ok为true，则表示截止时间已设置，在截止时间到达后，Context将被取消。
• Done方法返回一个只读的通道。当Context被取消或者到达截止时间时，这个通道会被关闭。
• Err方法返回Context被取消的原因。如果Context尚未被取消，返回nil；如果Context是因为超时而取消，返回context.DeadlineExceeded；如果Context是被手动取消，返回context.Canceled。
• Value方法用于获取与Context关联的键值对中的值。键通常是一个自定义类型，以避免命名冲突。

2. context 在不同阶段上下文切换中的基础使用

2.1 取消 goroutine

在一个复杂的并发程序中，可能会启动多个 goroutine 执行不同的任务。有时候，我们需要在某个特定的时刻取消其中一些或全部 goroutine。context可以很方便地实现这一点。

package main

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

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

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

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

在上述代码中，我们通过context.WithCancel创建了一个可取消的Context。在worker函数中，通过select语句监听ctx.Done()通道。当cancel函数被调用时，ctx.Done()通道被关闭，worker函数中的select语句会执行case <-ctx.Done()分支，从而优雅地停止worker。

2.2 设置截止时间

有时候，我们希望一个 goroutine 在一定时间内完成任务，如果超过这个时间，就取消该任务。context.WithTimeout可以满足这个需求。

package main

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

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

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

    go task(ctx)

    time.Sleep(3 * time.Second)
}

在这段代码中，我们通过context.WithTimeout创建了一个带有 1 秒超时的Context。在task函数中，通过select语句监听ctx.Done()通道和一个 2 秒的定时器通道。由于设置的超时时间是 1 秒，所以ctx.Done()通道会先被关闭，task函数会输出“task timed out”。

3. context 在不同阶段上下文切换中的进阶应用

3.1 跨层级传递 context

在实际应用中，一个 goroutine 可能会启动多个子 goroutine，这些子 goroutine 又可能启动更多的子 goroutine。在这种情况下，如何将取消信号或截止时间等上下文信息传递到最底层的 goroutine 呢？这就需要跨层级传递context。

package main

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

func subWorker(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("sub worker stopped")
            return
        default:
            fmt.Println("sub worker working")
            time.Sleep(100 * time.Millisecond)
        }
    }
}

func worker(ctx context.Context) {
    ctx, cancel := context.WithCancel(ctx)
    defer cancel()

    go subWorker(ctx)

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

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

    time.Sleep(1 * time.Second)
    cancel()
    time.Sleep(100 * time.Millisecond)
}

在上述代码中，main函数创建了一个可取消的Context并传递给worker函数。worker函数又创建了一个基于传入Context的新的可取消Context，并传递给subWorker函数。当main函数中的cancel函数被调用时，取消信号会层层传递，最终subWorker函数也会收到取消信号并停止工作。

3.2 使用 context 传递请求范围的值

除了取消信号和截止时间，context还可以用于在不同的 goroutine 之间传递请求范围的值。例如，在一个 Web 服务中，可能需要在不同的中间件和处理函数之间传递用户认证信息。

package main

import (
    "context"
    "fmt"
)

type userKey struct{}

func processRequest(ctx context.Context) {
    user := ctx.Value(userKey{}).(string)
    fmt.Printf("Processing request for user: %s\n", user)
}

func main() {
    ctx := context.WithValue(context.Background(), userKey{}, "John Doe")
    processRequest(ctx)
}

在这段代码中，我们定义了一个自定义类型userKey作为键，通过context.WithValue将用户信息“John Doe”与Context关联起来。在processRequest函数中，通过ctx.Value获取到用户信息并进行处理。

4. context 在 Web 开发中的上下文切换应用

4.1 处理 HTTP 请求

在 Go 的 Web 开发中，context常用于处理 HTTP 请求。net/http包在 Go 1.7 之后开始支持context。当一个 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():
        fmt.Fprintf(w, "request canceled\n")
        return
    case <-time.After(2 * time.Second):
        fmt.Fprintf(w, "request processed\n")
    }
}

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

在上述代码中，handler函数通过r.Context()获取与 HTTP 请求关联的Context。在处理请求时，通过select语句监听ctx.Done()通道和一个 2 秒的定时器通道。如果在 2 秒内请求被取消（例如客户端关闭连接），ctx.Done()通道会被关闭，函数会返回“request canceled”。

4.2 中间件中的 context 传递

在 Web 开发中，中间件是一个非常重要的概念。中间件可以在请求到达处理函数之前或之后执行一些通用的逻辑，如日志记录、认证等。在中间件中传递context可以确保上下文信息在整个请求处理过程中保持一致。

package main

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

func loggingMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        start := time.Now()
        ctx := context.WithValue(r.Context(), "startTime", start)
        r = r.WithContext(ctx)
        next.ServeHTTP(w, r)
        elapsed := time.Since(start)
        fmt.Printf("Request took %s\n", elapsed)
    })
}

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    startTime := ctx.Value("startTime").(time.Time)
    elapsed := time.Since(startTime)
    fmt.Fprintf(w, "Request processed in %s\n", elapsed)
}

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

在这段代码中，loggingMiddleware中间件通过context.WithValue将请求开始时间与Context关联起来，并通过r.WithContext将新的Context传递给下一个处理函数。在handler函数中，可以通过ctx.Value获取到请求开始时间并计算请求处理时间。

5. context 在数据库操作中的上下文切换应用

5.1 数据库查询的取消

在进行数据库查询时，如果查询时间过长，或者在查询过程中用户取消了操作，我们希望能够取消正在执行的查询。使用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(), 2*time.Second)
    defer cancel()

    rows, err := db.QueryContext(ctx, "SELECT * FROM my_table WHERE some_condition;")
    if err != nil {
        if err == context.DeadlineExceeded {
            fmt.Println("Query timed out")
        } else {
            fmt.Println("Query error:", err)
        }
        return
    }
    defer rows.Close()

    for rows.Next() {
        // 处理查询结果
    }

    if err := rows.Err(); err != nil {
        fmt.Println("Row error:", err)
    }
}

在上述代码中，我们通过db.QueryContext方法执行数据库查询，并传入一个带有 2 秒超时的Context。如果查询在 2 秒内没有完成，ctx.Done()通道会被关闭，db.QueryContext会返回context.DeadlineExceeded错误，从而取消查询。

5.2 事务中的 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(), 3*time.Second)
    defer cancel()

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

    _, err = tx.ExecContext(ctx, "UPDATE my_table SET some_column =? WHERE some_condition;")
    if err != nil {
        fmt.Println("Exec statement error:", err)
        tx.Rollback()
        return
    }

    err = tx.Commit()
    if err != nil {
        fmt.Println("Commit transaction error:", err)
    }
}

在这段代码中，我们通过db.BeginTx方法开始一个事务，并传入Context。在执行事务中的 SQL 语句时，使用tx.ExecContext方法并传入相同的Context。如果在事务执行过程中Context被取消（例如超时），可以通过检查err来决定是否回滚事务。

6. context 与 channel 的结合使用

6.1 利用 channel 实现 context 的复杂控制

虽然context本身提供了强大的取消和截止时间控制功能，但在某些复杂场景下，结合channel可以实现更灵活的控制。例如，我们可能希望根据不同的条件来决定是否取消某个 goroutine，而不仅仅依赖于截止时间或外部的取消信号。

package main

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

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    conditionCh := make(chan bool)

    go func() {
        for {
            select {
            case <-ctx.Done():
                fmt.Println("goroutine stopped due to context cancel")
                return
            case cond := <-conditionCh:
                if cond {
                    cancel()
                    fmt.Println("goroutine canceled due to condition")
                    return
                }
            }
        }
    }()

    time.Sleep(1 * time.Second)
    conditionCh <- true
    time.Sleep(1 * time.Second)
}

在上述代码中，我们创建了一个context和一个channel。在 goroutine 中，通过select语句监听ctx.Done()通道和conditionCh。当conditionCh接收到true时，调用cancel函数取消context，从而停止 goroutine。

6.2 使用 channel 传递 context 相关信息

有时候，我们可能需要在不同的 goroutine 之间传递与context相关的额外信息，而不仅仅是通过context.Value。channel可以很好地满足这个需求。

package main

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

type ContextInfo struct {
    Deadline time.Time
    Canceled bool
}

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    infoCh := make(chan ContextInfo)

    go func() {
        for {
            select {
            case <-ctx.Done():
                info := ContextInfo{
                    Deadline: ctx.Deadline(),
                    Canceled: true,
                }
                infoCh <- info
                return
            case <-time.After(100 * time.Millisecond):
                // 模拟工作
            }
        }
    }()

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

    info := <-infoCh
    fmt.Printf("Deadline: %v, Canceled: %v\n", info.Deadline, info.Canceled)
}

在这段代码中，我们定义了一个ContextInfo结构体来存储与context相关的信息。在 goroutine 中，当ctx.Done()通道被关闭时，将相关信息通过infoCh传递出去，主函数可以从infoCh中获取这些信息并进行处理。

7. context 使用中的常见问题与注意事项

7.1 避免内存泄漏

在使用context时，如果不正确地处理取消和截止时间，可能会导致内存泄漏。例如，如果一个 goroutine 没有正确监听ctx.Done()通道，当context被取消时，这个 goroutine 可能会继续运行，从而导致资源无法释放。

package main

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

func badWorker(ctx context.Context) {
    for {
        fmt.Println("bad worker working")
        time.Sleep(100 * time.Millisecond)
    }
}

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

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

在上述代码中，badWorker函数没有监听ctx.Done()通道，所以当cancel函数被调用时，badWorker函数不会停止，可能会导致内存泄漏。正确的做法是在badWorker函数中添加对ctx.Done()通道的监听。

7.2 合理设置截止时间

设置截止时间时，需要根据实际业务需求进行合理的设置。如果截止时间设置过短，可能会导致正常的任务被误取消；如果设置过长，可能会导致资源长时间被占用。

例如，在一个数据库查询中，如果查询的数据量较大，需要较长的时间来处理，那么设置一个较短的超时时间可能会导致查询经常失败。需要通过性能测试和实际业务场景来确定合适的截止时间。

7.3 注意 context 传递的层级

在跨层级传递context时，要确保context能够正确地传递到最底层的 goroutine。如果在中间某个层级丢失了context，可能会导致底层的 goroutine 无法接收到取消信号或截止时间信息。

在编写代码时，要仔细检查context在各个函数调用之间的传递情况，特别是在中间件或复杂的函数调用链中。

8. 总结 context 在不同阶段上下文切换中的应用

通过以上内容，我们详细了解了 Go 语言中context包在不同阶段上下文切换中的应用。从基础的取消 goroutine 和设置截止时间，到进阶的跨层级传递、在 Web 开发和数据库操作中的应用，以及与channel的结合使用，context为我们提供了强大而灵活的工具来管理并发任务的上下文。

在实际开发中，正确使用context可以提高程序的健壮性、可维护性和性能。同时，要注意避免常见的问题，如内存泄漏、不合理的截止时间设置和 context 传递层级错误等。

希望通过本文的介绍，读者能够对context在不同阶段上下文切换中的应用有更深入的理解，并在自己的 Go 语言项目中熟练运用context来解决复杂的并发问题。