通过源码理解Golang中的上下文Context(二)

Golang进阶·技术 · 2023-06-23

在上一篇文章中,我分享了有关context包的第一部分:valueCtxcancelCtx,我们在这篇文章中继续探索更多内容。

WithTimeout 和 WithDeadline

我们还是先来一个例子:

package main

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

func main() {
    cancelCtx, cancel := context.WithTimeout(context.Background(), time.Second*3)
    defer cancel()
    go task(cancelCtx)
    time.Sleep(time.Second * 4)
}

func task(ctx context.Context) {
    i := 1
    for {
        select {
        case <-ctx.Done():
            fmt.Println(ctx.Err())
            return
        default:
            fmt.Println(i)
            time.Sleep(time.Second * 1)
            i++
        }
    }
}

在上一篇文章中我们已经知道 cancelCtx 的行为,因此理解 WithTimeout 的工作原理非常简单,它接受一个超时时间,在此之后完成的通道将被关闭并且上下文会被取消。并且还会返回一个取消函数cancel,如果需要在超时之前取消上下文,可以调用该函数cancel

WithDeadline的用法和WithTimeout非常相似,我们来看一下源代码:

type timerCtx struct {
    cancelCtx
    timer *time.Timer 
    deadline time.Time
}

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

由于WithTimeoutWithDeadline之间有很多共同点,因此它们共享相同类型的上下文:timerCtx,它嵌入cancelCtx并定义了另外两个属性:timerdeadline

我们看一下创建timerCtx时会发生什么:

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    // Get deadline time of parent context. 
    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 { // 'err' field of the embedded cancelCtx is promoted 
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

相比WithCancleWithValueWithDeadline要复杂一些,我们来一点点的过一遍。

首先,parent.Deadline 将获取父上下文的截止时间。 Deadline 方法在 Context 接口中定义如下:

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    ...
}

在context包中,只有emptyCtxtimerCtx类型实现了该方法:

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}

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

因此,当我们调用parent.Deadline()时,如果父上下文也是timerCtx类型,并且实现了自己的Deadline()方法,那么就可以获得父上下文的截止时间。否则,如果父上下文是cancelCtxvalueCtx类型,那么最后会调用emptyCtxDeadline()方法,我们会得到time.Timebool类型的零值(0001-01-01 00:00:00 +0000 UTC 和 false)。

如果parentdeadline早于传入的deadline参数,则直接调用WithCancel(parent)返回一个cancelCtx。当传入的deadline合理时,我们需要创建一个timerCtx

//inside WithDeadline() function
...
c := &timerCtx{
    cancelCtx: newCancelCtx(parent),
    deadline:  d,
}
propagateCancel(parent, c)
...

在上面的代码中,又看到了propagateCancel方法,我在上一篇文章中已经讨论过它,如果你不明白,请参考这里

cancelCtx类似,timerCtx通过调用自己的cancel方法关闭done通道来发送上下文取消信号。取消上下文有两种情况:

  1. timeout cancel:超过期限时,自动关闭done通道;

    // inside WithDeadline function
    ...
    // timeout cancel
    c.timer = time.AfterFunc(dur, func() {
    c.cancel(true, DeadlineExceeded)
    })
    ...
  2. 手动cancel:调用返回的cancel函数,在截止时间之前关闭done通道

    // inside WithDeadline function
    ...
    // return the cancel function as the second return value
    return c, func() { c.cancel(true, Canceled) }
    ...

两种情况都会调用cancel方法。

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err) // close the done channel and set err field
    if removeFromParent {
        // Remove this timerCtx from its parent cancelCtx's children.
        // Note: timerCtx c's parent is c.cancelCtx.Context
        removeChild(c.cancelCtx.Context, c)
    }
    c.mu.Lock()
    // stop and clean the timer
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

timerCtx 实现 cancel 方法来停止和重置计时器,然后委托给 cancelCtx.cancel

到这里我们就对context里面重要的一些源码分析结束。

Context Golang进阶 Golang源码
Theme Jasmine by Kent Liao