Go学习系列Part8 深入学习Context

概览

Context主要用于在 API 边界之间传递截止日期、取消信号以及其他请求范围的值。它解决了在并发编程中常见的一些问题，比如如何优雅地取消长时间运行的操作，或者在多个 goroutine 之间共享和传递数据。

为什么需要Context？

goroutine 是轻量级的并发单元，可以很容易地启动大量的 goroutine 来处理并发任务。然而，当涉及到需要取消长时间运行的操作、管理超时或者传递跨多个 API 的数据时，仅仅使用 goroutine 就变得不够了。这时，就需要 Context 来帮助管理这些并发操作。

使用场景

客户端请求处理
存储和访问控制
性能优化
并发任务管理
连接池和资源管理
请求跟踪和日志记录

官方文档：pkg.go.dev/context#Con…

基础认识

接口定义

type Context interface {
    Deadline() (deadline time.Time, ok bool)    
    Done() <-chan struct{}
    
    Err() error
    // 返回Context结束的原因。只在Done对应的Channel关闭时返回非空值。
    // 如果Context被取消，返回 context.Canceled
    // 如果Context超时，返回context.DeadlineExceeded
    
    Value(key interface{}) interface{}
    // 从 Context 中获取键对应的值。如果未设置 key 对应的值则返回 nil。
    // 以相同 key 多次调用会返回相同的结果。
}

Deadline()：返回完成工作的截止时间，表示上下文应该被取消的时间。 ok为false时，表示没有设置截止时间。

Done()：返回一个Channel，用于判断上下文是否完成，当工作完成时被关闭，表示Context被取消。如果无法取消，则可能返回 nil。多次调用，返回同一个Channel。

创建

context包提供两个创建的函数：

func TODO() Context：返回一个非nil但空的Context，不确定要使用哪个上下文时，可以将其用作占位符。
func Background() Context：返回一个非nil但空的Context。不会被取消，没有值，且没有截止时间。通常在主函数、初始化、测试时使用，并作为传入请求的顶级上下文。

派生

通过func WithValue(parent Context, key, val interface{}) Context方法，可以让Context携带键值对。返回一个与传入Context为派生关系的Context。

parent：父Context
key：一个不可导出的空接口类型 interface{} 的值，用作键。
val：键值，可以是任何类型。

通过func WithoutCancel(parent Context) Context可以获取取消父项时未取消的父项的副本。返回的Context不返回Deadline或Err，其Done的channel为nil。

带取消功能的Context

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

WithCancel()：返回一个继承parent的新Context和一个取消函数，调用取消函数时，发送一个取消信号给ctx及其派生Context。用于创建带取消功能的Context。
WithCancelCause：基于WithCancel拓展的，在取消时传递一个取消原因。

ctx, cancel := context.WithCancelCause(parent)
cancel(myError)
ctx.Err() // returns context.Canceled
context.Cause(ctx) // returns myError

带期限的Context

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)
func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc)

WithDeadline：当截止时间到达时，或者当调用 cancel 函数时，它会发送一个取消信号给 ctx 及其所有派生的子上下文。 2.WithDeadlineCause：类似WithDeadline，在超过截止日期时设置返回Context的原因。

带超时限制的Context

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, CancelFunc)

WithTimeout：用于创建一个具有超时时间限制的上下文（context）。当超时时间到达时，或者当调用取消函数时，新的上下文会被取消，同时它的 Done 通道会被关闭，并且 Err 方法将返回 context.DeadlineExceeded。
WithTimeoutCause：类似WithTimeout, 在超时到期时设置返回Context的原因

总结

Context 的创建是先通过BackGround()创建空Context，然后在此基础上，调用With开头的函数，不断的派生Context，从而生成一颗Context树，BackGround()创建的空Context为根。

使用实践

传递共享数据

通过WithValue()实现：

// UserIDKey 是用于在 context 中存储用户 ID 的键
type UserIDKey struct{}

// handleRequest 模拟处理 HTTP 请求的函数
func handleRequest(ctx context.Context) {
	// 从 context 中获取用户 ID
	userID := ctx.Value(UserIDKey{})
	if userID == nil {
		fmt.Println("User ID not found in context")
		return
	}

	// 假设 userID 是一个 int 类型，我们需要将其断言回 int
	userIDValue, ok := userID.(int)
	if !ok {
		fmt.Println("Failed to assert user ID to int")
		return
	}

	// 假设这是业务逻辑的一部分，我们在这里使用用户 ID
	fmt.Printf("Processing request for user: %d\n", userIDValue)

	// 调用另一个函数，并将相同的 context 传递给它
	processData(ctx)
}

// processData 是处理数据的函数，它也使用相同的 context
func processData(ctx context.Context) {
	// 从 context 中获取用户 ID
	userID := ctx.Value(UserIDKey{})
	if userID == nil {
		fmt.Println("User ID not found in context")
		return
	}

	// 假设 userID 是一个 int 类型，我们将其断言回 int
	userIDValue, ok := userID.(int)
	if !ok {
		fmt.Println("Failed to assert user ID to int")
		return
	}

	// 在这里，我们使用用户 ID 进行一些数据处理
	fmt.Printf("Processing data for user: %d\n", userIDValue)
}

func main() {
	// 创建一个带有用户 ID 的 context
	ctx := context.WithValue(context.Background(), UserIDKey{}, 123)

	// 调用 handleRequest 函数，并将带有用户 ID 的 context 传递给它
	handleRequest(ctx)
}

虽然 context 可以用于传递一些元数据，但它并不是用来存储大量数据或用于在 goroutine 之间共享数据的。对于大数据或复杂的共享状态，应该使用其他机制，如通道（channel）、互斥锁（mutex）或其他的并发安全的数据结构。

取消信号

结合HTTP请求，实现一个请求取消功能。模拟了一个长时间运行的 HTTP 请求，并使用 context 来设置超时时间并在超时时取消请求。

// 模拟长时间运行的 HTTP 请求
func longRunningRequest(ctx context.Context, url string) error {
	// 创建一个新的 HTTP 请求
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return err
	}

	// 发送请求
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		// 如果是因为上下文被取消而导致的错误，则直接返回
		if err == context.Canceled {
			return err
		}
		return fmt.Errorf("failed to make request: %w", err)
	}
	defer resp.Body.Close()

	// 读取响应体（这里只是简单地读取并丢弃，实际情况下会处理响应内容）
	_, err = io.Copy(io.Discard, resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read response: %w", err)
	}

	return nil
}

func main() {
	// 假设我们有一个需要很长时间才能响应的 URL
	url := "http://example.com/long-running-request"

	// 创建一个带有 5 秒超时的上下文
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel() // 确保最终取消上下文

	// 启动 goroutine 来发送 HTTP 请求
	go func() {
		err := longRunningRequest(ctx, url)
		if err != nil {
			if err == context.DeadlineExceeded {
				fmt.Println("Request timed out")
			} else if err == context.Canceled {
				fmt.Println("Request cancelled")
			} else {
				fmt.Printf("Request failed: %v\n", err)
			}
		} else {
			fmt.Println("Request succeeded")
		}
	}()

	// 主 goroutine 等待一段时间以便看到请求的结果
	time.Sleep(10 * time.Second)
        // 输出 Request succeeded
}

longRunningRequest 函数来模拟发送一个长时间运行的 HTTP 请求，使用http.NewRequestWithContext来创建一个带有 context 的 HTTP 请求。
当请求在5秒内没有响应，此时，context在5秒后超时，触发取消，因此会看到"Request timed out"输出，表示请求因超时而被取消。

携带截止日期并跨多个 goroutine 传递

定义了一个 longRunningTask 函数，它模拟一个需要 5 秒才能完成的任务。创建了一个带有 3 秒截止日期的 context。然后，启动一个 goroutine 来执行 longRunningTask，并传入这个 context。等待 4 秒后取消 context。由于 longRunningTask 的截止日期超过了我们设置的 3 秒，所以调用 cancel() 时，longRunningTask 会收到取消信号，并打印出 "Task 1 cancelled"。

// longRunningTask 模拟一个长时间运行的任务
func longRunningTask(ctx context.Context, taskID int) {
	select {
	case <-time.After(5 * time.Second):
		fmt.Printf("Task %d completed\n", taskID)
	case <-ctx.Done():
		fmt.Printf("Task %d cancelled\n", taskID)
		// 如果需要，可以在这里清理资源
	}
}

func main() {
	// 创建一个带有 3 秒截止日期的 context
	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(3*time.Second))
	defer cancel() // 确保在函数返回时调用 cancel

	// 启动一个 goroutine 来执行一个长时间运行的任务
	go longRunningTask(ctx, 1)

	// 等待一段时间，以便看到任务是否完成或被取消
	select {
	case <-time.After(4 * time.Second):
		fmt.Println("Main: Waited for 4 seconds. Cancelling the context...")
		cancel() // 取消 context，这将导致 longRunningTask 中的任务被取消
	case <-ctx.Done():
		fmt.Println("Main: Context was already done")
	}

	// 等待一段时间以确保 goroutine 完成
	time.Sleep(1 * time.Second)
}

Go语言深入学习Context

概览