Golang基础笔记十三之context

本文首发于公众号：Hunter后端

原文链接：Golang基础笔记十三之context

在 Golang 里，context 包提供了很多比如传递截止时间、取消信号、传递数据等操作的标准方式，用于在跨 API 边界、进程和 goroutine之间进行。

这一篇笔记详细介绍一下 context 包相关的一些操作。

以下是本篇笔记目录：

1. Context 接口及作用
2. 取消传播
3. 超时控制
4. 截止时间
5. 传递数据

1、Context 接口及作用

1. Context 接口

Context 是 context 包下的一个接口，其定义如下：

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

1. Deadline() 返回上下文被取消的时间
2. Done() 返回一个通道，当上下文被取消时关闭
3. Error() 返回上下文被取消的原因
4. Value() 可以获取相应的键值对数据

下面所有的操作都是基于 context.Context 实现。

2. context 相关函数

我们可以使用 context 创建一些函数，用来实现取消、超时控制、传递数据等操作。

1) context.Background()

根上下文，可以作为所有上下文的起点，当我们需要实现取消、超时控制等功能，都需要定义一个父级上下文，这个时候我们就可以使用 context.Backgroud() 来创建。

2) context.TODO()

当我们不确定使用哪个上下文时，就可以使用 context.TODO()，但是在源代码中，它的实现与 context.Backgroud() 是一样的逻辑。

3) context.WithCancel(parent)

创建可以取消的上下文，参数是父级上下文，当我们调用一个函数，并且希望在某些时候取消这个调用，比如超时，这个时候我们可以使用这个函数，并手动进行取消。

4) context.WithTimeout(parent, timeout)

创建有超时的上下文，参数是父级上下文和 time.Duration，当我们调用一个函数，并且希望在多久以后可以自动取消，可以使用这个函数。

5) context.WithDeadline(parent, d)

创建有截止时间的上下文，参数是父级上下文和 time.Time，当我们调用一个函数，并且希望在某个具体的时间点可以自动取消，可以使用这个函数。

6) context.WithValue(parent, key, value)

创建带有键值对的上下文，我们希望通过上下文传递数据，就可以使用这个函数。

这里介绍了 context 相关的一些函数，接下来我们以具体的代码为示例，分别用这些函数来实现对应的功能。

2、 取消传播

我们使用 context.WithCancel() 函数实现取消传播的功能。

其实就是取消函数执行的操作，为什么会说取消传播，因为上下文可以在函数调用中一层一层传递，当我们取消了根上下文，所有调用链中的上下文都会被取消。

context.WithCancel() 的使用代码示例如下：

ctx, cancel := context.WithCancel(context.Background())

这个函数返回两个结果，一个是上下文参数，一个是取消函数，我们可以使用取消函数在特定节点取消这个上下文，这里取消操作需要我们手动执行。

下面我们实现一个取消操作，我们调用两个函数，这两个函数随机执行一段时间，然后某个函数先返回结果，返回之后我们立马执行上下文的取消操作，

package main

import (
    "context"
    "fmt"
    "math/rand"
    "time"
)

func F1(ctx context.Context, ch chan string) {
    sleepSeconds := rand.Intn(3)
    time.Sleep(time.Duration(sleepSeconds) * time.Second)
    ch <- "f1 result"
}

func F2(ctx context.Context, ch chan string) {
    sleepSeconds := rand.Intn(3)
    time.Sleep(time.Duration(sleepSeconds) * time.Second)
    ch <- "f2 result"
}

func CallF1(ctx context.Context, ch chan string) {
    chF1 := make(chan string)
    go F1(ctx, chF1)

    select {
    case result := <-chF1:
        ch <- result
    case <-ctx.Done():
        fmt.Println("F1 函数调用超时")
    }
}

func CallF2(ctx context.Context, ch chan string) {
    chF2 := make(chan string)
    go F2(ctx, chF2)

    select {
    case result := <-chF2:
        ch <- result
    case <-ctx.Done():
        fmt.Println("F2 函数调用超时")
    }
}

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

    ch1 := make(chan string)
    ch2 := make(chan string)
    go CallF1(ctx, ch1)
    go CallF2(ctx, ch2)

    select {
    case r1 := <-ch1:
        fmt.Println("f1 调用完成: ", r1)
        cancel()
    case r2 := <-ch2:
        fmt.Println("f2 调用完成: ", r2)
        cancel()
    }
}

在这里，我们目标调用函数为 F1() 和 F2()，使用 CallF1() 和 CallF2() 作为中间调用函数，在其中使用 select-case 等待目标函数返回结果，并同时监听 ctx.Done() 判断 ctx 是否已经取消。

同时在 main() 函数中监听 ch1 和 ch2 判断哪个通道先返回结果，监听到返回结果马上取消 ctx 上下文，main() 函数就可以接着往下执行，避免两个目标函数其中一个长时间执行。

这里需要注意的是，中间函数 CallF1() 和 CallF2() 虽然被取消了，但 F1()、F2() 作为执行的 goroutine 并没有取消执行，如果有取消的需求，可以在 F() 函数内部伺机监听 ctx.Done() 以提前退出函数。

3、 超时控制

我们使用 context.WithTimeout(parent, timeout) 可以实现超时控制。

下面实现一个功能为，调用某个函数，并给一秒的超时时间，超过时间则进行超时处理，避免长时间堵塞，其中目标执行函数为 TargetFunc，其中，随机 sleep 三秒内的时间。

整体代码如下：

package main

import (
    "context"
    "fmt"
    "math/rand"
    "time"
)

func TargetFunc() string {
    sleepSeconds := rand.Intn(3)
    time.Sleep(time.Duration(sleepSeconds) * time.Second)
    return "result"
}

func CallFunc(ch chan string) {
    ch <- TargetFunc()
}

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

    ch := make(chan string)
    defer close(ch)

    go CallFunc(ch)

    select {
    case result := <-ch:
        fmt.Println("result: ", result)
    case <-ctx.Done():
        fmt.Println("函数调用超时")
    }
}

其中，TargetFunc 函数内部的逻辑为随机休息三秒内的时间，可能导致超时，也可能不超时。

CallFunc() 函数用作中间函数，用于获取目标函数调用结果并通过 channel 返回。

在 main 函数中，首先定义一个有超时时间的上下文，设置了一秒超时，然后定义一个通道，将其传给 CallFunc() 用于返回数据。

之后通过 select-case 操作，进入循环等待状态，用于判断 goroutine 中的 channel 和 ctx 的超时哪个先返回结果。

如果 TargetFunc 函数调用时间过长，ctx 超时控制的上下文先到期，那么则会打印出 函数调用超时 的信息，否则会打印出函数调用返回的结果。

4、 截止时间

我们可以使用 context.WithDeadline(parent, d) 函数实现一个有截止时间的上下文，它的调用参数第二个为 d，是一个具体的 time.Time 类型。

但其实，在背后的源码中，context.WithTimeout() 函数会将其中的 time.Duration 参数处理通过 time.Now().Add(timeout) 的方式处理成 time.Time 然后再调用 WithDeadline 函数，所以超时控制和截止时间这两个函数在本质上的逻辑是一致的。

所以这里我们的代码示例，也只是把输入的参数修改一下，即可实现功能，如下：

deadline := time.Now().Add(1 * time.Second)
ctx, cancel := context.WithDeadline(context.Background(), deadline)

5、 传递数据

我们可以通过 context.WithValue() 函数实现传递数据的功能，这个函数接收三个参数，父级上下文，key 和 value，以下是使用示例：

ctx := context.Background()

ctxWithValue := context.WithValue(ctx, "str1", "value1")
ctxWithValue = context.WithValue(ctxWithValue, "slice", []int{1, 2, 3})

在获取数据的时候，我们可以使用 ctx.Value() 函数获取：

func ProcessCtxValue(ctx context.Context) {
    value1 := ctx.Value("str1")
    fmt.Println("str1 value: ", value1)

    value2 := ctx.Value("slice")
    fmt.Println("slice value: ", value2)
}

但是在真正用到这些数据的时候，我们还需要对 value 进行类型判断，可以直接如下操作：

func ProcessCtxValue(ctx context.Context) {
    value1, ok := ctx.Value("str1").(string)
    if ok {
        fmt.Println("str1 value: ", value1)
    }

    value2 := ctx.Value("slice").([]int)
    if ok {
        fmt.Println("slice value: ", value2, value2[1])
    }
}