Update the signal type pbench handles

It should be SIGTERM, SIGINT, and SIGQUIT
prestodb · Jun 20, 2024 · 4a2ae07 · 4a2ae07
1 parent f21bf29
commit 4a2ae07
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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.