Go 中的 Context实现原理以及正确使用方式

Chicken Run

1431人浏览 · 2024-11-05 22:07:38

Chicken Run · 2024-11-05 22:07:38 发布

在 Go 语言中，Context 包是一种非常常用的工具，它被用来管理 goroutine 之间的通信和取消。本文将深入探讨Context 包的基本原理，包括使用场景、原理和一些最佳实践。

1. 基本原理

1.1 Context 包的介绍

在 Go 语言中，Context 包是用于传递请求范围数据、取消信号和截止时间的机制。它通常被用来处理 goroutine 之间的通信和取消。Context 包是 Go 语言内置的，它可以很方便地使用，而不需要额外的依赖。
Context 包是一个轻量级的工具，它提供了一个标准的接口，用于在 goroutine 之间传递请求范围的数据、取消信号和截止时间。Context 包实现了一种类似于树状结构的数据结构，其中每个节点都表示一个请求范围。每个节点都有一个唯一的 key-value 对，其中 key 是一个 interface{} 类型的值，而 value 则是任何类型的值。Context 包还提供了一个可选的超时机制，用于在一定时间后自动取消请求。
Context 包的核心是一个 Context 接口，它定义了一些方法，用于获取请求范围数据、取消请求和处理超时。

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

Deadline() 方法返回截止时间和一个布尔值，指示截止时间是否已经设置。
Done() 方法返回一个只读的 channel，当请求被取消或超时时，该 channel 将被关闭。
Err() 方法返回一个错误，指示为什么请求被取消。
Value() 方法返回与给定key相关联的值，如果没有值，则返回 nil。

Context 包还提供了两个用于创建 Context 的函数：WithContext 和 Background。Background 函数返回一个空的 Context，而 WithContext 函数则根据给定的父 Context 创建一个新的 Context。
Context 包的基本原理是通过在 goroutine 之间传递 Context 来实现请求范围数据、取消信号和截止时间的管理。当一个 goroutine 创建了一个新的 goroutine 时，它将 Context 作为参数传递给新的 goroutine。新的goroutine 可以使用这个 Context 来访问请求范围数据、接收取消信号和处理超时。

1.2 Context 的创建

在 Golang 中，Context 可以通过 WithCancel、WithDeadline、WithTimeout 和 WithValue 等函数来创建。下面分别介绍这些函数的用法和注意事项。

1.2.1 WithCancel

WithCancel 函数可以用于创建一个 Context 对象，并返回一个可取消的上下文和一个取消函数。当调用取消函数时，会通知所有的 Context 对象和其子 Context 对象，使它们都取消执行。

 func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "time"
 )
 
 func main() {
     parent := context.Background()
     ctx, cancel := context.WithCancel(parent)
     go func() {
         select {
         case <-ctx.Done():
             fmt.Println(ctx.Err())
             return
         case <-time.After(5 * time.Second):
             fmt.Println("work done")
         }
     }()
     time.Sleep(10 * time.Second)
     cancel()
     time.Sleep(1 * time.Second)
 }

在上面的代码中，我们首先使用 context.Background() 函数创建一个根 Context 对象 parent，然后使用 WithCancel 函数创建一个子 Context 对象 ctx，并返回一个可取消的上下文和一个取消函数 cancel。接下来，我们在一个 goroutine 中使用 select 语句监听 Context 对象的 Done 方法和 time.After 函数的返回值，如果 Done 方法返回一个非 nil 的 error，则说明 Context 已经被取消，否则说明 time.After 函数已经超时。在主函数中，我们调用 cancel 函数来通知 Context 对象和其子 Context 对象，使它们都取消执行。最后，我们使用 time.Sleep 函数让程序等待一段时间，以便观察 Context 的执行情况。

1.2.2 WithDeadline

WithDeadline 函数可以用于创建一个 Context 对象，并返回一个截止时间和一个取消函数。当超过截止时间时，会自动通知所有的 Context 对象和其子 Context 对象，使它们都取消执行。

func WithDeadline(parent Context, deadline time.Time) (ctx Context, cancel CancelFunc)

下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "time"
 )
 
 func main() {
     parent := context.Background()
     ctx, cancel := context.WithDeadline(parent, time.Now().Add(5*time.Second))
     go func() {
         select {
         case <-ctx.Done():
             fmt.Println(ctx.Err())
             return
         case <-time.After(10 * time.Second):
             fmt.Println("work done")
         }
     }()
     time.Sleep(20 * time.Second)
     cancel()
     time.Sleep(1 * time.Second)
 }

在上面的代码中，我们首先使用 context.Background() 函数创建一个根 Context 对象 parent，然后使用 WithDeadline 函数创建一个子 Context 对象 ctx，并返回一个截止时间和一个取消函数 cancel。接下来，我们在一个 goroutine 中使用 select 语句监听 Context 对象的 Done 方法和 time.After 函数的返回值，如果 Done 方法返回一个非 nil 的 error，则说明 Context 已经被取消，否则说明 time.After 函数已经超时。在主函数中，我们调用 cancel 函数来通知 Context 对象和其子 Context 对象，使它们都取消执行。最后，我们使用 time.Sleep 函数让程序等待一段时间，以便观察 Context 的执行情况。

1.2.3 WithTimeout

WithTimeout 函数可以用于创建一个 Context 对象，并返回一个超时时间和一个取消函数。当超过超时时间时，会自动通知所有的 Context 对象和其子 Context 对象，使它们都取消执行。

func WithTimeout(parent Context, timeout time.Duration) (ctx Context, cancel CancelFunc)

下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "time"
 )
 
 func main() {
     parent := context.Background()
     ctx, cancel := context.WithTimeout(parent, 5*time.Second)
     go func() {
         select {
         case <-ctx.Done():
             fmt.Println(ctx.Err())
             return
         case <-time.After(10 * time.Second):
             fmt.Println("work done")
         }
     }()
     time.Sleep(20 * time.Second)
     cancel()
     time.Sleep(1 * time.Second)
 }

在上面的代码中，我们首先使用 context.Background() 函数创建一个根 Context 对象 parent，然后使用 WithTimeout 函数创建一个子 Context 对象 ctx，并返回一个超时时间和一个取消函数 cancel。接下来，我们在一个 goroutine 中使用 select 语句监听 Context 对象的 Done 方法和 time.After 函数的返回值，如果 Done 方法返回一个非 nil 的 error，则说明 Context 已经被取消，否则说明 time.After 函数已经超时。在主函数中，我们调用 cancel 函数来通知 Context 对象和其子 Context 对象，使它们都取消执行。最后，我们使用 time.Sleep 函数让程序等待一段时间，以便观察 Context 的执行情况。

1.2.4 WithValue

WithValue 函数可以用于创建一个 Context 对象，并返回一个包含指定值的 Context 对象。这个值可以是任意类型的数据，可以是基本类型、结构体或者指针等。需要注意的是，这个值只在当前 Context 对象及其子 Context 对象中有效，对于其他 Context 对象来说是不可见的。

func WithValue(parent Context, key interface{}, val interface{}) Context

下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
 )
 
 type userKey struct{}
 
 func main() {
     parent := context.Background()
     ctx := context.WithValue(parent, userKey{}, "admin")
     go func() {
         if user, ok := ctx.Value(userKey{}).(string); ok {
             fmt.Printf("user is %s\n", user)
         } else {
             fmt.Println("user is not found")
         }
     }()
     select {}
 }

在上面的代码中，我们首先使用 context.Background() 函数创建一个根 Context 对象 parent，然后使用 WithValue 函数创建一个子 Context 对象 ctx，并返回一个包含指定值的 Context 对象。接下来，我们在一个 goroutine 中使用 ctx.Value 函数获取 Context 对象中的值，并判断其类型是否为字符串类型。如果是，则输出其值，否则输出 “user is not found”。在主函数中，我们使用select语句使程序一直运行，以便观察 Context 的执行情况。

2. Context 的使用场景

2.1 并发控制

一个很典型的使用场景是，当我们需要同时启动多个 goroutine 进行任务处理时，我们可以使用 Context 来控制这些 goroutine 的执行。在每个 goroutine 中，我们都可以检测 Context 对象是否被取消，如果是，则退出 goroutine 的执行，否则继续执行。
下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "sync"
 )
 
 func worker(ctx context.Context, wg *sync.WaitGroup) {
     defer wg.Done()
     for {
         select {
         default:
             fmt.Println("work")
         case <-ctx.Done():
             return
         }
     }
 }
 
 func main() {
     parent := context.Background()
     ctx, cancel := context.WithCancel(parent)
     var wg sync.WaitGroup
     for i := 0; i < 3; i++ {
         wg.Add(1)
         go worker(ctx, &wg)
     }
     cancel()
     wg.Wait()
 }

在上面的代码中，我们首先使用 context.Background() 函数创建一个根 Context 对象 parent，然后使用 WithCancel 函数创建一个子 Context 对象 ctx，并返回一个取消函数 cancel。接下来，我们使用 sync.WaitGroup 来等待所有的 goroutine 执行完成。在主函数中，我们启动了三个 goroutine 来执行任务，同时使用 cancel 函数来通知这些 goroutine 取消执行。最后，我们使用Wait方法等待所有的 goroutine 执行完成。

2.2 超时控制

另一个典型的使用场景是，当我们需要对一个操作设置一个超时时间时，我们可以使用 Context 来控制这个操作的执行时间。在操作执行超时时，我们可以通知 Context 对象和其子 Context 对象取消执行。
下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "time"
 )
 
 func work(ctx context.Context) {
     for {
         select {
         default:
             fmt.Println("work")
         case <-ctx.Done():
             fmt.Println("work done")
             return
         }
     }
 }

 func main() {
     parent := context.Background()
     ctx, cancel := context.WithTimeout(parent, time.Second*5)
     defer cancel()
     work(ctx)
 }

在上面的代码中，我们首先使用 context.Background() 函数创建一个根 Context 对象 parent，然后使用 WithTimeout 函数创建一个子 Context 对象 ctx，并返回一个取消函数 cancel。在 work 函数中，我们启动一个无限循环，不断输出 “work”。同时，我们使用 select 语句来等待 Context 对象被取消。在主函数中，我们使用 defer 语句调用 cancel 函数，以确保 Context 对象被取消。由于我们在 WithTimeout 函数中设置了一个 5 秒的超时时间，因此当程序运行超过 5 秒时，work 函数就会停止执行。

2.3 数据库连接

在使用数据库连接时，我们通常需要保证连接池中的连接数量不会超过一定的阈值。如果连接池中的连接数量超过了阈值，则需要等待连接释放后再进行操作。在这种情况下，我们可以使用 Context 来控制连接的生命周期。
下面是一个示例代码：

package main
 
 import (
     "context"
     "database/sql"
     "fmt"
     "sync"
     "time"
 
     _ "github.com/go-sql-driver/mysql"
 )
 
 const maxConn = 5
 
 func main() {
     db, err := sql.Open("mysql", "root:root@tcp(127.0.0.1:3306)/test")
     if err != nil {
         panic(err)
     }
     defer db.Close()
 
     ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
     defer cancel()
 
     connCh := make(chan *sql.Conn, maxConn)
     var wg sync.WaitGroup
     for i := 0; i < maxConn; i++ {
         wg.Add(1)
         go func() {
             defer wg.Done()
             for {
                 select {
                 case <-ctx.Done():
                     return
                 default:
                     if len(connCh) < maxConn {
                         conn, err := db.Conn(ctx)
                         if err != nil {
                             fmt.Println(err)
                             return
                         }
                         connCh <- conn
                     }
                 }
             }
         }()
     }
     wg.Wait()
 }

在上面的代码中，我们首先使用 sql.Open 函数打开一个 MySQL 数据库的连接，并返回一个 DB 对象 db。接下来，我们使用 WithTimeout 函数创建一个 Context 对象 ctx，并设置一个超时时间为 5 秒。同时，我们创建一个容量为 maxConn 的 channel 对象 connCh，用于存储数据库连接。在g oroutine 中，我们使用 select 语句等待 Context 对象被取消。在每次循环中，我们检查连接池中连接的数量是否超过了阈值，如果没有，则使用 db.Conn 函数从连接池中获取一个新的连接，并将其存储到 connCh 中。最后，我们使用 sync.WaitGroup 等待所有的 goroutine 执行完成。

####2.4 HTTP 请求

在使用 HTTP 请求时，我们通常需要设置一个超时时间，以确保请求能够在规定的时间内得到响应。在这种情况下，我们可以使用 Context 来控制HTTP请求的执行时间。
下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "io/ioutil"
     "net/http"
     "time"
 )
 
 func main() {
     client := http.DefaultClient
     ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
     defer cancel()
     req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://www.example.com", nil)
     if err != nil {
         fmt.Println(err)
         return
     }
 
     resp, err := client.Do(req)
     if err != nil {
         fmt.Println(err)
         return
     }
     defer resp.Body.Close()
 
     body, err := ioutil.ReadAll(resp.Body)
     if err != nil {
         fmt.Println(err)
         return
     }
 
     fmt.Println(string(body))
 }

在上面的代码中，我们首先使用 http.DefaultClient 创建一个 HTTP 客户端对象 client。接下来，我们使用 WithTimeout 函数创建一个 Context 对象 ctx，并设置一个超时时间为 5 秒。同时，我们使用 http.NewRequestWithContext 函数创建一个 HTTP 请求对象 req，并将 Context 对象 ctx 作为参数传递给该函数。在 Do 函数中，我们会自动将 Context 对象 ctx 传递给 HTTP 请求，并在超时时间到达后自动取消该请求。

####2.5 gRPC 请求

在使用 gRPC 请求时，我们通常需要设置一个超时时间，以确保请求能够在规定的时间内得到响应。在这种情况下，我们可以使用 Context 来控制 gRPC 请求的执行时间。
下面是一个示例代码：

package main
 
 import (
     "context"
     "fmt"
     "log"
     "time"
 
     pb "github.com/example/helloworld"
     "google.golang.org/grpc"
 )
 
 const (
     address     = "localhost:50051"
     defaultName = "world"
 )
 
 func main() {
     conn, err := grpc.Dial(address, grpc.WithInsecure())
     if err != nil {
         log.Fatalf("did not connect: %v", err)
     }
     defer conn.Close()
 
     c := pb.NewGreeterClient(conn)
 
     ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
     defer cancel()
 
     r, err := c.SayHello(ctx, &pb.HelloRequest{Name: defaultName})
     if err != nil {
         log.Fatalf("could not greet: %v", err)
     }
     log.Printf("Greeting: %s", r.GetMessage())
 }

在上面的代码中，我们首先使用 grpc.Dial 函数创建一个 gRPC 客户端连接对象 conn。接下来，我们使用 pb.NewGreeterClient 函数创建一个 GreeterClient 对象 c。然后，我们使用 WithTimeout 函数创建一个 Context 对象 ctx，并设置一个超时时间为 5 秒。最后，我们使用 GreeterClient 对象 c 的 SayHello 函数发送一个 gRPC 请求，并将 Context 对象 ctx 作为参数传递给该函数。在 SayHello 函数中，我们会自动将 Context 对象 ctx 传递给 gRPC 请求，并在超时时间到达后自动取消该请求。

Go Context 到底放第一个参数传，还是放结构体里？

作为函数的第一个参数

优点：

明确性：将context作为第一个参数，清晰地表明了函数执行的上下文依赖，增强了代码的可读性和意图表达。这种方式符合Go的设计哲学，即显式优于隐式。
易于测试：测试时可以轻松创建并传递一个自定义的context.Context实例，便于控制测试中的超时和取消逻辑，无需修改函数签名或结构体定义。
标准一致性：Go的标准库广泛采用了这种模式，比如net/http包中的ServeHTTP方法。遵循这一标准使得代码风格统一，便于其他开发者理解和维护。

缺点：

参数列表增长：对于参数较多的函数，额外增加一个context.Context可能会让函数签名显得冗长，特别是当多个函数嵌套调用时，每一层都需要传递context。
侵入性：虽然增加了灵活性，但也意味着每一个需要考虑取消或超时逻辑的函数都需要调整，对于既有代码库的改造成本较高。

嵌入到结构体中

优点：

减少函数签名复杂度：将context.Context作为一个字段嵌入到结构体中，可以减少函数参数的数量，使函数签名更加简洁。
封装性：对于内部逻辑复杂的服务，将context隐藏在结构体内部，可以对外提供更加抽象和友好的接口，提高代码的封装性。

缺点：

测试复杂度增加：如果结构体中的context字段不是公开的，测试时可能需要通过构造特定的结构体实例来传递特定的上下文信息，这可能使得测试代码变得复杂。
灵活性降低：一旦context作为结构体的一部分，函数调用时就失去了直接控制context的能力，比如无法在运行时轻易改变超时时间或取消策略。

实践建议

常规操作：对于大多数情况，遵循Go标准库的推荐，将context作为函数的第一个参数传递是最佳选择。这样做既体现了Go的简洁和明了，也便于维护和测试。
高度封装的服务：在设计高度封装的内部服务或复杂的API时，可以考虑将context嵌入到结构体中，尤其是当需要在整个服务生命周期内管理上下文时。但需权衡好封装性和测试便利性之间的关系。
混合使用：在某些场景下，你可能会发现结合两者使用的效果更好。例如，在服务初始化阶段，将context作为结构体字段管理，而在服务的具体操作函数中，依然将context作为第一个参数传递，以保持操作的灵活性。

总之，选择将context放在哪里，应基于项目的具体需求、代码的可读性和维护性综合考量。无论哪种方式，关键在于理解并充分利用context机制，以提升程序的健壮性和可维护性。