From 4a2ae07ce77e57cd5d723385117d56049fd043c7 Mon Sep 17 00:00:00 2001
From: "Yiqun (Ethan) Zhang" <github@yzhang.io>
Date: Wed, 19 Jun 2024 20:15:58 -0500
Subject: [PATCH] Update the signal type pbench handles

It should be SIGTERM, SIGINT, and SIGQUIT
---
 cmd/loadjson/main.go | 5 +++--
 cmd/replay/main.go   | 5 +++--
 cmd/save/main.go     | 5 +++--
 stage/stage.go       | 3 ++-
 4 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/cmd/loadjson/main.go b/cmd/loadjson/main.go
index 6a444765..fc820ee2 100644
--- a/cmd/loadjson/main.go
+++ b/cmd/loadjson/main.go
@@ -15,6 +15,7 @@ import (
 	"pbench/utils"
 	"reflect"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -60,8 +61,8 @@ func Run(_ *cobra.Command, args []string) {
 	log.Info().Int("parallelism", Parallelism).Send()
 	ctx, cancel := context.WithCancel(context.Background())
 	timeToExit := make(chan os.Signal, 1)
-	signal.Notify(timeToExit, os.Interrupt, os.Kill)
-	// Handle SIGKILL and SIGINT. When ctx is canceled, in-progress MySQL transactions and InfluxDB operations will roll back.
+	signal.Notify(timeToExit, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+	// Handle SIGINT, SIGTERM, and SIGQUIT. When ctx is canceled, in-progress MySQL transactions and InfluxDB operations will roll back.
 	go func() {
 		sig := <-timeToExit
 		if sig != nil {
diff --git a/cmd/replay/main.go b/cmd/replay/main.go
index 4578bc82..bcb35567 100644
--- a/cmd/replay/main.go
+++ b/cmd/replay/main.go
@@ -13,6 +13,7 @@ import (
 	"pbench/presto"
 	"pbench/utils"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -43,8 +44,8 @@ func Run(_ *cobra.Command, args []string) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	timeToExit := make(chan os.Signal, 1)
-	signal.Notify(timeToExit, os.Interrupt, os.Kill)
-	// Handle SIGKILL and SIGINT. When ctx is canceled, in-progress MySQL transactions and InfluxDB operations will roll back.
+	signal.Notify(timeToExit, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+	// Handle SIGINT, SIGTERM, and SIGQUIT. When ctx is canceled, in-progress MySQL transactions and InfluxDB operations will roll back.
 	go func() {
 		sig := <-timeToExit
 		if sig != nil {
diff --git a/cmd/save/main.go b/cmd/save/main.go
index 256556dc..019b355f 100644
--- a/cmd/save/main.go
+++ b/cmd/save/main.go
@@ -15,6 +15,7 @@ import (
 	"pbench/utils"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -46,8 +47,8 @@ func Run(_ *cobra.Command, args []string) {
 	log.Info().Int("parallelism", Parallelism).Send()
 	ctx, cancel := context.WithCancel(context.Background())
 	timeToExit := make(chan os.Signal, 1)
-	signal.Notify(timeToExit, os.Interrupt, os.Kill)
-	// Handle SIGKILL and SIGINT. When ctx is canceled, in-progress MySQL transactions and InfluxDB operations will roll back.
+	signal.Notify(timeToExit, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+	// Handle SIGINT, SIGTERM, and SIGQUIT. When ctx is canceled, in-progress MySQL transactions and InfluxDB operations will roll back.
 	go func() {
 		sig := <-timeToExit
 		if sig != nil {
diff --git a/stage/stage.go b/stage/stage.go
index 0f40edd3..896a54cf 100644
--- a/stage/stage.go
+++ b/stage/stage.go
@@ -20,6 +20,7 @@ import (
 	"strconv"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/rs/zerolog"
@@ -131,7 +132,7 @@ func (s *Stage) Run(ctx context.Context) int {
 	results := make([]*QueryResult, 0, len(s.Queries)+len(s.QueryFiles))
 	s.States.resultChan = make(chan *QueryResult, 16)
 	timeToExit := make(chan os.Signal, 1)
-	signal.Notify(timeToExit, os.Interrupt, os.Kill)
+	signal.Notify(timeToExit, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 	// Each goroutine we spawn will increment this wait group (count-down latch). We may start a goroutine for running
 	// a benchmark stage, or write query output to disk asynchronously.
 	// This wait group is propagated to the descendant stages.