From deefa3a9f14c99fd2ed3e3ccfade6ba282589558 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Thu, 17 Oct 2024 13:11:27 +0200
Subject: [PATCH] fn: add ContextGuard from tapd repo

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.
---
 fn/context_guard.go | 113 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 fn/context_guard.go

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
+}