Go Context常见用法的拓展应用

Go Context 基础回顾

在深入探讨拓展应用之前，先简要回顾一下 Go Context 的基础概念和常见用法。

Context 是 Go 语言中用于管理请求范围的截止时间、取消信号以及传递请求特定值的工具。在 Go 1.7 引入后，它成为构建健壮并发程序的关键组件。

最常见的创建 Context 的方式是使用 context.Background 和 context.TODO。context.Background 通常用于整个应用程序的根 Context，而 context.TODO 用于暂时不确定使用哪种 Context 的情况。

例如：

package main

import (
    "context"
    "fmt"
)

func main() {
    ctx := context.Background()
    todo := context.TODO()
    fmt.Printf("Background context: %v\n", ctx)
    fmt.Printf("TODO context: %v\n", todo)
}

在这个简单示例中，我们创建了 context.Background 和 context.TODO，并打印它们。这两个 Context 本身不携带任何截止时间、取消信号或值，但可以作为更复杂 Context 的基础。

Context 主要通过 WithCancel、WithDeadline 和 WithTimeout 函数来创建携带取消信号或截止时间的 Context。

WithCancel 允许手动取消 Context：

package main

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

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    go func(ctx context.Context) {
        for {
            select {
            case <-ctx.Done():
                fmt.Println("Context cancelled")
                return
            default:
                fmt.Println("Working...")
                time.Sleep(1 * time.Second)
            }
        }
    }(ctx)
    time.Sleep(3 * time.Second)
    cancel()
    time.Sleep(1 * time.Second)
}

在上述代码中，我们通过 context.WithCancel 创建了一个可取消的 Context。在 goroutine 中，我们通过 select 语句监听 ctx.Done() 通道，当 Context 被取消时，ctx.Done() 通道会被关闭，从而退出 goroutine。

WithDeadline 和 WithTimeout 则用于设置截止时间。WithDeadline 接受一个绝对时间作为截止时间，而 WithTimeout 接受一个相对时间作为超时时间。

package main

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

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    go func(ctx context.Context) {
        select {
        case <-ctx.Done():
            fmt.Println("Context timed out")
            return
        default:
            fmt.Println("Working...")
            time.Sleep(3 * time.Second)
        }
    }(ctx)
    time.Sleep(4 * time.Second)
}

这里使用 context.WithTimeout 创建了一个 2 秒超时的 Context。在 goroutine 中，由于操作时间超过了 2 秒，最终会因为 Context 超时，ctx.Done() 通道被关闭而退出。

拓展应用：跨服务调用的 Context 传递

在微服务架构中，一个请求可能会涉及多个服务之间的调用。在这种情况下，正确传递 Context 至关重要，以确保整个请求生命周期内的截止时间和取消信号能够有效传递。

假设我们有两个服务，ServiceA 和 ServiceB，ServiceA 调用 ServiceB。

首先定义 ServiceB：

package main

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

func ServiceB(ctx context.Context) {
    select {
    case <-ctx.Done():
        fmt.Println("ServiceB: Context cancelled or timed out")
        return
    default:
        fmt.Println("ServiceB: Working...")
        time.Sleep(3 * time.Second)
        fmt.Println("ServiceB: Done")
    }
}

然后是 ServiceA 调用 ServiceB：

package main

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

func ServiceA() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    go ServiceB(ctx)
    time.Sleep(4 * time.Second)
}

在 ServiceA 中，我们创建了一个 2 秒超时的 Context，并传递给 ServiceB。由于 ServiceB 的操作时间超过了 2 秒，ServiceB 会因为 Context 超时收到取消信号而提前结束。

通过这种方式，我们确保了从 ServiceA 发起的整个请求在 2 秒内完成，无论涉及多少层服务调用，Context 的截止时间和取消信号都能正确传递。

拓展应用：在数据库操作中使用 Context

数据库操作往往是应用程序中耗时较长的部分，合理使用 Context 可以有效管理数据库操作的生命周期。

以 SQLite 数据库为例，假设我们有一个查询操作：

package main

import (
    "context"
    "database/sql"
    "fmt"
    _ "github.com/mattn/go - sqlite3"
    "time"
)

func queryDatabase(ctx context.Context, db *sql.DB) {
    rows, err := db.QueryContext(ctx, "SELECT * FROM users")
    if err != nil {
        if err == context.Canceled || err == context.DeadlineExceeded {
            fmt.Println("Database query cancelled or timed out")
        } else {
            fmt.Printf("Database query error: %v\n", err)
        }
        return
    }
    defer rows.Close()
    for rows.Next() {
        var id int
        var name string
        err := rows.Scan(&id, &name)
        if err != nil {
            fmt.Printf("Scan error: %v\n", err)
            return
        }
        fmt.Printf("User: ID %d, Name %s\n", id, name)
    }
    err = rows.Err()
    if err != nil {
        fmt.Printf("Rows error: %v\n", err)
    }
}

在上述代码中，我们使用 db.QueryContext 方法，它接受一个 Context。如果在查询过程中 Context 被取消或超时，数据库操作会被中断，并返回相应的错误。

调用这个函数：

func main() {
    db, err := sql.Open("sqlite3", "test.db")
    if err != nil {
        fmt.Printf("Failed to open database: %v\n", err)
        return
    }
    defer db.Close()
    ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
    defer cancel()
    go queryDatabase(ctx, db)
    time.Sleep(2 * time.Second)
}

这里我们创建了一个 1 秒超时的 Context，并传递给 queryDatabase 函数。如果查询操作在 1 秒内未完成，Context 会超时，数据库操作会被取消。

拓展应用：在分布式任务队列中使用 Context

在分布式任务队列场景下，任务可能会在不同的节点上执行，并且可能需要根据全局的取消信号或截止时间来终止任务。

假设我们使用一个简单的基于 Redis 的任务队列。首先定义任务处理函数：

package main

import (
    "context"
    "fmt"
    "github.com/go - redis/redis/v8"
    "time"
)

func processTask(ctx context.Context, rdb *redis.Client, taskID string) {
    select {
    case <-ctx.Done():
        fmt.Printf("Task %s cancelled\n", taskID)
        return
    default:
        fmt.Printf("Processing task %s...\n", taskID)
        time.Sleep(3 * time.Second)
        fmt.Printf("Task %s completed\n", taskID)
    }
}

然后是从任务队列中取出任务并处理的逻辑：

func main() {
    rdb := redis.NewClient(&redis.Options{
        Addr:     "localhost:6379",
        Password: "",
        DB:       0,
    })
    ctx := context.Background()
    for {
        taskID, err := rdb.RPop(ctx, "task_queue").Result()
        if err != nil {
            if err == redis.Nil {
                time.Sleep(1 * time.Second)
                continue
            }
            fmt.Printf("Error getting task from queue: %v\n", err)
            return
        }
        taskCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
        go processTask(taskCtx, rdb, taskID)
        time.Sleep(1 * time.Second)
        cancel()
    }
}

在上述代码中，我们从 Redis 任务队列中取出任务，并为每个任务创建一个 2 秒超时的 Context。如果任务处理时间超过 2 秒，Context 会超时，任务会被取消。

这种方式确保了在分布式任务队列环境下，即使任务在不同节点执行，也能根据统一的 Context 进行有效管理。

拓展应用：Context 与日志记录

在应用程序开发中，日志记录是非常重要的部分。结合 Context 可以使日志记录更加丰富和有意义，特别是在跟踪请求流程方面。

我们可以通过在 Context 中携带请求 ID 等信息，然后在日志记录中使用这些信息。

首先定义一个中间件来为每个请求创建一个携带请求 ID 的 Context：

package main

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

const requestIDKey = "requestID"

func withRequestID(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        requestID := fmt.Sprintf("%d", rand.Int63())
        ctx := context.WithValue(r.Context(), requestIDKey, requestID)
        r = r.WithContext(ctx)
        next.ServeHTTP(w, r)
    })
}

然后在处理函数中记录日志并使用 Context 中的请求 ID：

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    requestID := ctx.Value(requestIDKey).(string)
    fmt.Printf("Request ID %s received\n", requestID)
    time.Sleep(1 * time.Second)
    fmt.Fprintf(w, "Response for Request ID %s\n", requestID)
}

在主函数中使用中间件和处理函数：

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

在这个示例中，每个请求都被分配了一个唯一的请求 ID，并存储在 Context 中。处理函数通过从 Context 中获取请求 ID 来记录日志，这样在日志中就可以方便地跟踪每个请求的处理流程。

拓展应用：Context 在 WebSocket 中的应用

WebSocket 是一种在 Web 应用中实现双向通信的协议。在使用 WebSocket 时，合理使用 Context 可以有效管理连接的生命周期。

假设我们有一个简单的 WebSocket 服务器：

package main

import (
    "context"
    "fmt"
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

func serveWs(ctx context.Context, w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println(err)
        return
    }
    defer conn.Close()
    for {
        select {
        case <-ctx.Done():
            fmt.Println("WebSocket connection cancelled")
            return
        default:
            _, _, err := conn.ReadMessage()
            if err != nil {
                if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                    log.Printf("error: %v", err)
                }
                return
            }
            err = conn.WriteMessage(websocket.TextMessage, []byte("Message received"))
            if err != nil {
                log.Printf("error: %v", err)
                return
            }
        }
    }
}

在主函数中创建一个可取消的 Context 并传递给 WebSocket 处理函数：

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        serveWs(ctx, w, r)
    })
    go func() {
        time.Sleep(5 * time.Second)
        cancel()
    }()
    log.Fatal(http.ListenAndServe(":8080", nil))
}

在这个示例中，我们为 WebSocket 连接创建了一个可取消的 Context。当 Context 被取消时（这里模拟 5 秒后取消），WebSocket 连接会被关闭，从而有效管理了 WebSocket 连接的生命周期。

拓展应用：Context 与资源管理

在一些复杂的应用场景中，可能需要管理多个资源，并且这些资源的生命周期需要与 Context 相关联。

假设我们有一个需要管理文件资源和数据库连接资源的场景：

package main

import (
    "context"
    "database/sql"
    "fmt"
    "io/ioutil"
    _ "github.com/mattn/go - sqlite3"
    "os"
    "time"
)

func manageResources(ctx context.Context) {
    file, err := ioutil.TempFile("", "example")
    if err != nil {
        fmt.Printf("Failed to create temp file: %v\n", err)
        return
    }
    defer os.Remove(file.Name())
    db, err := sql.Open("sqlite3", "test.db")
    if err != nil {
        fmt.Printf("Failed to open database: %v\n", err)
        return
    }
    defer db.Close()
    select {
    case <-ctx.Done():
        fmt.Println("Resources management cancelled")
        return
    default:
        fmt.Println("Managing resources...")
        time.Sleep(3 * time.Second)
        fmt.Println("Resources management done")
    }
}

在主函数中创建一个超时的 Context 并传递给资源管理函数：

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    go manageResources(ctx)
    time.Sleep(4 * time.Second)
}

在这个示例中，当 Context 超时，资源管理函数会提前结束，并且相关的文件资源和数据库连接资源会被正确清理。这确保了在应用程序的复杂资源管理场景下，资源的生命周期能够与 Context 紧密关联，避免资源泄漏等问题。

通过以上多种拓展应用场景的介绍和代码示例，我们可以看到 Go Context 在不同领域和复杂场景下的强大作用，合理运用 Context 可以极大地提升应用程序的健壮性和可维护性。