go context源码分析

 

 

源码面前，了无秘密。最近由于工作需要使用gprc，深入了解了下context这个包。context是go语言内置包，goroutine提供上下文，创建者可以比较方便的通知所以线程结束。（也仅仅是比较方便而已）

 

廉价的代价

import threading, time, inspect, ctypes

def _async_raise(tid, exctype):
    tid = ctypes.c_long(tid)
    if not inspect.isclass(exctype):
        exctype = type(exctype)
    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
    if res == 0:
        raise ValueError("invalid thread id")
    elif res != 1:
        ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)

def happy():
    while True:
        print("being happy!!!")
        time.sleep(1)

def main():
    fucker = threading.Thread(target=happy)
    fucker.start()
    time.sleep(3)
    _async_raise(fucker.ident, SystemExit)
    print("I don't care if you are being happy")
    fucker.join()

if __name__=="__main__":
    main()

我们可以通过这种粗鲁的方式结束子进程(很多人喜欢粗鲁...)当然，这种结束方式并不优雅，但是其中隐含了一种潜在的权力，就是我可以通过某种方式强行结束子线程(无论子线程是否同意)，这来源于操作系统提供的接口，大多数编程语言的线程对应了操作系统的线程，IEEEIEEE标准1003.1c中定义了线程的标准，pthread_exit调用用于结束线程。大部分类unix操作系统支持该标准(windows也部分支持)。但是go语言的线程和这些完全不同，go语言实现了非常廉价的线程。创建、销毁都非常轻量级，当然他们完全是用户级线程，go语言没有这样的语法获取线程句柄 goroutine:=go func(){}(),主线程无法主动了解线程的运行情况，例如线程是否正在运行还是已经运行结束，当然也无法主动结束子线程的运行。所以唯一可以选择的方式便是通过某种方式相互沟通，运行结束同时主线程已经运行结束，主线程通知子线程退出运行等等。。。

 

Context接口

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

如果设置了deadline,Deadline函数返回dealine,ok为true，否则ok为false,

Done函数返回一个channel，当deadline超时、主动执行了cancel函数时channel关闭

Err()当channel未关闭时，返回nil；当channel关闭时，返回错误

Value()获取上下文参数

上层通过context向线程传递上下文信息，线程需要通过Done方法主动判断是否应该结束线程

emptyCtx类型

type emptyCtx int

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

func (*emptyCtx) Done() <-chan struct{} {
	return nil
}

func (*emptyCtx) Err() error {
	return nil
}

func (*emptyCtx) Value(key interface{}) interface{} {
	return nil
}

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

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

实现了Context接口定义的方法，只返回空值，context包定义了两个emptyCtx类型的包内变量，background,todo。可以通过Background()和TODO()函数返回这两个变量。background用于作为所有context的顶层context,todo用于传递一个默认的context。

cancelCtx类型

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
}

 ancelCtx类型嵌入了Context接口，ancelCtx类型定义并覆盖了Context的Done(),Err()方法，同时还定义了cancel方法。

Done()方法返回cancelCtx类型的done成员变量。由于符合Context接口的Done只有在这里重新定义，所以调用Done()，不是调用了emptyCtx类型的Done()方法就是调用了cancelCtx类型的Done方法。

其中children成员变量保存了所有以此cancelCtx为祖先节点的距离最近的cancelCtx类型和timerCtx类型。当调用cancel方法时，cancelCtx同时调用children保存的所有canceler接口的cancel方法。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}

WithCancel函数创建一个Context，同时返回一个内部调用cancel方法的CancelFunc函数。

newCancelCtx创建一个以Context成员变量为parent的cancelCtx类型。

propagateCancel函数将cancelCtx类型加入到距离自己最近的cancelCtx类型的祖先节点（或者timerCtx类型,timerCtx类型嵌入了cancelCtx类型,类似于继承了cancelCtx）

func propagateCancel(parent Context, child canceler) {
	if parent.Done() == nil {
		return // parent is never canceled
	}
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err)
		} else {
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

propagateCancel 用于将子节点加入到父节点的children成员变量中，如果父节点继承或定义的Done()方法返回nil,例如继承emptyCtx类型或者自定义的不可取消的Context接口类型，直接返回。因为父节点不可以取消，所以没有必要将child节点保存到相应的祖先节点里。

parentCancelCtx查询child的祖先节点中距离最近的CancelCtx节点。如果CancelCtx节点已经被取消了，调用child节点的cancel方法，如果没有被取消，将自节点加入到children成员变量里。

关于propagateCancel函数最后开启的线程，用于用户自定义的Context类型，当用户自定义了一个可取消的Context类型，为了保证当自定义的Context被取消之后，以自定义类型为父节点的cancelCtx和timerCtx类型可以被调用cancel方法。case<-child.Done()代码表示如果子节点已经取消，无论父节点有没有被取消，这个线程都可以结束了。如果不使用自定义的Context类型，这个线程是毫无必要的

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
		}
	}
}

parentCancelCtx函数，如果parent是*cancelCtx，则返回，如果不是，则寻找嵌入在parent内的*cancelCtx类型,如果类型无法识别（自定以的Context类型或者emptyCtx）返回nil,false

 

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)
	}
}

当用户调用返回的CancelFunc 时，实际调用了相应cancelCtx或者timerCtx的cancel方法。必须传入一个非nil的err，获取mu锁后，关闭成员变量done，调用children保存的节点的cancel方法，会从上而下的获取mu锁，所以不会导致死锁。调用完成后设置children=nil,释放mu锁。如果removeFromParent=true,从祖先节点的children中移除。removeChild会申请祖先节点的mu锁，所以必须先释放当前mu锁，否则可能会导致死锁。

type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

timerCtx嵌入了cancelCtx类型，继承了cancelCtx的Done方法，覆盖了cancelCtx的cancel方法，定义了Deadline方法。

Deadline方法返回成员变量deadline。

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	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) }
}

WithDeadline函数设置超时时间，返回Context和CancelFunc

当祖先节点的超时时间在当前之前，则不需要设置deadline,直接调用WithCancel函数创建一个cancelCtx节点。

propagateCancel函数用于将timerCtx节点保存到祖先节点的children成员变量里，当祖先节点cancel后，调用保存在children的子孙节点cancel方法。

如果节点已经超时，取消当前节点，否则，申请mu锁，调用time.AfterFunc延后执行cancel方法，返回Context和CancelFunc,CancelFunc调用timerCtx的cancel方法。

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()
}

timerCtx类型的cancel方法,首先调用cancelCtx的cancel方法（注意传入的第一个参数是false）,因为保存在父节点中children成员变量里的的canceler接口是timerCtx类型,如果传入c.cancelCtx.cancel方法的removeFromParent为true,则调用removeChild函数的第二个参数为c.cancelCtx，导致删除错误。

如果removeFromParent为true,调用removeChild函数删除祖先节点保存的当前节点。

停止延后执行的函数，设置timer为nil

type valueCtx struct {
	Context
	key, val interface{}
}

valueCtx类型嵌入了Context接口，同时有key,val两个interface{}成员变量。定义了Value方法

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

WithValue返回保存了参数的上下文，key必须课比较的，返回一个保存了参数的上下文

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

因为Context有层次结构，下层节点保存了上层节点，所以可以查询以前添加的参数，直到调用了emptyCtx的Value方法。