context: add AfterFunc (original) (raw)

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
}