context: add AfterFunc

[neild]

Edit: The latest version of this proposal is #57928 (comment).

---

This proposal originates in discussion on #36503.

Contexts carry a cancellation signal. (For simplicity, let us consider a context past its deadline to be cancelled.)

Using a context's cancellation signal to terminate a blocking call to an interruptible but context-unaware function is tricky and inefficient. For example, it is possible to interrupt a read or write on a net.Conn or a wait on a sync.Cond when a context is cancelled, but only by starting a goroutine to watch for cancellation and interrupt the blocking operation. While goroutines are reasonably efficient, starting one for every operation can be inefficient when operations are cheap.

I propose that we add the ability to register a function which is called when a context is cancelled.

package context

// OnDone arranges for f to be called in a new goroutine after ctx is cancelled.
// If ctx is already cancelled, f is called immediately.
// f is called at most once.
//
// Calling the returned CancelFunc waits until any in-progress call to f completes,
// and stops any future calls to f.
// After the CancelFunc returns, f has either been called once or will not be called.
//
// If ctx has a method OnDone(func()) CancelFunc, OnDone will call it.
func OnDone(ctx context.Context, f func()) CancelFunc

OnDone permits a user to efficiently take some action when a context is cancelled, without the need to start a new goroutine in the common case when operations complete without being cancelled.

OnDone makes it simple to implement the merged-cancel behavior proposed in #36503:

func WithFirstCancel(ctx1, ctx2 context.Context) (context.Context, context.CancelFunc) {
	ctx, cancel := context.WithCancel(ctx1)
	stopf := context.OnDone(ctx2, func() {
		cancel()
	})
	return ctx, func() {
		cancel()
		stopf()
	}
}

Or to stop waiting on a sync.Cond when a context is cancelled:

func Wait(ctx context.Context, cond *sync.Cond) error {
	stopf := context.OnDone(ctx, cond.Broadcast)
	defer stopf()
	cond.Wait()
	return ctx.Err()
}

The OnDone func is executed in a new goroutine rather than synchronously in the call to CancelFunc that cancels the context because context cancellation is not expected to be a blocking operation. This does require the creation of a goroutine, but only in the case where an operation is cancelled and only for a limited time.

The CancelFunc returned by OnDone both provides a mechanism for cleaning up resources consumed by OnDone, and a synchronization mechanism. (See the ContextReadOnDone example below.)

Third-party context implementations can provide an OnDone method to efficiently schedule OnDone funcs. This mechanism could be used by the context package itself to improve the efficiency of third-party contexts: Currently, context.WithCancel and context.WithDeadline start a new goroutine when passed a third-party context.

---

Two more examples; first, a context-cancelled call to net.Conn.Read using the APIs available today:

// ContextRead demonstrates bounding a read on a net.Conn with a context
// using the existing Done channel.
func ContextRead(ctx context.Context, conn net.Conn, b []byte) (n int, err error) {
	errc := make(chan error)
	donec := make(chan struct{})
        // This goroutine is created on every call to ContextRead, and runs for as long as the conn.Read call.
	go func() {
		select {
		case <-ctx.Done():
			conn.SetReadDeadline(time.Now())
			errc <- ctx.Err()
		case <-donec:
			close(errc)
		}
	}()
	n, err = conn.Read(b)
	close(donec)
	if ctxErr := <-errc; ctxErr != nil {
		conn.SetReadDeadline(time.Time{})
		err = ctxErr
	}
	return n, err
}

And with context.OnDone:

func ContextReadOnDone(ctx context.Context, conn net.Conn, b []byte) (n int, err error) {
	var ctxErr error
        // The OnDone func runs in a new goroutine, but only when the context expires while the conn.Read is in progress.
	stopf := context.OnDone(ctx, func() {
		conn.SetReadDeadline(time.Now())
		ctxErr = ctx.Err()
	})
	n, err = conn.Read(b)
	stopf()
        // The call to stopf() ensures the OnDone func is finished modifying ctxErr.
	if ctxErr != nil {
		conn.SetReadDeadline(time.Time{})
		err = ctxErr
	}
	return n, err
}

[gopherbot]

Change https://go.dev/cl/462855 mentions this issue: context: add OnDone

[neild]

It is worth noting that OnDone is a performance optimization. We can write it today in terms of the existing context package. The benefit of adding it to the package is that it permits us to take action on context cancellation without leaving a goroutine sitting around waiting for the cancel to happen.

func OnDone(ctx context.Context, f func()) context.CancelFunc {
	stopc := make(chan struct{}, 1)
	donec := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			f()
		case <-stopc:
		}
		close(donec)
	}()
	return func() {
		select {
		case stopc <- struct{}{}:
		default:
		}
		<-donec
	}
}

[ianlancetaylor]

The ContextReadOnce example is still a bit awkward, in the sense that there are a few things that have to be done exactly right to avoid any problems. Perhaps we can tighten up the idea, at the cost of another closure.

// Try executes fn while watching ctx.  If ctx is cancelled, Try calls cancel and returns ctx.Err().
// Otherwise, Try returns the result of fn.
// This is implemented using an internal implementation of OnDone as described above.
func Try(fn func() error, cancel func()) error

// Example of using Try.
func ContextReadOnDone(ctx context.Context, conn net.Conn, b []byte) (n int, err error) {
    err = context.Try(func() error {
       n, err = conn.Read(b)
    }, func() {
        conn.SetReadDeadline(time.Now())
    })
    return n, err
}

[neild]

Try does look quite elegant to use. I don't think it would work for the WithFirstCancel case, and in general Try can be implemented in terms of WithDone but not vice-versa, so if we were to have only one I'd vote for OnDone.

func Try(ctx context.Context, fn func() error, cancel func()) error {
    stopf := context.OnDone(ctx, cancel)
    err := fn()
    stopf()
    if ctx.Err() != nil {
        return ctx.Err()
    }
    return err
}

[bcmills]

The Try function looks a lot like the os/exec functionality added for #50436.

• Try is analogous to (*exec.Cmd).Wait.
• The fn argument is analogous to the subprocess itself.
• The cancel argument is analogous to the exec.Cmd.Cancel field.

[bcmills]

Based on that analogy, I would suggest a slightly different implementation of Try in terms of OnDone:

func Try(ctx context.Context, fn func() error, cancel func()) error {
	canceled := false
	stop := context.OnDone(ctx, func() {
		cancel()
		canceled = true
	})
	err := fn()
	stop()
	if canceled && err == nil {
		return ctx.Err()
	}
	return err
}

Notably: in case of cancellation I would prefer to still return the result of fn if it is non-nil, since it may contain more detail about the result, but I would return ctx.Err() if fn returns nil in case the result of fn is spurious.

[rsc]

This proposal has been added to the active column of the proposals project
and will now be reviewed at the weekly proposal review meetings.
— rsc for the proposal review group

[rsc]

This seems similar to having time.AfterFunc to avoid making a goroutine that calls time.Sleep and then the function. Here we avoid making a goroutine that receives from ctx.Done and then calls the function. Perhaps it should be context.After or AfterFunc?

context.After(ctx, func() { println("ctx is done!") })

reads nicely to me. It's nice that this would let people implement context.Merge themselves at no efficiency cost compared to the standard library.

We should specify that f is always run in a goroutine by itself, at least semantically, even in the case where

// If ctx is already cancelled, f is called immediately.

f shouldn't be called by context.After in that case either.