Golang 中的 Context 如何讓代碼更智能、更安全、更易於擴展

Go 開發中，context 包已經成爲一個必不可少的工具。它提供了一種在不同的 goroutine 之間傳遞請求範圍內變量、取消信號和截止時間的方法。通過合理地使用 context，我們可以使代碼變得更智能、更安全，並且更易於擴展。本文將詳細探討 context 的作用以及如何在實際開發中應用它。

什麼是 Context？

context 是 Go 1.7 引入的一個標準庫，它主要用於在 goroutine 之間傳遞請求範圍內的變量以及控制信號。context 主要有以下幾個功能：

取消信號傳遞：可以通過 context 傳遞取消信號，用於取消正在進行的操作。
截止時間傳遞：可以設定一個截止時間，超時後會自動取消操作。
請求範圍內變量傳遞：可以在 context 中傳遞一些與請求相關的變量。

context 的基本使用

創建 Context

context 包提供了四種創建 context 的方法：

context.Background()：返回一個空的 Context，一般用於主函數、初始化和測試。
context.TODO()：返回一個空的 Context，表示目前還不知道用什麼 Context 時使用。
context.WithCancel(parent)：返回一個可取消的 context 和一個取消函數 cancel。
context.WithDeadline(parent, deadline) 和 context.WithTimeout(parent, timeout)：返回一個帶有超時功能的 context。

package main

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

func main() {
    // 創建一個帶有取消功能的 context
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()  // 確保在函數結束時取消 context

    gofunc() {
        time.Sleep(2 * time.Second)
        cancel()  // 2 秒後取消 context
    }()

    select {
    case <-time.After(3 * time.Second):
        fmt.Println("Operation completed")
    case <-ctx.Done():
        fmt.Println("Operation cancelled")
    }
}

在上述代碼中，創建了一個帶有取消功能的 context。在另一個 goroutine 中，在 2 秒後取消了 context，所以 select 語句會在 ctx.Done() 信道接收到取消信號時執行相應的操作。

如何使用 Context 使代碼更智能

通過使用 context，可以在不同的 goroutine 之間傳遞控制信號和變量，這樣可以減少全局變量的使用，使代碼更加模塊化和智能。例如，在處理 HTTP 請求時，可以將請求的上下文傳遞給所有處理函數，從而確保在請求取消時，所有相關的操作都能及時地響應並終止。

package main

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

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    fmt.Println("Handler started")
    defer fmt.Println("Handler ended")

    select {
    case <-time.After(5 * time.Second):
        fmt.Fprintf(w, "Hello, World!")
    case <-ctx.Done():
        err := ctx.Err()
        fmt.Println("Handler cancelled:", err)
        http.Error(w, err.Error(), http.StatusInternalServerError)
    }
}

func main() {
    http.HandleFunc("/", handler)
    http.ListenAndServe(":8080", nil)
}

在上述代碼中，將 HTTP 請求的 context 傳遞給處理函數 handler，這樣當請求被取消時，處理函數可以及時響應並終止操作。

如何使用 Context 使代碼更安全

context 提供的取消和超時機制，可以有效地防止資源泄漏和殭屍進程。例如，在進行數據庫查詢或網絡請求時，如果操作超時或請求被取消，可以及時終止操作，釋放資源。

package main

import (
    "context"
    "database/sql"
    "fmt"
    "log"
    "time"

    _ "github.com/lib/pq"
)

func queryWithTimeout(ctx context.Context, db *sql.DB, query string) (*sql.Rows, error) {
    ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
    defer cancel()

    return db.QueryContext(ctx, query)
}

func main() {
    connStr := "user=username db
    db, err := sql.Open("postgres", connStr)
    if err != nil {
        log.Fatal(err)
    }
    defer db.Close()

    ctx := context.Background()
    rows, err := queryWithTimeout(ctx, db, "SELECT * FROM mytable")
    if err != nil {
        log.Fatal("Query failed:", err)
    }
    defer rows.Close()

    for rows.Next() {
        var id int
        var name string
        if err := rows.Scan(&id, &name); err != nil {
            log.Fatal(err)
        }
        fmt.Printf("id: %d, name: %s\n", id, name)
    }
}

在上述代碼中，在數據庫查詢時使用了 context.WithTimeout，這樣如果查詢時間超過 2 秒，查詢操作會自動取消並返回超時錯誤，從而避免了長時間阻塞。

如何使用 Context 使代碼更易於擴展

通過使用 context，可以方便地在不同的函數之間傳遞信息，而不需要修改函數簽名。這使得代碼更易於擴展和維護。例如，可以在 context 中傳遞一些用戶認證信息或請求 ID，從而簡化函數參數。

package main

import (
    "context"
    "fmt"
)

type key int

const requestIDKey key = 0

func withRequestID(ctx context.Context, requestID string) context.Context {
    return context.WithValue(ctx, requestIDKey, requestID)
}

func requestIDFromContext(ctx context.Context) (string, bool) {
    requestID, ok := ctx.Value(requestIDKey).(string)
    return requestID, ok
}

func handleRequest(ctx context.Context) {
    if requestID, ok := requestIDFromContext(ctx); ok {
        fmt.Println("Handling request with ID:", requestID)
    } else {
        fmt.Println("No request ID found in context")
    }
}

func main() {
    ctx := context.Background()
    ctx = withRequestID(ctx, "12345")
    handleRequest(ctx)
}

在上述代碼中，使用 context.WithValue 在 context 中存儲了一個請求 ID，然後在處理函數中提取並使用這個請求 ID。這使得可以在不修改函數簽名的情況下，方便地傳遞和使用請求範圍內的變量。

Context 的高級使用

傳遞元數據

有時需要在多個 goroutine 之間傳遞元數據（如請求 ID、用戶認證信息等）。context 可以用來安全地傳遞這些信息。

package main

import (
    "context"
    "fmt"
)

type key int

const requestIDKey key = 0

// withRequestID 將請求 ID 存儲在 context 中
func withRequestID(ctx context.Context, requestID string) context.Context {
    return context.WithValue(ctx, requestIDKey, requestID)
}

// requestIDFromContext 從 context 中提取請求 ID
func requestIDFromContext(ctx context.Context) (string, bool) {
    requestID, ok := ctx.Value(requestIDKey).(string)
    return requestID, ok
}

func handleRequest(ctx context.Context) {
    if requestID, ok := requestIDFromContext(ctx); ok {
        fmt.Println("Handling request with ID:", requestID)
    } else {
        fmt.Println("No request ID found in context")
    }
}

func main() {
    ctx := context.Background()
    ctx = withRequestID(ctx, "12345")
    handleRequest(ctx)
}

在上述代碼中，使用 context.WithValue 在 context 中存儲了一個請求 ID，然後在處理函數中提取並使用這個請求 ID。這使得我們可以在不修改函數簽名的情況下，方便地傳遞和使用請求範圍內的變量。

處理併發操作

在處理併發操作時，context 可以幫助控制 goroutine 的生命週期，確保在請求取消時能夠正確地終止所有相關操作。

package main

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

func worker(ctx context.Context, name string) {
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("%s: received cancellation signal\n", name)
            return
        default:
            fmt.Printf("%s: working...\n", name)
            time.Sleep(1 * time.Second)
        }
    }
}

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

    go worker(ctx, "worker1")
    go worker(ctx, "worker2")

    time.Sleep(3 * time.Second)
    fmt.Println("Cancelling context")
    cancel()

    // 等待一段時間以確保所有 goroutine 都能收到取消信號
    time.Sleep(1 * time.Second)
}

在上述代碼中，創建了兩個併發執行的 worker goroutine，並使用 context.WithCancel 創建了一個可取消的 context。當主函數調用 cancel() 時，所有的 worker 都會接收到取消信號並停止工作。

使用 Context 進行超時控制

在處理外部資源（如網絡請求、數據庫查詢）時，設置超時是非常重要的。使用 context 可以輕鬆地實現超時控制。

package main

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

func fetchURL(ctx context.Context, url string) error {
    req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
    if err != nil {
        return err
    }

    client := &http.Client{}
    resp, err := client.Do(req)
    if err != nil {
        return err
    }
    defer resp.Body.Close()

    fmt.Printf("Fetched %s: %s\n", url, resp.Status)
    returnnil
}

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

    url := "https://www.example.com"
    if err := fetchURL(ctx, url); err != nil {
        fmt.Println("Error fetching URL:", err)
    }
}

在上述代碼中，使用 context.WithTimeout 創建了一個帶有超時功能的 context，並將其傳遞給 fetchURL 函數。如果請求超過了 2 秒，context 將自動取消請求，避免了長時間的阻塞。

Context 的最佳實踐

在使用 Context 時，有一些重要的最佳實踐需要遵循：

不要將 Context 存儲在結構體中

// 錯誤示例
type Service struct {
    ctx context.Context    // 不要這樣做
}

// 正確示例
type Service struct {
    // ... 其他字段
}

func (s *Service) DoSomething(ctx context.Context) error {
    // 在方法參數中傳遞 context
}

Context 應該是函數的第一個參數

// 推薦的方式
func DoSomething(ctx context.Context, arg string) error {
    // ...
}

// 不推薦的方式
func DoSomething(arg string, ctx context.Context) error {
    // ...
}

使用 context.WithValue 時要謹慎

// 推薦：使用自定義類型作爲 key
type contextKey string
const userIDKey contextKey = "userID"

// 不推薦：直接使用內置類型作爲 key
ctx = context.WithValue(ctx, "userID", "123") // 避免這樣做

常見陷阱和注意事項

避免傳遞 nil context：總是使用 context.Background() 或 context.TODO() 作爲起點
注意 context 取消的傳播：父 context 取消時，所有子 context 都會被取消
合理使用超時設置：避免設置過長或過短的超時時間
正確處理 context.Done()：在使用 select 語句時，確保正確處理取消信號

總結

通過合理地使用 context，可以使 Go 代碼變得更智能、更安全，並且更易於擴展。 context 提供的取消信號、超時機制和變量傳遞功能，使得可以更好地控制並管理 goroutine 之間的交互，從而編寫出更加健壯和可靠的程序。在實際開發中，建議儘量使用 context 處理與請求範圍相關的操作，以提高代碼的可維護性和擴展性。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/uW7NW-BMk32fDxLiM8PByQ