proposal: context: WithInterrupt for lower overhead checking of cancellation

[duckbrain]

Proposal Details

Abstract

This proposal introduces WithInterrupt, a new function in the context package that provides an atomic-based, low-overhead cancellation check for performance-sensitive code.

Background

The context package provides mechanisms for request-scoped cancellation, by checking ctx.Err() or <-ctx.Done(), but in tight loops these can introduce significant overhead.

A common workaround is to periodically check ctx.Err() using a counter rather than in every iteration, but this introduces a trade-off between responsiveness and efficiency. A better approach is to use sync/atomic to maintain a separate cancellation flag, but this both a non-obvious solution and must be implemented manually each time, though this code isn't that difficult to write once you understand the problem.

Proposal

Add WithInterrupt to the context package (implementation as an example)

// WithInterrupt returns a dirived context that pionts to the parent but has a new Done channel in a similar fasion to [WithCancel].
// The returned isDone function can be called to efficiently check the context's cancelation status.
// This is useful for tight loops that need an interrupt where checking `ctx.Err()` or `<-ctx.Done()` adds too much overhead.
//
// Like [WithCancel], canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this [Context] complete.
func WithInterrupt(parent Context) (ctx Context, isDone func() bool, cancel func()) {
	var flag uint32
	ctx, cancel = WithCancel(parent)

	go func() {
		<-ctx.Done()
		atomic.StoreUint32(&flag, 1)
	}()

	isDone = func() bool {
		return atomic.LoadUint32(&flag) != 0
	}
	return
}

Use Case Example

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

for {
	if isDone() {
		return ctx.Err()
	}
	// Perform computation
}

[seankhliao]

It seems you can do this with just

var b atomic.Bool
context.AfterFunc(ctx, func() { b.Store(true) })

for {
  if b.Load() {
    _ = ctx.Err()
  }
}

That is, does this really need to be in the standard library? https://go.dev/doc/faq#x_in_std

[adonovan]

If this is a performance hack, could you provide a comparison of timings, ideally a CPU profile?

I wonder why the ordinary channel poll operation is too slow, and whether it could simply be made faster. Aselect { case <- ch: default: } statement boils down to a call to runtime.selectnbrecv, which calls chanrecv(false), which does the following checks:

• that the channel is non-nil
• that we're not in a synctest bubble
• that the channel has no timer
• that the receive is nonblocking
• that the channel element type is struct{}

and then does an atomic load.

The last two could in principle be eliminated statically by the compiler. The rest seem truly dynamic, albeit fairly predictable branches.

[gabyhelp]

Related Issues

• context: add AfterFunc #57928 (closed)
• proposal: add context cancellation mechanisms to os/signal #16472 (closed)
• proposal: context: add WithCancelReason #46273 (closed)
• proposal: context.WithSignal #21521 (closed)
• proposal: context: enable first class citizenship of third party implementations #28728 (closed)
• context: add WithoutCancel #40221 (closed)
• proposal: context: new method WithCallbackOnCancel(...) #57325 (closed)
• proposal: context: add Scoped #44751 (closed)

Related Documentation

• Package context > func WithCancel ¶ 1.7
• Package context > func WithTimeout ¶ 1.7

_{(Emoji vote if this was helpful or unhelpful; more detailed feedback welcome in this discussion.)}

[neild]

As @seankhliao says, it seems that you can do exactly this today with context.AfterFunc.

Also, ctx.Err() consists of a single mutex lock/unlock. It's pretty fast, but checking an atomic bool is admittedly faster. If we want to make ctx.Err faster, perhaps we could change it to use an atomic operation rather than a mutex. (Either change the internal cancelCtx.err value to an atomic.Value, or add a separate atomic.Bool tracking when the cancelCtx is canceled.)

[neild]

The last two could in principle be eliminated statically by the compiler. The rest seem truly dynamic, albeit fairly predictable branches.

FWIW, on my M1 MacBook an atomic.Value.Load is ~2ns and a select { case <-ch: default: } is ~4ns. The current Context.Err implementation (one mutex lock/unlock) is ~13ns.

Selecting on context.Done is ~6ns, because context.Done itself includes an atomic load of the lazily-created done channel.

So Err is pretty fast as it is, but we could make it quite a bit faster by switching to atomics.

[gopherbot]

Change https://go.dev/cl/653795 mentions this issue: context: use atomic operation in ctx.Err