Timeout using context package in Go

#go

Key takeaways

context.WithTimeout can be used in a timeout implementation.
WithDeadline returns CancelFunc that tells an operation to abandon its work.
timerCtx implements cancel() by stopping its timer then delegating to cancelCtx.cancel, and cancelCtx closes the context.
ctx.Done returns a channel that's closed when work done on behalf of this context should be canceled.

context.WithTimeout

The context package as the standard library was moved from the golang.org/x/net/context package in Go 1.7. This allows the use of contexts for cancellation, timeouts, and passing request-scoped data in other library packages.

context.WithTimeout can be used in a timeout implementation.



func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

For example, you could implement it as follow (Go playground):



package main

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

func execute(ctx context.Context) error {
    proc1 := make(chan struct{}, 1)
    proc2 := make(chan struct{}, 1)

    go func() {
        // Would be done before timeout
        time.Sleep(1 * time.Second)
        proc1 <- struct{}{}
    }()

    go func() {
        // Would not be executed because timeout comes first
        time.Sleep(3 * time.Second)
        proc2 <- struct{}{}
    }()

    for i := 0; i < 3; i++ {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-proc1:
            fmt.Println("process 1 done")
        case <-proc2:
            fmt.Println("process 2 done")

        }
    }

    return nil
}

func main() {
    ctx := context.Background()
    ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
    defer cancel()

    if err := execute(ctx); err != nil {
        log.Fatalf("error: %#v\n", err)
    }
    log.Println("Success to process in time")
}

Canceling this context releases resources associated with it, so you should call cancel as soon as the operations running in this Context complete.



ctx := context.Background()
ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()

Cancel notification after timeout is received from ctx.Done(). Done returns a channel that's closed when work done on behalf of this context should be canceled. WithTimeout arranges for Done to be closed when the timeout elapses.



select {
case <-ctx.Done():
    return ctx.Err()
}

When you execute this code, you will get the following result. A function call that can be completed in 1s will be finished, but a function call that can be done after 3s will not be executed because a timeout occurs in 2s.



$ go run main.go
process 1 done
2021/12/28 12:32:59 error: context.deadlineExceededError{}
exit status 1

In this way, you can implement timeout easily.

Deep dive into context.WithTimeout

Here's a quick overview.

WithTimeout is a wrapper function for WithDeadline.



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

WithDeadline returns CancelFunc that tells an operation to abandon its work. Internally, a function that calls timerCtx.cancel(), a function that's not exported, will be returned.



func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    // (omit)
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    // (omit)
    return c, func() { c.cancel(true, Canceled) }
}

// (omit)

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

    deadline time.Time
}

A timerCtx carries a timer and a deadline, and embeds a cancelCtx.



type cancelCtx struct {
    Context

    mu       sync.Mutex            // protects following fields
    done     atomic.Value          // of 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
}

timerCtx implements cancel() by stopping its timer then delegating to cancelCtx.cancel.



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
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        c.done.Store(closedchan)
    } else {
        close(d)
    }
    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)
    }
}

In the function the context is closed.

Conclusion

I explained how to implement timeout with context package, and dived into internal implementation in it. I hope this helps you understand the Go implementation.

Top comments (3)

Sebastian Bogan • Jan 1 '22 • Edited

I may be worthwhile mentioning that the second goroutine is executed (as opposed to the comments). Only execute() is gone and thus the print is not happening.

If the second goroutine would be an "endless" loop without checking ctx.Done(), it would be leaked.