深入理解 Golang Context 包：提高程序并发控制与超时控制

在 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 请求，并在超时时间到达后自动取消该请求。

3. 总结

通过本文的介绍，我们了解了 Golang 的 Context 包的基本原理、使用方法和示例代码。Context 包是 Go 语言中实现并发控制和超时控制的重要工具之一，可以帮助我们更加灵活地控制程序的执行。在实际应用中，我们可以使用 Context 包来实现一些复杂的功能，例如控制数据库连接池、处理 HTTP 请求和 gRPC 请求等。通过学习和使用 Context 包，我们可以更好地实现并发控制和超时控制，提高程序的可靠性和稳定性。
在使用 Context 包时，需要注意以下几点：

1. Context 对象是线程安全的，可以被多个 goroutine 同时访问。
2. 可以使用 WithCancel、WithDeadline 和 WithTimeout 函数来创建 Context 对象，并使用 context.Background 函数创建根 Context 对象。
3. 在进行 goroutine 间的通信时，可以使用 Context 对象来传递消息和控制 goroutine 的执行。
4. 在使用第三方库时，需要注意该库是否支持 Context 功能，以便正确地使用 Context 包。
5. 在使用 Context 包时，需要避免出现 Context 滥用的情况，应该根据实际需要仅仅传递必要的 Context 对象，避免将 Context 对象传递到太多的函数中，导致程序难以维护和调试。

通过本文的介绍，相信你已经对 Golang 的 Context 包有了更深入的了解和掌握。希望你可以在实际应用中灵活运用 Context 包，提高程序的效率和可靠性。

转自：https://juejin.cn/post/7228071726206353465