tsungjen.io

Intro

There are two classic methods in Go to handle concurrency. One is WaitGroup, the another one is Context. In this article, we are going to talk about Context and explore the source code to see how it works behind the scenes. After the source code review, we will gain more insight of how to use Context properly.

source code: src/context/context.go

Interface

First, Take a look at the interface


type Context interface {

	Deadline() (deadline time.Time, ok bool)

	Done() <-chan struct{}

	Err() error

    Value(key interface{}) interface{}
    
}

There are 4 methods for this interface.

Deadline: Deadline is a method to get the deadline of the context. When deadline is due, the context will raise a cancellation automatically. The second return value ok is to indicate whether the context has a deadline or not. If ok equals false, there is no deadline for the context
Done: Done is a method to get a read-only channel. We can read the channel and wait for its parent context to raise the cancellation
Err: Err returns the reason to cancel
Value: To get the value binded to the context. (We can bind the value to the context with a key)

There is a basic implimentation, emptyCtx, for Context and 2 objects are initiated to use, Background and TODO.


// An emptyCtx is never canceled, has no values, and has no deadline. It is not
// struct{}, since vars of this type must have distinct addresses.
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
}

func (e *emptyCtx) String() string {
	switch e {
	case background:
		return "context.Background"
	case todo:
		return "context.TODO"
	}
	return "unknown empty Context"
}

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

// Background returns a non-nil, empty Context. It is never canceled, has no
// values, and has no deadline. It is typically used by the main function,
// initialization, and tests, and as the top-level Context for incoming
// requests.
func Background() Context {
	return background
}

// TODO returns a non-nil, empty Context. Code should use context.TODO when
// it's unclear which Context to use or it is not yet available (because the
// surrounding function has not yet been extended to accept a Context
// parameter).
func TODO() Context {
	return todo
}

We can see emptyCtx do nothing. All the methods simply return nil. We can not set the deadline, cancel or bind the value to the Context.
TODO is temporary used when we are not sure which Context implimentation to use.
Background is a used as a root Context and to create the child Context with following 4 functions.


func WithCancel(parent Context) (ctx Context, cancel CancelFunc) 
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithValue(parent Context, key, val interface{}) Context

Let’s go through the functions one by one.

WithCacnel

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

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

// A cancelCtx can be canceled. When canceled, it also cancels any children
// that implement canceler.
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
}

WithCancel will receive a parent context and return a cancelCtx object along with a cancel function. In the function we will

Create a cancelCtx
Call propagateCancel to arrange this context to be canceled when parent is canceled.

A cancelCtx has a

Context to point to its parent
mutex to ensure thread-safe property (lock the context when we r/w the context),
done to push the cancellation signal
children to store its future children (any context can be a parent context)
err to store the reason for cancellation.

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
}

The implementation is quite simple and the lock makes it thread-safe

After a look to cancelCtx, Take a look at the propagateCancel function.

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

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

What propagateCancel does is to

Check parent is already canceled (parent can be canceled, meanwhile, we call WithCancel function). If its parent is canceled, cancel itself as well.
If parent is not canceled. Check if parent is cancelCtx, if yes add itself as its parent’s child. Otherwise, start a go routine to listen to a done signal and see if its parent is done (the goroutine will trigger the cacnellation when parent is done).

WithDeadline

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

WithDeadline returns a copy of the parent context with the deadline adjusted to be no later than d. If the parent’s deadline is already earlier than d, WithDeadline(parent, d) is semantically equivalent to parent. The returned context’s Done channel is closed when the deadline expires, when the returned cancel function is called, or when the parent context’s Done channel is closed, whichever happens first.

Take a look at timerCtx

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

timerCtx is built on cancelCtx, it has

cancelCtx created with its parent. It takes care of the cancellation process
deadline to store a deadline, we can get it by calling Deadline method

The cancellation is quite the same as cancelCtx except it needs to stop the timer when cancellation is triggered

WithTimeout

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

WithTimeout is a the same as WithDeadline. It takes a time.Duration as an input, the deadline will be now + duration.

WithValue

// WithValue returns a copy of parent in which the value associated with key is
// val.
//
// Use context Values only for request-scoped data that transits processes and
// APIs, not for passing optional parameters to functions.
//
// The provided key must be comparable and should not be of type
// string or any other built-in type to avoid collisions between
// packages using context. Users of WithValue should define their own
// types for keys. To avoid allocating when assigning to an
// interface{}, context keys often have concrete type
// struct{}. Alternatively, exported context key variables' static
// type should be a pointer or 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}
}

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

// stringify tries a bit to stringify v, without using fmt, since we don't
// want context depending on the unicode tables. This is only used by
// *valueCtx.String().
func stringify(v interface{}) string {
	switch s := v.(type) {
	case stringer:
		return s.String()
	case string:
		return s
	}
	return "<not Stringer>"
}

func (c *valueCtx) String() string {
	return contextName(c.Context) + ".WithValue(type " +
		reflectlite.TypeOf(c.key).String() +
		", val " + stringify(c.val) + ")"
}

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

WithValue returns a copy of parent in which the value associated with key is val

valueCtx has a

context to point to its parent
key, value to store a k-v pair

The lookup for the value is a recursive process. A goroutine will recursively find the key, if the key is not existed in all the nodes, the root (Background) will return a nil (emptyCtx implementation).

Noted that valueCtx doesn’t ensure the val is thread-safe.

Conclusion

After we explore the source code of Context, we have conclusion as belows

Do not put Context inside the structure. Pass it as a parameter
Make the Context the first paramter in a function
Do not pass nil when you passing Context. It will cause panic if someone uses nil to create child Context. Pass context.TODO if you dont know what to pass.
Use valueCtx properly, pass the required data only (such as request meta data). valueCtx doesn’t ensure thread-safe for value, passing unnecessary data may cause dirty data issues.
Context itself is thread-safe, feel free pass it to different goroutines.

Go Context Source Code review