dial-nodes: add timeout flag

A recent refactor got rid of a default timeout for the dial nodes
subcommand. This change adds a flag -timeout that can be passed into the
command so that it doesn't have to wait the full 30 seconds to get a
timeout error back.

Changelog: changed

3 files changed, 56 insertions, 11 deletions

diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index fc67c7471..0ff0baa4a 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -24,7 +24,7 @@ type subcmd interface {
 }
 
 const (
-	defaultDialTimeout        = 30 * time.Second
+	defaultDialTimeout        = 10 * time.Second
 	paramVirtualStorage       = "virtual-storage"
 	paramRelativePath         = "repository"
 	paramAuthoritativeStorage = "authoritative-storage"
diff --git a/cmd/praefect/subcmd_dial_nodes.go b/cmd/praefect/subcmd_dial_nodes.go
index 06d7719ec..71054d6ad 100644
--- a/cmd/praefect/subcmd_dial_nodes.go
+++ b/cmd/praefect/subcmd_dial_nodes.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"io"
+	"time"
 
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/nodes"
@@ -18,14 +19,29 @@ func newDialNodesSubcommand(w io.Writer) *dialNodesSubcommand {
 }
 
 type dialNodesSubcommand struct {
-	w io.Writer
+	w       io.Writer
+	timeout time.Duration
 }
 
 func (s *dialNodesSubcommand) FlagSet() *flag.FlagSet {
-	return flag.NewFlagSet(dialNodesCmdName, flag.ExitOnError)
+	fs := flag.NewFlagSet(dialNodesCmdName, flag.ExitOnError)
+	fs.DurationVar(&s.timeout, "timeout", 0, "timeout for dialing gitaly nodes")
+	fs.Usage = func() {
+		printfErr("Description:\n" +
+			"	This command attempts to reach all Gitaly nodes.\n")
+		fs.PrintDefaults()
+	}
+
+	return fs
 }
 
 func (s *dialNodesSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
-	ctx := context.Background()
+	if s.timeout == 0 {
+		s.timeout = defaultDialTimeout
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
+	defer cancel()
+
 	return nodes.PingAll(ctx, conf, nodes.NewTextPrinter(s.w), false)
 }
diff --git a/cmd/praefect/subcmd_dial_nodes_test.go b/cmd/praefect/subcmd_dial_nodes_test.go
index 9817b32ad..fd2877f60 100644
--- a/cmd/praefect/subcmd_dial_nodes_test.go
+++ b/cmd/praefect/subcmd_dial_nodes_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
@@ -44,10 +45,11 @@ func TestSubCmdDialNodes(t *testing.T) {
 	}
 
 	for _, tt := range []struct {
-		name string
-		conf config.Config
-		resp *gitalypb.ServerInfoResponse
-		logs string
+		name   string
+		conf   config.Config
+		resp   *gitalypb.ServerInfoResponse
+		logs   string
+		errMsg string
 	}{
 		{
 			name: "2 virtuals, 2 storages, 1 node",
@@ -97,6 +99,26 @@ func TestSubCmdDialNodes(t *testing.T) {
 				"SUCCESS: confirmed Gitaly storage \"2\" in virtual storages [storage-1] is served",
 				"SUCCESS: node configuration is consistent!",
 			}), ""),
+			errMsg: "",
+		},
+		{
+			name: "node unreachable",
+			conf: config.Config{
+				VirtualStorages: []*config.VirtualStorage{
+					{
+						Name: "default",
+						Nodes: []*config.Node{
+							{
+								Storage: "1",
+								Address: "unix:///unreachable/socket",
+							},
+						},
+					},
+				},
+			},
+			resp:   nil,
+			logs:   "",
+			errMsg: "the following nodes are not healthy: unix:///unreachable/socket",
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
@@ -105,10 +127,17 @@ func TestSubCmdDialNodes(t *testing.T) {
 
 			output := &bytes.Buffer{}
 
-			cmd := newDialNodesSubcommand(output)
-			require.NoError(t, cmd.Exec(nil, tt.conf))
+			cmd := dialNodesSubcommand{w: output, timeout: 1 * time.Second}
+
+			err := cmd.Exec(nil, tt.conf)
+
+			if tt.errMsg == "" {
+				require.NoError(t, err)
+				require.Equal(t, tt.logs, output.String())
+				return
+			}
 
-			require.Equal(t, tt.logs, output.String())
+			require.Equal(t, tt.errMsg, err.Error())
 		})
 	}
 }