从零开始学习 go - Context

从零开始学习技术
64 阅读4分钟

参考链接

Go语言高级特性：Context深入解读mp.weixin.qq.com/s/_IkMo8uBB…

学习目的

go 使用并发非常方便，但是并发应该通过什么来进行交流呢，这就是 Context

创建 Context

package main

import (
    "context"
    "fmt"
)

func main() {
    // 创建一个空的Context
    ctx := context.Background()

    // 创建一个带有取消信号的Context
    _, cancel := context.WithCancel(ctx)
    defer cancel() // 延迟调用cancel，确保在函数退出时取消Context

    fmt.Println("Context创建成功")
}

传递 Context

创建了一个带有取消信号的context，并将它传递给一个 goroutine 中的任务。

调用cancel函数，可以发送取消信号，中断任务的执行。

package main

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

func worker(ctx context.Context) {
    for {
       select {
       case <-ctx.Done():
          fmt.Println("收到取消信号，任务结束")
          return

       default:
          fmt.Println("正在执行任务...")
          time.Sleep(1 * time.Second)
       }

    }
}

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

    ctxWithCancel, cancel := context.WithCancel(ctx)

    defer cancel()

    go worker(ctxWithCancel)

    time.Sleep(3 * time.Second)
    cancel() // 发送取消信号
    time.Sleep(1 * time.Second)
}

使用 Context 实现取消

在这个例子中，使用time.After函数模拟一个长时间运行的任务。

如果任务在 2 秒内完成，就打印“任务完成”，否则在接收到取消信号后打印“收到取消信号，任务结束”。

package main

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

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

    go func() {
       select {
       case <-ctxWithCancel.Done():
          fmt.Println("收到取消信号，任务结束")
       case <-time.After(2 * time.Second):
          fmt.Println("任务完成")
       }
    }()

    time.Sleep(1 * time.Second)

    cancel() // 发送取消信号

    time.Sleep(1 * time.Second)
}

使用 Context 实现超时控制

在上面例子中，使用context.WithTimeout(parent, duration)函数创建了一个在 2 秒后自动取消的context

如果任务在 1 秒内完成，就打印“任务完成”，否则在 2 秒后打印“超时，任务结束”。

package main

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

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

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

    defer cancel()

    select {
    case <-ctxWithTimeout.Done():
       fmt.Println("超时，任务结束")
    case <-time.After(1 * time.Second):
       fmt.Println("任务完成")
    }
}

使用 WithValue 传递值

在上面例子中，使用context.WithValue(parent, key, value)函数在context中传递了一个键值对。

使用ctx.Value(key)函数可以获取传递的值。

package main

import (
    "context"
    "fmt"
)

type key string

func main() {
    ctx := context.WithValue(context.Background(), key("name"), "Alice")
    value := ctx.Value(key("name"))
    fmt.Println("Name:", value)
}

使用 Context 处理 HTTP 请求

使用http.Request结构体中的Context()方法获取到请求的context，并在处理请求时设置了 3 秒的超时时间。\

如果请求在 3 秒内完成，就返回“Hello, World!”，否则返回一个错误。

package main

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

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

    select {
    case <-time.After(3 * time.Second):
       fmt.Fprintln(w, "Hello, World!")
    case <-ctx.Done():
       err := ctx.Err()
       fmt.Println("处理请求:", err)
       http.Error(w, err.Error(), http.StatusInternalServerError)
    }
}

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

处理 HTTP 请求的超时与取消

使用context.WithTimeout(parent, duration)函数为每个请求设置了 2 秒的超时时间。

如果请求不能在 2 秒内完成，就返回一个错误。

package main

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

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

    select {
    case <-time.After(3 * time.Second):
       fmt.Fprintln(w, "Hello, World!")
    case <-ctx.Done():
       err := ctx.Err()
       fmt.Println("处理请求:", err)
       http.Error(w, err.Error(), http.StatusInternalServerError)
    }
}

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

使用 Context 控制数据库查询的超时

使用context.WithTimeout(parent, duration)函数为数据库查询设置了 2 秒的超时时间，确保在 2 秒内完成查询操作。

如果查询超时，程序会收到context的取消信号，从而中断数据库查询。

package main

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

    _ "github.com/go-sql-driver/mysql"
)

func main() {
    db, err := sql.Open("mysql", "username:password@tcp(localhost:3306)/database")

    if err != nil {
       fmt.Println(err)
       return
    }

    defer db.Close()

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

    defer cancel()

    rows, err := db.QueryContext(ctx, "SELECT * FROM users")

    if err != nil {
       fmt.Println("查询出错:", err)
       return
    }

    defer rows.Close()

    // 处理查询结果
}

使用 Context 取消长时间运行的数据库事务

在这个例子中，使用context.WithTimeout(parent, duration)函数为数据库事务设置了 2 秒的超时时间。

同时使用 goroutine 监听context的取消信号。

在 2 秒内事务没有完成，程序会收到context的取消信号，从而回滚事务。

package main

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

    _ "github.com/go-sql-driver/mysql"
)

func main() {

    db, err := sql.Open("mysql", "username:password@tcp(localhost:3306)/database")

    if err != nil {
       fmt.Println(err)
       return
    }
    defer db.Close()

    tx, err := db.Begin()
    if err != nil {
       fmt.Println(err)
       return
    }

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

    defer cancel()

    go func() {
       <-ctx.Done()

       fmt.Println("接收到取消信号，回滚事务")

       tx.Rollback()
    }()

    // 执行长时间运行的事务操作
    // ...

    err = tx.Commit()
    if err != nil {

       fmt.Println("提交事务出错:", err)

       return
    }

    fmt.Println("事务提交成功")
}

实现自定义的 Context 类型

在这个例子中，实现了一个自定义的MyContext类型，它包含了一个RequestID字段。

通过WithRequestID函数，可以在原有的context上附加一个RequestID值，然后在需要的时候获取这个值。

package main

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

type MyContext struct {
    context.Context
    RequestID string
}

func WithRequestID(ctx context.Context, requestID string) *MyContext {
    return &MyContext{
       Context:   ctx,
       RequestID: requestID,
    }
}

func (c *MyContext) Deadline() (deadline time.Time, ok bool) {
    return c.Context.Deadline()
}

func (c *MyContext) Done() <-chan struct{} {
    return c.Context.Done()
}

func (c *MyContext) Err() error {
    return c.Context.Err()
}

func (c *MyContext) Value(key interface{}) interface{} {
    if key == "requestID" {
       return c.RequestID
    }
    return c.Context.Value(key)
}

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

    ctxWithRequestID := WithRequestID(ctx, "12345")

    requestID := ctxWithRequestID.Value("requestID").(string)

    fmt.Println("Request ID:", requestID)
}

Context 的使用技巧

1.用的时候要记得主动调用 context.WithCancel 或 context.WithTimeout 等函数创建的 Context 的 Cancel 函数来通知子 goroutines 停止工作。

这样可以确保资源被及时释放，避免 goroutine 泄漏。

2.不要在函数签名中传递 Context

避免在函数的参数列表中传递 context.Context，除非确实需要用到它。

如果函数的逻辑只需要使用 Context 的部分功能，那么最好只传递它需要的具体信息，而不是整个 Context。

3.避免在结构体中滥用 Context

avatar
文章
阅读
粉丝