Golang Context

I’m learning golang and I have reached chapter covering contexts. I initially understood the idea behind contexts but to further widen my understanding I’m gonna cover custom implementation here.

Context is an interface

Golang context is an interface with following methods:

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

Why is it required

Go has in it’s core the idea of goroutine, a lightweight user-space managed thread that is cheap and efficient. Goroutines are usually communicating with each other using go channels.

While goroutine cover launching and execution of our code, what it does not handle is cancellation.

func main() {
	var ch chan int

	longRunningFunc := func(ch chan int) {
		time.Sleep(time.Second * 5)
		ch <- 1
	}

	go longRunningFunc(ch)

	select {
	case i := <-ch:
		fmt.Println(i)
	}

}

Custom cancellation interface

For now we will pretend that there is no Context interface, and provide our own implementation called MyContext

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

The idea is that across all our code-base we use this interface to handle following goroutine aspects:

• Identifying when we must forcefully cancel the goroutine execution
• Identifying until when we can run the goroutine (deadline method)
• Capturing the error (err method)
• Capturing context variable (Value method)

MyContext implementation

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

type MyContextImpl struct {
	deadline time.Time
	done     chan struct{}
	ok       bool
	err      error
	kv       map[interface{}]interface{}
}

func (ctx MyContextImpl) Deadline() (time.Time, bool) {
	return ctx.deadline, ctx.ok
}

func (ctx MyContextImpl) Done() <-chan struct{} {
	return ctx.done
}

func (ctx MyContextImpl) Err() error {
	return ctx.err
}

func (ctx MyContextImpl) Value(key interface{}) interface{} {
	return ctx.kv[key]
}

func Background() MyContext {
	onClose := func(m *MyContextImpl) {
		select {
		case <-m.done:
			m.err = errors.New("Done called")
		}
	}
	ctx := MyContextImpl{
		ok: false,
	}
	ctx.done = make(chan struct{})
	ctx.kv = make(map[interface{}]interface{})
	go onClose(&ctx)
	return ctx
}

func WithCancel(c MyContext) (MyContext, func()) {
	ctx := c.(MyContextImpl)
	n := &MyContextImpl{
		deadline: ctx.deadline,
		done:     ctx.done,
		ok:       ctx.ok,
		err:      ctx.err,
		kv:       ctx.kv,
	}

	return n, func() {
		n.err = errors.New("Canceled")
		close(n.done)
	}
}

func WithTimeout(myContext MyContext, timeout time.Duration) (MyContext, func()) {
	c := myContext.(MyContextImpl)
	n := &MyContextImpl{
		deadline: time.Now().Add(timeout),
		done:     c.done,
		ok:       true,
		err:      c.err,
		kv:       c.kv,
	}

	go func() {
		time.Sleep(timeout)
		n.err = errors.New("Timed out")
		close(n.done)
	}()

	return n, func() {
		n.err = errors.New("Canceled")
	}
}

Using our MyContext within our code

Now we have updated our code to use MyContext. We can see that we have only used ctx.Done() channel, to use MyContext. All other details behind MyContext are hidden in implementation.

func main() {
	ch := make(chan int)
	ctx := Background()

	longRunningFunc := func(ch chan int, ctx MyContext) {
		select {
		case <-time.After(time.Second * 5):
			ch <- 1
		case <-ctx.Done():
			fmt.Println("Canceled")
		}
	}

	go longRunningFunc(ch, ctx)

	select {
	case i := <-ch:
		fmt.Println(i)
	}
}

Canceling goroutine

There are times, when we run goroutine, but want to cancel it’s execution based on the caller’s condition.

func main() {
	ch := make(chan int)
	ctx, cancel := WithCancel(Background())

	longRunningFunc := func(ch chan int, ctx MyContext) {
		select {
		case <-time.After(time.Second * 5):
			ch <- 1
		case <-ctx.Done():
			fmt.Println("Canceled")
		}
	}

	go longRunningFunc(ch, ctx)
	cancel() //right away

	select {
	case i := <-ch:
		fmt.Println(i)
	case <-ctx.Done():
		
	}

	if e := ctx.Err(); e != nil {
		fmt.Println(e)
	}
}

Setting timeout for our goroutine

func main() {
	ch := make(chan int)
	ctx, _ := WithTimeout(Background(), time.Second*2)

	longRunningFunc := func(ch chan int, ctx MyContext) {
		select {
		case <-time.After(time.Second * 5):
			ch <- 1
		case <-ctx.Done():
			fmt.Println("Canceled")
		}
	}

	go longRunningFunc(ch, ctx)

	select {
	case i := <-ch:
		fmt.Println(i)
	case <-ctx.Done():
	}

	if e := ctx.Err(); e != nil {
		fmt.Println(e)
	}
}