From 47dae851a9513bb5e3d98b59a33aef909a2bf125 Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Wed, 7 Jun 2023 09:36:53 -0600
Subject: [PATCH] feat: add connection test for startup (#1832)

When callers want to ensure their instance is in fact reachable, the
--run-connection-test flag will attempt to connection to all registered
instances.

Fixes #348
---
 cmd/root.go                  |  8 +++++++
 cmd/root_test.go             | 14 +++++++++++
 internal/proxy/proxy.go      | 13 ++++++++++
 internal/proxy/proxy_test.go | 46 ++++++++++++++++++++++++++++++++++++
 4 files changed, 81 insertions(+)

diff --git a/cmd/root.go b/cmd/root.go
index ffabed926..b831b4f0b 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -435,6 +435,10 @@ is the target account.`)
 address returned by the SQL Admin API. In most cases, this flag should not be used.
 Prefer default of public IP or use --private-ip instead.`)
 
+	pflags.BoolVar(&c.conf.RunConnectionTest, "run-connection-test", false, `Runs a connection test
+against all specified instances. If an instance is unreachable, the Proxy exits with a failure
+status code.`)
+
 	// Global and per instance flags
 	pflags.StringVarP(&c.conf.Addr, "address", "a", "127.0.0.1",
 		"(*) Address to bind Cloud SQL instance listeners.")
@@ -474,6 +478,10 @@ func parseConfig(cmd *Command, conf *proxy.Config, args []string) error {
 	}
 
 	if conf.FUSEDir != "" {
+		if conf.RunConnectionTest {
+			return newBadCommandError("cannot run connection tests in FUSE mode")
+		}
+
 		if err := proxy.SupportsFUSE(); err != nil {
 			return newBadCommandError(
 				fmt.Sprintf("--fuse is not supported: %v", err),
diff --git a/cmd/root_test.go b/cmd/root_test.go
index 312071d41..6283e56fd 100644
--- a/cmd/root_test.go
+++ b/cmd/root_test.go
@@ -412,6 +412,13 @@ func TestNewCommandArguments(t *testing.T) {
 				AutoIP: true,
 			}),
 		},
+		{
+			desc: "using the run-connection-test flag",
+			args: []string{"--run-connection-test", "proj:region:inst"},
+			want: withDefaults(&proxy.Config{
+				RunConnectionTest: true,
+			}),
+		},
 	}
 
 	for _, tc := range tcs {
@@ -1057,6 +1064,13 @@ func TestNewCommandWithErrors(t *testing.T) {
 				"p:r:i?private-ip=true",
 			},
 		},
+		{
+			desc: "run-connection-test with fuse",
+			args: []string{
+				"--run-connection-test",
+				"--fuse", "myfusedir",
+			},
+		},
 	}
 
 	for _, tc := range tcs {
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index 98cf89b83..461db9597 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -233,6 +233,10 @@ type Config struct {
 	// OtherUserAgents is a list of space separate user agents that will be
 	// appended to the default user agent.
 	OtherUserAgents string
+
+	// RunConnectionTest determines whether the Proxy should attempt a connection
+	// to all specified instances to verify the network path is valid.
+	RunConnectionTest bool
 }
 
 // dialOptions interprets appropriate dial options for a particular instance
@@ -571,6 +575,15 @@ func (c *Client) Serve(ctx context.Context, notify func()) error {
 		return c.serveFuse(ctx, notify)
 	}
 
+	if c.conf.RunConnectionTest {
+		c.logger.Infof("Connection test started")
+		if _, err := c.CheckConnections(ctx); err != nil {
+			c.logger.Errorf("Connection test failed")
+			return err
+		}
+		c.logger.Infof("Connection test passed")
+	}
+
 	exitCh := make(chan error)
 	for _, m := range c.mnts {
 		go func(mnt *socketMount) {
diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go
index 57f013a96..8d509295e 100644
--- a/internal/proxy/proxy_test.go
+++ b/internal/proxy/proxy_test.go
@@ -665,3 +665,49 @@ func TestCheckConnections(t *testing.T) {
 		t.Fatalf("CheckConnections number of connections: want = %v, got = %v", want, got)
 	}
 }
+
+func TestRunConnectionCheck(t *testing.T) {
+	in := &proxy.Config{
+		Addr: "127.0.0.1",
+		Port: 50002,
+		Instances: []proxy.InstanceConnConfig{
+			{Name: "proj:region:pg"},
+		},
+		RunConnectionTest: true,
+	}
+	d := &fakeDialer{}
+	c, err := proxy.NewClient(context.Background(), d, testLogger, in)
+	if err != nil {
+		t.Fatalf("proxy.NewClient error: %v", err)
+	}
+	defer func(c *proxy.Client) {
+		err := c.Close()
+		if err != nil {
+			t.Log(err)
+		}
+	}(c)
+	go func() {
+		// Serve alone without any connections will still verify that the
+		// provided instances are reachable.
+		_ = c.Serve(context.Background(), func() {})
+	}()
+
+	verifyDialAttempts := func() error {
+		var tries int
+		for {
+			tries++
+			if tries == 10 {
+				return errors.New("failed to verify dial tries after 10 tries")
+			}
+			if got := d.dialAttempts(); got > 0 {
+				return nil
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+
+	if err := verifyDialAttempts(); err != nil {
+		t.Fatal(err)
+	}
+
+}