DEV Community

Kazuki Higashiguchi
Kazuki Higashiguchi

Posted on

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)
Enter fullscreen mode Exit fullscreen mode

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")
}
Enter fullscreen mode Exit fullscreen mode

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()
Enter fullscreen mode Exit fullscreen mode

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()
}
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

In this way, you can implement timeout easily.

Deep dive into context.WithTimeout

Here's a quick overview.

A diagram describing the inside of context package

WithTimeout is a wrapper function for WithDeadline.

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

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
}
Enter fullscreen mode Exit fullscreen mode

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
}
Enter fullscreen mode Exit fullscreen mode

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)
    }
}
Enter fullscreen mode Exit fullscreen mode

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.

Discussion (3)

Collapse
sebogh profile image
Sebastian Bogan • Edited on

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.

Collapse
tfutada profile image
Takashi Futada

Is it a forgotten sender?

Collapse
hgsgtk profile image
Kazuki Higashiguchi Author

Thank you so much for your point out! I'm a little busy with work right now, so I'll watch it this weekend.