mirror of
https://github.com/lightningnetwork/lnd.git
synced 2024-11-19 01:43:16 +01:00
deefa3a9f1
This adds the ContextGuard struct from the taproot-assets repository that has been in use there for a while. The context guard allows for easy creation of contexts that depend on a central wait group and quit channel. A context can either be created with a timeout and quit, meaning it will cancel either on reaching the timeout or when the central quit channel is closed. Or a context can be created to block and use a timeout, meaning it will _not_ cancel on quit but rather block the shutdown until it is completed (or times out). The third way is to create a context that just cancels on quit with no timeout.
114 lines
2.7 KiB
Go
114 lines
2.7 KiB
Go
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
|
|
}
|