From 55ab2f5845a1f6680de2f24fb01687aabf84422b Mon Sep 17 00:00:00 2001
From: bufdev <4228796+bufdev@users.noreply.github.com>
Date: Sun, 6 Oct 2024 22:16:58 -0400
Subject: [PATCH] Clean up interrupt package (#3379)

---
 private/pkg/app/app.go                        |  7 ++--
 private/pkg/interrupt/interrupt.go            | 33 +++----------------
 ...nterrupt_windows.go => interrupt_other.go} |  8 ++---
 private/pkg/interrupt/interrupt_unix.go       | 14 +++-----
 4 files changed, 17 insertions(+), 45 deletions(-)
 rename private/pkg/interrupt/{interrupt_windows.go => interrupt_other.go} (76%)

diff --git a/private/pkg/app/app.go b/private/pkg/app/app.go
index c52c08fae0..0f9011de1d 100644
--- a/private/pkg/app/app.go
+++ b/private/pkg/app/app.go
@@ -326,8 +326,11 @@ func Main(ctx context.Context, f func(context.Context, Container) error) {
 // The run will be stopped on interrupt signal.
 // The exit code can be determined using GetExitCode.
 func Run(ctx context.Context, container Container, f func(context.Context, Container) error) error {
-	ctx, cancel := interrupt.WithCancel(ctx)
-	defer cancel()
+	ctx, cancel := interrupt.NotifyContext(ctx)
+	go func() {
+		<-ctx.Done()
+		cancel()
+	}()
 	if err := f(ctx, container); err != nil {
 		printError(container, err)
 		return err
diff --git a/private/pkg/interrupt/interrupt.go b/private/pkg/interrupt/interrupt.go
index 2f62494a56..50f04e1496 100644
--- a/private/pkg/interrupt/interrupt.go
+++ b/private/pkg/interrupt/interrupt.go
@@ -20,33 +20,10 @@ import (
 	"os/signal"
 )
 
-var signals = append(
-	[]os.Signal{
-		os.Interrupt,
-	},
-	extraSignals...,
-)
-
-// WithCancel returns a context that is cancelled if interrupt signals are sent.
-func WithCancel(ctx context.Context) (context.Context, context.CancelFunc) {
-	signalC, closer := NewSignalChannel()
-	ctx, cancel := context.WithCancel(ctx)
-	go func() {
-		<-signalC
-		closer()
-		cancel()
-	}()
-	return ctx, cancel
-}
+var interruptSignals = append([]os.Signal{os.Interrupt}, extraSignals...)
 
-// NewSignalChannel returns a new channel for interrupt signals.
-//
-// Call the returned function to cancel sending to this channel.
-func NewSignalChannel() (<-chan os.Signal, func()) {
-	signalC := make(chan os.Signal, 1)
-	signal.Notify(signalC, signals...)
-	return signalC, func() {
-		signal.Stop(signalC)
-		close(signalC)
-	}
+// NotifyContext returns a new [context.Context] from [signal.NotifyContext]
+// with the appropriate interrupt signals.
+func NotifyContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	return signal.NotifyContext(ctx, interruptSignals...)
 }
diff --git a/private/pkg/interrupt/interrupt_windows.go b/private/pkg/interrupt/interrupt_other.go
similarity index 76%
rename from private/pkg/interrupt/interrupt_windows.go
rename to private/pkg/interrupt/interrupt_other.go
index 8589a5404b..5a5dd97924 100644
--- a/private/pkg/interrupt/interrupt_windows.go
+++ b/private/pkg/interrupt/interrupt_other.go
@@ -12,8 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//go:build windows
-// +build windows
+//go:build !unix
 
 package interrupt
 
@@ -22,7 +21,6 @@ import "os"
 // extraSignals are signals beyond os.Interrupt that we want to be handled
 // as interrupts.
 //
-// For unix-like platforms, this adds syscall.SIGTERM, although this is only
-// tested on darwin and linux, which buf officially supports. Other unix-like
-// platforms should have this as well, however.
+// For unix-like platforms, this adds syscall.SIGTERM. No other signals
+// are added for other platforms.
 var extraSignals = []os.Signal{}
diff --git a/private/pkg/interrupt/interrupt_unix.go b/private/pkg/interrupt/interrupt_unix.go
index def8fd59c3..903ca1ddc7 100644
--- a/private/pkg/interrupt/interrupt_unix.go
+++ b/private/pkg/interrupt/interrupt_unix.go
@@ -12,10 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Matching the unix-like build tags in the Golang source i.e. https://github.com/golang/go/blob/912f0750472dd4f674b69ca1616bfaf377af1805/src/os/file_unix.go#L6
-
-//go:build aix || darwin || dragonfly || freebsd || (js && wasm) || linux || netbsd || openbsd || solaris
-// +build aix darwin dragonfly freebsd js,wasm linux netbsd openbsd solaris
+//go:build unix
 
 package interrupt
 
@@ -27,9 +24,6 @@ import (
 // extraSignals are signals beyond os.Interrupt that we want to be handled
 // as interrupts.
 //
-// For unix-like platforms, this adds syscall.SIGTERM, although this is only
-// tested on darwin and linux, which buf officially supports. Other unix-like
-// platforms should have this as well, however.
-var extraSignals = []os.Signal{
-	syscall.SIGTERM,
-}
+// For unix-like platforms, this adds syscall.SIGTERM. No other signals
+// are added for other platforms.
+var extraSignals = []os.Signal{syscall.SIGTERM}