build source

vendor/github.com/jackc/pgx/v5/pgconn/ctxwatch/context_watcher.go

@@ -0,0 +1,72 @@
+package ctxwatch
+
+import (
+	"context"
+	"sync"
+)
+
+// ContextWatcher watches a context and performs an action when the context is canceled. It can watch one context at a
+// time.
+type ContextWatcher struct {
+	handler Handler
+
+	// Lock protects the members below.
+	lock sync.Mutex
+	// Stop is the handle for an "after func". See [context.AfterFunc].
+	stop func() bool
+	done chan struct{}
+}
+
+// NewContextWatcher returns a ContextWatcher. onCancel will be called when a watched context is canceled.
+// OnUnwatchAfterCancel will be called when Unwatch is called and the watched context had already been canceled and
+// onCancel called.
+func NewContextWatcher(handler Handler) *ContextWatcher {
+	cw := &ContextWatcher{
+		handler: handler,
+	}
+
+	return cw
+}
+
+// Watch starts watching ctx. If ctx is canceled then the onCancel function passed to NewContextWatcher will be called.
+func (cw *ContextWatcher) Watch(ctx context.Context) {
+	cw.lock.Lock()
+	defer cw.lock.Unlock()
+
+	if cw.stop != nil {
+		panic("watch already in progress")
+	}
+
+	if ctx.Done() != nil {
+		cw.done = make(chan struct{})
+		cw.stop = context.AfterFunc(ctx, func() {
+			cw.handler.HandleCancel(ctx)
+			close(cw.done)
+		})
+	}
+}
+
+// Unwatch stops watching the previously watched context. If the onCancel function passed to NewContextWatcher was
+// called then onUnwatchAfterCancel will also be called.
+func (cw *ContextWatcher) Unwatch() {
+	cw.lock.Lock()
+	defer cw.lock.Unlock()
+
+	if cw.stop != nil {
+		if !cw.stop() {
+			<-cw.done
+			cw.handler.HandleUnwatchAfterCancel()
+		}
+		cw.stop = nil
+		cw.done = nil
+	}
+}
+
+type Handler interface {
+	// HandleCancel is called when the context that a ContextWatcher is currently watching is canceled. canceledCtx is the
+	// context that was canceled.
+	HandleCancel(canceledCtx context.Context)
+
+	// HandleUnwatchAfterCancel is called when a ContextWatcher that called HandleCancel on this Handler is unwatched.
+	HandleUnwatchAfterCancel()
+}