This commit adds and passes context around to (almost) every function. This allows cancelling (almost) every function call globally.
56 lines
1.4 KiB
Go
56 lines
1.4 KiB
Go
package cancel
|
|
|
|
import (
|
|
"context"
|
|
)
|
|
|
|
// WithContext executes f and returns the returns the return value and nil.
|
|
//
|
|
// If the context is closed before f returns, invokes cancel and returns f(), ctx.Err().
|
|
//
|
|
// In general, WithContext always waits for f() to return even if cancel was called.
|
|
// As a special case if a closed context is passed, f is not invoked.
|
|
//
|
|
// allowcancel must be called by f exactly once, as soon as the cancel function may be invoked.
|
|
func WithContext[T any](ctx context.Context, f func(allowcancel func()) T, cancel func()) (t T, err error) {
|
|
t, _, err = WithContext2(ctx, func(start func()) (T, struct{}) {
|
|
return f(start), struct{}{}
|
|
}, cancel)
|
|
return
|
|
}
|
|
|
|
// WithContext2 is exactly like WithContext, but takes a function returning two parameters.
|
|
func WithContext2[T1, T2 any](ctx context.Context, f func(start func()) (T1, T2), cancel func()) (t1 T1, t2 T2, err error) {
|
|
// context is already closed, don't even try invoking it.
|
|
if err := ctx.Err(); err != nil {
|
|
return t1, t2, err
|
|
}
|
|
|
|
cancelable := make(chan struct{}, 1)
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
defer close(cancelable)
|
|
|
|
t1, t2 = f(func() {
|
|
cancelable <- struct{}{}
|
|
})
|
|
}()
|
|
|
|
select {
|
|
case <-done:
|
|
// the function has exited regularly
|
|
// nothing to be done
|
|
case <-ctx.Done():
|
|
|
|
// context was cancelled
|
|
<-cancelable
|
|
cancel()
|
|
|
|
// still wait for it to be done!
|
|
<-done
|
|
err = ctx.Err()
|
|
}
|
|
return
|
|
}
|