golang context包源码分析

章鱼编程
963 阅读5分钟

golang 源码分析系列之——Context包

Context包

context.Context 是 Go 语言在 1.7 版本中引入标准库的接口。引入context包主要包含两个方面的作用

  1. golang 没有thread local 变量，在不同goroutine 间参数传递没有标准的解决方案
  2. 不同go routine 间的控制管理(启动、取消、超时等)

那我们看下golang 的context包如何解决上面两个问题 首先我们看下context 接口的定义

type Context interface {
    // 返回context，超时时间，如果有设置超时时间
    Deadline() (deadline time.Time, ok bool)
    // 关闭这个context的channel
    Done() <-chan struct{}
    // 当返回值不为nil，表示context 已经cancel 或者超时
    Err() error
    // 存储在里面的key
    Value(key interface{}) interface{}
}

总共四个接口： 1. Deadline() 返回context的超时时间，第二个返回值表示是否支持超时 2. Done() 返回context关闭的channel, 这也是实现跨go routine 任务取消、停止的关键 3. Err() 当context 关闭、超时时，err会记录，其他情况Err()返回nil 4. Value() 根据key 返回value 这是context 作为参数传递载体的关键，下面会详细分析它的结构

Context：参数传递

Context作为参数传递载体，主要使用Value方法，在Context 包里对应的实现类是ValueCtx, 我们先看下ValueCtx源码

// A valueCtx carries a key-value pair. It implements Value for that key and
// delegates all other calls to the embedded Context.
type valueCtx struct {
    Context
    key, val interface{}
}

func WithValue(parent Context, key, val interface{}) Context {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    if key == nil {
        panic("nil key")
    }
    if !reflectlite.TypeOf(key).Comparable() {
        panic("key is not comparable")
    }
    return &valueCtx{parent, key, val}
}

...

func (c *valueCtx) Value(key interface{}) interface{} {
    if c.key == key {
        return c.val
    }
    return c.Context.Value(key)
}

创建ValueCtx的方法是context.WithValue(), 从方法和结构体的定义我们不难发现，ValueCtx底层的存储结构实际上是一个链表。 valueCtx结构图

package main

import (
    "context"
    "fmt"
)

const (
    KEY1 = "key1"
    KEY2 = "key2"
)

func main() {
    // Pass a context with a timeout to tell a blocking function that it
    // should abandon its work after the timeout elapses.
    ctx1 := context.WithValue(context.Background(), KEY1, 100)
    ctx2 := context.WithValue(ctx1, KEY2, 200)
        printVal(ctx2)
    
}

func printVal(ctx context.Context) {
    v1 := ctx.Value(KEY1)
    v2 := ctx.Value(KEY2)
    fmt.Println("v1=", v1) // 输出: v1= 100
    fmt.Println("v2=", v2) // 输出: v2= 200
}

不过有点坑的是valueCtx 不支持链式调用

Context：go routine 控制

context 实现go routing 控制主要有两种，关闭控制和超时控制，本节讲下关闭控制，下一节讲下超时控制 context 实现关闭(cancel)控制的实现类叫cancelCtx, 创建方法是

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

...

type cancelCtx struct {
    Context

    mu       sync.Mutex            // protects following fields
    done     chan struct{}         // created lazily, closed by first cancel call
    children map[canceler]struct{} // set to nil by the first cancel call
    err      error                 // set to non-nil by the first cancel call
}

func (c *cancelCtx) Value(key interface{}) interface{} {
    if key == &cancelCtxKey {
        return c
    }
    return c.Context.Value(key)
}

func (c *cancelCtx) Done() <-chan struct{} {
    c.mu.Lock()
    if c.done == nil {
        c.done = make(chan struct{})
    }
    d := c.done
    c.mu.Unlock()
    return d
}

func (c *cancelCtx) Err() error {
    c.mu.Lock()
    err := c.err
    c.mu.Unlock()
    return err
}

type stringer interface {
    String() string
}

func contextName(c Context) string {
    if s, ok := c.(stringer); ok {
        return s.String()
    }
    return reflectlite.TypeOf(c).String()
}

func (c *cancelCtx) String() string {
    return contextName(c.Context) + ".WithCancel"
}

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    if c.err != nil {
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err
    if c.done == nil {
        c.done = closedchan
    } else {
        close(c.done)
    }
    for child := range c.children {
        // NOTE: acquiring the child's lock while holding parent's lock.
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {
        removeChild(c.Context, c)
    }
}

cancelCtx 的源码也不是很难读，主要是propagateCancel代码有点绕，我会在源码分析章进行解读，现在先不用关注。这里我们主要注意几点

  1. context 之间存在父子关系，
  2. WithCancel(parent Context) 返回子context，父context 取消的时候会同时取消子context, 这也是go routine 关闭控制的关键

好了，我们看下demo吧

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    // Pass a context with a timeout to tell a blocking function that it
    // should abandon its work after the timeout elapses.
    parentCtx, parentCancel := context.WithCancel(context.Background())
    childCtx, childCancel := context.WithCancel(parentCtx)
    
    go parentJob(parentCtx, parentCancel)

    go childJob(childCtx, childCancel)

    time.Sleep(time.Duration(5)*time.Second)
    
}

func parentJob(ctx context.Context, cancel context.CancelFunc) {
    // do someting
    time.Sleep(time.Duration(2)*time.Second)

    fmt.Println("parent job cancled.")
    cancel()
}

func childJob(ctx context.Context, cancel context.CancelFunc) {
    // do someting
    for {
        select {
        case <- ctx.Done():
            fmt.Println("child job cancled because of parent job cancel.")
            return
        default :
            // do some thing, or cancel child job
        }
    }
}

如果大家希望深入理解，最好去看一下包，比如net, database/sql 等的源码，加深理解

Context：超时控制

context 实现超时控制, 先放下源码

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c)
    dur := time.Until(d)
    if dur <= 0 {
        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 {
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}


// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
// implement Done and Err. It implements cancel by stopping its timer then
// delegating to cancelCtx.cancel.
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.


    deadline time.Time
}


func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}


func (c *timerCtx) String() string {
    return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
        c.deadline.String() + " [" +
        time.Until(c.deadline).String() + "])"
}


func (c *timerCtx) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err)
    if removeFromParent {
        // Remove this timerCtx from its parent cancelCtx's children.
        removeChild(c.cancelCtx.Context, c)
    }
    c.mu.Lock()
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}


// WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete:
//
//     func slowOperationWithTimeout(ctx context.Context) (Result, error) {
//         ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
//         defer cancel()  // releases resources if slowOperation completes before timeout elapses
//         return slowOperation(ctx)
//     }
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

看下如何实现超时控制的demo

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    // Pass a context with a timeout to tell a blocking function that it
    // should abandon its work after the timeout elapses.
    parentCtx, parentCancel := context.WithTimeout(context.Background(), 3*time.Second)
    childCtx, childCancel := context.WithCancel(parentCtx)
    
    go parentJob(parentCtx, parentCancel)

    go childJob(childCtx, childCancel)

    time.Sleep(time.Duration(5)*time.Second)
    
}

func parentJob(ctx context.Context, cancel context.CancelFunc) {
    // do someting
    time.Sleep(time.Duration(2)*time.Second)

    fmt.Println("parent job cancled.")
    cancel()
}

func childJob(ctx context.Context, cancel context.CancelFunc) {
    // do someting
    for {
        select {
        case <- ctx.Done():
            fmt.Println("child job cancled. reason = ", ctx.Err().Error())
            return
        default :
            // do some thing, or cancel child job
        }
    }
}

上面的demo，因为parentJob超时时间是3s, childjob 在3s会被取消而终止任务

Context 包源码分析

好了，下面我们看下context包的源码吧，上面源码基本放的差不多了，为了方便大家理解，画了下context 类图 context 类图 上面valueCtx 的代码和CancelCtx、timerCtx 代码已经放了，也不是很难懂，下面主要看下比较复杂点的函数

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    if done == nil {
        return // parent is never canceled
    }

    // 这块代码不是必须的，但是能看出功力，因为如果父context已经取消，就直接返回了，不用加锁了
    select {
    case <-done:
        // parent is already canceled
        child.cancel(false, parent.Err())
        return
    default:
    }


    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // parent has already been canceled
            // 这个必须有，因为在多协程环境，parentCancelCtx返回正常，到p.mu.Lock 之间，context 可能会取消
            child.cancel(false, p.err)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        
        atomic.AddInt32(&goroutines, +1)
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

好了，下次分享time包吧

golang社区

知识星球，一起golang: t.zsxq.com/nUR723R 博客: blog.17ledong.com/

avatar
文章
阅读
粉丝