diff --git a/fn/context_guard.go b/fn/context_guard.go new file mode 100644 index 000000000..23914702e --- /dev/null +++ b/fn/context_guard.go @@ -0,0 +1,113 @@ +package fn + +import ( + "context" + "sync" + "time" +) + +var ( + // DefaultTimeout is the default timeout used for context operations. + DefaultTimeout = 30 * time.Second +) + +// ContextGuard is an embeddable struct that provides a wait group and main quit +// channel that can be used to create guarded contexts. +type ContextGuard struct { + DefaultTimeout time.Duration + Wg sync.WaitGroup + Quit chan struct{} +} + +func NewContextGuard() *ContextGuard { + return &ContextGuard{ + DefaultTimeout: DefaultTimeout, + Quit: make(chan struct{}), + } +} + +// WithCtxQuit is used to create a cancellable context that will be cancelled +// if the main quit signal is triggered or after the default timeout occurred. +func (g *ContextGuard) WithCtxQuit() (context.Context, func()) { + return g.WithCtxQuitCustomTimeout(g.DefaultTimeout) +} + +// WithCtxQuitCustomTimeout is used to create a cancellable context that will be +// cancelled if the main quit signal is triggered or after the given timeout +// occurred. +func (g *ContextGuard) WithCtxQuitCustomTimeout( + timeout time.Duration) (context.Context, func()) { + + timeoutTimer := time.NewTimer(timeout) + ctx, cancel := context.WithCancel(context.Background()) + + g.Wg.Add(1) + go func() { + defer timeoutTimer.Stop() + defer cancel() + defer g.Wg.Done() + + select { + case <-g.Quit: + + case <-timeoutTimer.C: + + case <-ctx.Done(): + } + }() + + return ctx, cancel +} + +// CtxBlocking is used to create a cancellable context that will NOT be +// cancelled if the main quit signal is triggered, to block shutdown of +// important tasks. The context will be cancelled if the timeout is reached. +func (g *ContextGuard) CtxBlocking() (context.Context, func()) { + return g.CtxBlockingCustomTimeout(g.DefaultTimeout) +} + +// CtxBlockingCustomTimeout is used to create a cancellable context with a +// custom timeout that will NOT be cancelled if the main quit signal is +// triggered, to block shutdown of important tasks. The context will be +// cancelled if the timeout is reached. +func (g *ContextGuard) CtxBlockingCustomTimeout( + timeout time.Duration) (context.Context, func()) { + + timeoutTimer := time.NewTimer(timeout) + ctx, cancel := context.WithCancel(context.Background()) + + g.Wg.Add(1) + go func() { + defer timeoutTimer.Stop() + defer cancel() + defer g.Wg.Done() + + select { + case <-timeoutTimer.C: + + case <-ctx.Done(): + } + }() + + return ctx, cancel +} + +// WithCtxQuitNoTimeout is used to create a cancellable context that will be +// cancelled if the main quit signal is triggered. +func (g *ContextGuard) WithCtxQuitNoTimeout() (context.Context, func()) { + ctx, cancel := context.WithCancel(context.Background()) + + g.Wg.Add(1) + go func() { + defer cancel() + defer g.Wg.Done() + + select { + case <-g.Quit: + + case <-ctx.Done(): + } + }() + + return ctx, cancel +}