From 6b3b0a63e9a98d2091037ab9c5f6b96ff9893dae Mon Sep 17 00:00:00 2001
From: Filipe Menezes <filipe.menezes@mongodb.com>
Date: Fri, 10 Jun 2022 11:10:48 +0100
Subject: [PATCH] CLOUDP-125763: Fix race condition for sighandle

---
 internal/sighandle/sighandle.go | 52 ++++++++++++++++++++-------------
 1 file changed, 32 insertions(+), 20 deletions(-)

diff --git a/internal/sighandle/sighandle.go b/internal/sighandle/sighandle.go
index 528401e9e9..724c262d01 100644
--- a/internal/sighandle/sighandle.go
+++ b/internal/sighandle/sighandle.go
@@ -15,44 +15,56 @@
 package sighandle
 
 import (
+	"context"
 	"os"
 	"os/signal"
 )
 
 type handler struct {
-	sig      []os.Signal
-	f        func(os.Signal)
-	c        chan os.Signal
-	notified bool
+	sig         []os.Signal
+	handlerFunc func(os.Signal)
+	sigChannel  chan os.Signal
+	cancel      func()
 }
 
-func (h *handler) routine() {
-	h.f(<-h.c)
+func (h *handler) routine(ctx context.Context) {
+	select {
+	case <-ctx.Done():
+		// cancel this goroutine
+	case s := <-h.sigChannel:
+		h.handlerFunc(s)
+	}
 }
 
-func (h *handler) notify() {
-	if h.notified {
-		h.reset()
-	}
-	h.notified = true
-	h.c = make(chan os.Signal, 1)
-	go h.routine()
-	signal.Notify(h.c, h.sig...)
+func (h *handler) notify(ctx context.Context) {
+	go h.routine(ctx)
+	signal.Notify(h.sigChannel, h.sig...)
 }
 
 func (h *handler) reset() {
 	signal.Reset(h.sig...)
-	h.notified = false
+	h.cancel()
 }
 
-var std = &handler{}
+var std *handler
 
 func Notify(f func(os.Signal), sig ...os.Signal) {
-	std.f = f
-	std.sig = sig
-	std.notify()
+	Reset()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	std = &handler{
+		handlerFunc: f,
+		sig:         sig,
+		sigChannel:  make(chan os.Signal, 1),
+		cancel:      cancel,
+	}
+	std.notify(ctx)
 }
 
 func Reset() {
-	std.reset()
+	if std != nil {
+		std.reset()
+		std = nil
+	}
 }