rk呓语

reposkeeper

Golang Context 源码解析

Golang Context 是一个较为常用的组件。经常用作在协程中附带值,或者传递结束信号。

Context 的实现并不难理解。但是 Context 的使用,还是有颇多的限制,需要在业务代码中嵌入对于 Context 的处理才可以。

Context 可能用在什么地方呢?

比如:我们有一个 HTTP 服务器,有一个接口比较耗时。这时,我们可以在接口中实现一个 Context 并层层传递到后面执行的函数中。这样,当超时时间到,或者客户端取消请求时,通过 Context 可以通知被调用函数结束后面的计算。

再比如:一个 daemon 服务,当接收到 SIGINT 时,那么就使用 Context 通知所有运行的协程结束工作,优雅退出。

我们先来看一个例子:

package main

import (
    "context"
    "fmt"
)

func main() {
    gen := func(ctx context.Context) <-chan int {
        dst := make(chan int)
        n := 1
        go func() {
            for {
                select {
                case <-ctx.Done():
                    return // returning not to leak the goroutine
                case dst <- n:
                    n++
                }
            }
        }()
        return dst
    }

    ctx, cancel := context.WithCancel(context.Background())
    defer cancel() // cancel when we are finished consuming integers

    for n := range gen(ctx) {
        fmt.Println(n)
        if n == 5 {
            break
        }
    }
}

上面是官方文档中的一个例子。当从 gen() 中获取了5个元素的时候,则退出循环。此时 defer 会触发 cnacel(), 则 goroutine 会一并退出;

Context 实现

context

context 整体实现的结构是如上图所描述的。

首先,定义了一个,Context Interface,所有的实现都要基于这个 interface。

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key interface{}) interface{}
}
  • Deadline();如果设置了取消时间,那么会将时间返回,如果没有设置,则返回 ok == false
  • Done();如果 Context 被取消,则返回一个关闭的 channel,如果 Context 不能取消,则返回 nil
  • Err(); 如果 Context 没有被取消,则返回 nil,如果取消了,则返回错误,多次调用返回同样的错误
  • Value(key interface{});返回和key相关的value内容

Context Interface 要求上述几个接口,Context 的具体实现均基于此。

emptyCtx

emptyCtx 是第一个实现 Context Interface 的数据类型。但是如其名,它只是一个空的 Context。主要是用来给其他真正有作用的 Context 当做 parent 来用;

var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

首先用 emptyCtx 生成了两个供外部使用的 parent context:Background 和 TODO;
这两个 Context 可以作为所有其他 Context 的 Parent 来使用;比如: WithCancel

WithCancel

// 创建一个带有取消能力的 Context
// 参数传入 parent context
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := newCancelCtx(parent) // 新建一个 CancelContext
    propagateCancel(parent, &c) // 向 parent context 注册 cancel 函数
    return &c, func() { c.cancel(true, Canceled) } // 返回
}

// 返回一个初始化的 CancelContext
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

// 向 parent context 注册 cancel
// 可以在 parent 触发 cancel 的时候,同时 触发 child 的 cancel
func propagateCancel(parent Context, child canceler) {
    if parent.Done() == nil {
        return // 如果 parent 不具备取消能力,则直接返回(比如:TODO、Background)
    }
    if p, ok := parentCancelCtx(parent); ok { // 找到最近的 parent cancelCtx
        p.mu.Lock()
        if p.err != nil {
            // parent 已经取消了,则标记 child 取消
            child.cancel(false, p.err)
        } else {
           // 如果 parent 没有 child ,则初始化
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            // 将自己加到 parent 中
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
      // 如果没找到最近的 cancelContext,那么说明自己就是 parent
      // 启动协程 等待信号
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

// 沿着 Context 的生成链路,去查找最近的一个 cancelCtx 并返回
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    for {
        switch c := parent.(type) {
        case *cancelCtx:
            return c, true
        case *timerCtx:
            return &c.cancelCtx, true
        case *valueCtx:
            parent = c.Context
        default:
            return nil, false
        }
    }
}

上面的一段代码可以看到,创建一个 CancelContext 需要几个步骤:

  1. 初始化 context
  2. 将自己注册到最近的 parent 上;
  3. 如果没有 parent,则自己负担 parent 的职责;

所以,看到这里,可以知晓,并不是一类 Context 才能相互继承。整个继承的链条中,可以夹杂多种Context;

WithDeadline & WithTimeout

除去 CancelContext 还有三种 Context:WithDeadline、WithTimeout、WithValue。
其中,WithTimeout 仅仅是 WithDeadline 的变种,其具体实现是一样的。


func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        // 如果 parent 有 deadline,还比自己更短,那自己也省得定时了
        // 直接创建 WithCancel 就好
        return WithCancel(parent)
    }
    // 创建一个 timerCtx,是 cancelCtx 包含了一个 deadline
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    // 注册自己的 cancel
    propagateCancel(parent, c)
    dur := time.Until(d)  // 拿到距离 Deadline 的中间时间
    if dur <= 0 { // 已经触发 Deadline
        c.cancel(true, DeadlineExceeded) // deadline has already passed
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
       // 注册 timer 函数,过了 dur 时间后,就自动调用 cancel
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

上面就是整个 WithDeadline 的实现,其实相对比较容易理解。仅仅基于 WithCancel 上增加了定时取消的功能。

WithValue

最后一类 Context 就是用来传递值的。在实际使用中,只有 WithValue 这一类的 Context 可以携带值来传递。但是因为只要是使用 valueCtx 创建的 child context,不管是什么类型,都可以找到所传递的值

func (c *valueCtx) Value(key interface{}) interface{} {
    if c.key == key { // 如果当前 context 的 key 和要找寻的 key 一样,则返回 value
        return c.val
    }
    return c.Context.Value(key) // 如果不是,则调用 parent 的 Value()
}

上面的源码可以看到,Value 函数如果没有在当前 context 中找到 key,会一直沿着继承树去查找,知道找到最上面的 emptyCtx。
这样,就可以达到不管在下面的哪一层,都可以使用上面传递的 key value。

Context 在日常使用中,对于信号传递,值传递时,十分好用。日常开发中,如果派生的 Context ⛓太长,可以使用 fmt.Println(ctx) 直接打印整个 Context 的链条。可以用作 Debug 的一种手段。

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

    ctx2, _ := context.WithCancel(
        context.WithValue(
            context.WithValue(ctx, "T-key1", "Value 1"),
            "T-key2",
            "Value 2"))

    fmt.Println(ctx2)
    cancel()
}
// output => 
// context.Background.WithCancel.WithValue(type string, val Value 1).WithValue(type string, val Value 2).WithCancel

Golang sync.Map 源码解析

上一篇

Golang sync.Mutex 源码解析

下一篇
评论
发表评论 说点什么
还没有评论
120
0