praefect: Sub-command interruption with SIGINT

Execution of any sub-commands should be possible to interrupt by
sending SIGINT signal to the process. It is so for all except the
'serve' command as it has its own handling mechanism.

2 files changed, 47 insertions, 13 deletions

diff --git a/internal/cli/praefect/main.go b/internal/cli/praefect/main.go
index 1b2ec299e..b1fcf2007 100644
--- a/internal/cli/praefect/main.go
+++ b/internal/cli/praefect/main.go
@@ -15,10 +15,12 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 
 	"github.com/urfave/cli/v2"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/praefect/service"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/version"
+	"golang.org/x/exp/slices"
 )
 
 func init() {
@@ -37,6 +39,8 @@ const (
 
 // NewApp returns a new praefect app.
 func NewApp() *cli.App {
+	interrupt := make(chan os.Signal, 1)
+
 	return &cli.App{
 		Name:    progname,
 		Usage:   "a gitaly proxy",
@@ -78,6 +82,29 @@ func NewApp() *cli.App {
 				Usage: "load configuration from `FILE`",
 			},
 		},
+		Before: func(appCtx *cli.Context) error {
+			// Praefect service manages os.Interrupt on its own, by making a "table-flip".
+			// That is why the signal listening is omitted if there are no arguments passed
+			// (old-fashioned method of starting Praefect service) or 'serve' sub-command
+			// is invoked. Other sub-commands require signal to be properly handled.
+			args := appCtx.Args().Slice()
+			if len(args) == 0 || slices.Contains(args, "serve") {
+				return nil
+			}
+
+			signal.Notify(interrupt, os.Interrupt)
+			go func() {
+				if _, ok := <-interrupt; ok {
+					os.Exit(130) // indicates program was interrupted
+				}
+			}()
+
+			return nil
+		},
+		After: func(*cli.Context) error {
+			close(interrupt)
+			return nil
+		},
 	}
 }
 
diff --git a/internal/cli/praefect/subcmd_check_test.go b/internal/cli/praefect/subcmd_check_test.go
index ba8804625..c970b8520 100644
--- a/internal/cli/praefect/subcmd_check_test.go
+++ b/internal/cli/praefect/subcmd_check_test.go
@@ -181,19 +181,17 @@ Checking check 3...Failed (warning) error: i failed but not too badly
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			var stdout, stderr bytes.Buffer
-			app := NewApp()
-			app.Writer = &stdout
-			app.ErrWriter = &stderr
-			for i, cmd := range app.Commands {
-				if cmd.Name == "check" {
-					app.Commands[i] = newCheckCommand(tc.checks)
-					break
-				}
-			}
-
 			t.Run("quiet", func(t *testing.T) {
-				stdout.Reset()
+				var stdout, stderr bytes.Buffer
+				app := NewApp()
+				app.Writer = &stdout
+				app.ErrWriter = &stderr
+				for i, cmd := range app.Commands {
+					if cmd.Name == "check" {
+						app.Commands[i] = newCheckCommand(tc.checks)
+						break
+					}
+				}
 				err := app.Run(append([]string{progname, "-config", confPath, "check", "-q"}, tc.args...))
 				assert.Equal(t, tc.expectedError, err)
 				if len(tc.args) == 0 {
@@ -203,7 +201,16 @@ Checking check 3...Failed (warning) error: i failed but not too badly
 			})
 
 			t.Run("normal", func(t *testing.T) {
-				stdout.Reset()
+				var stdout, stderr bytes.Buffer
+				app := NewApp()
+				app.Writer = &stdout
+				app.ErrWriter = &stderr
+				for i, cmd := range app.Commands {
+					if cmd.Name == "check" {
+						app.Commands[i] = newCheckCommand(tc.checks)
+						break
+					}
+				}
 				err := app.Run(append([]string{progname, "-config", confPath, "check"}, tc.args...))
 				assert.Equal(t, tc.expectedError, err)
 				if len(tc.args) == 0 {