From 73c556a33af88edae1409a2cda2dfdb068edd4af Mon Sep 17 00:00:00 2001 From: Eno Compton Date: Tue, 21 Nov 2023 11:23:49 -0700 Subject: [PATCH 1/3] feat: add support for wait command To help ensure the Proxy is up and ready, this commit adds a wait command with an optional `--max` flag to set the maximum time to wait. By default when invoking this command: ./cloud-sql-proxy wait The Proxy will wait up to the maximum time for the /startup endpoint to respond. This command requires that the Proxy be started in another process with the HTTP health check enabled. If an alternate health check port or address is used, as in: ./cloud-sql-proxy --http-address 0.0.0.0 \ --http-port 9191 Then the wait command must also be told to use the same custom values: ./cloud-sql-proxy wait --http-address 0.0.0.0 \ --http-port 9191 By default the wait command will wait 30 seconds. To alter this value, use: ./cloud-sql-proxy wait --max 10s Fixes #1117 --- cmd/root.go | 95 +++++++++++++++++++++++++++++++++++++++++++----- cmd/wait_test.go | 61 +++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 9 deletions(-) create mode 100644 cmd/wait_test.go diff --git a/cmd/root.go b/cmd/root.go index bf770a1e8..f048a83a7 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -280,8 +280,38 @@ Localhost Admin Server /quitquitquit. The admin server exits gracefully when it receives a POST request at /quitquitquit. -(*) indicates a flag that may be used as a query parameter +Waiting for Startup + + Sometimes it is necessary to wait for the Proxy to start. + + To help ensure the Proxy is up and ready, the Proxy includes adds a wait + subcommand with an optional --max flag to set the maximum time to wait. + + By default when invoking this command: + + ./cloud-sql-proxy wait + + The Proxy will wait up to the maximum time for the /startup endpoint to + respond. This command requires that the Proxy be started in another + process with the HTTP health check enabled. If an alternate health + check port or address is used, as in: + + ./cloud-sql-proxy \ + --http-address 0.0.0.0 \ + --http-port 9191 + + Then the wait command must also be told to use the same custom values: + ./cloud-sql-proxy wait \ + --http-address 0.0.0.0 \ + --http-port 9191 + + By default the wait command will wait 30 seconds. To alter this value, + use: + + ./cloud-sql-proxy wait --max 10s + +(*) indicates a flag that may be used as a query parameter ` const envPrefix = "CSQL_PROXY" @@ -314,9 +344,44 @@ func instanceFromEnv(args []string) []string { return args } +const ( + waitMaxFlag = "max" + httpAddressFlag = "http-address" + httpPortFlag = "http-port" +) + +func runWaitCmd(c *cobra.Command, _ []string) error { + a, _ := c.Flags().GetString(httpAddressFlag) + p, _ := c.Flags().GetString(httpPortFlag) + addr := fmt.Sprintf("http://%v:%v/startup", a, p) + + wait, err := c.Flags().GetDuration(waitMaxFlag) + if err != nil { + // This error should always be nil. If the error occurs, it means the + // wait flag name has changed where it was registered. + return err + } + c.SilenceUsage = true + + t := time.After(wait) + for { + select { + case <-t: + return errors.New("command failed to complete successfully") + default: + resp, err := http.Get(addr) + if err != nil || resp.StatusCode != http.StatusOK { + time.Sleep(time.Second) + break + } + return nil + } + } +} + // NewCommand returns a Command object representing an invocation of the proxy. func NewCommand(opts ...Option) *Command { - cmd := &cobra.Command{ + rootCmd := &cobra.Command{ Use: "cloud-sql-proxy INSTANCE_CONNECTION_NAME...", Version: versionString, Short: "cloud-sql-proxy authorizes and encrypts connections to Cloud SQL.", @@ -325,7 +390,7 @@ func NewCommand(opts ...Option) *Command { logger := log.NewStdLogger(os.Stdout, os.Stderr) c := &Command{ - Command: cmd, + Command: rootCmd, logger: logger, cleanup: func() error { return nil }, conf: &proxy.Config{ @@ -336,7 +401,19 @@ func NewCommand(opts ...Option) *Command { o(c) } - cmd.Args = func(cmd *cobra.Command, args []string) error { + var waitCmd = &cobra.Command{ + Use: "wait", + RunE: runWaitCmd, + } + waitFlags := waitCmd.PersistentFlags() + waitFlags.DurationP( + waitMaxFlag, "m", + 30*time.Second, + "maximum amount of time to wait for startup", + ) + rootCmd.AddCommand(waitCmd) + + rootCmd.Args = func(cmd *cobra.Command, args []string) error { // If args is not already populated, try to read from the environment. if len(args) == 0 { args = instanceFromEnv(args) @@ -358,9 +435,9 @@ func NewCommand(opts ...Option) *Command { return nil } - cmd.RunE = func(*cobra.Command, []string) error { return runSignalWrapper(c) } + rootCmd.RunE = func(*cobra.Command, []string) error { return runSignalWrapper(c) } - pflags := cmd.PersistentFlags() + pflags := rootCmd.PersistentFlags() // Override Cobra's default messages. pflags.BoolP("help", "h", false, "Display help information for cloud-sql-proxy") @@ -405,9 +482,9 @@ the Proxy will then pick-up automatically.`) "Enable Prometheus HTTP endpoint /metrics on localhost") pflags.StringVar(&c.conf.PrometheusNamespace, "prometheus-namespace", "", "Use the provided Prometheus namespace for metrics") - pflags.StringVar(&c.conf.HTTPAddress, "http-address", "localhost", + pflags.StringVar(&c.conf.HTTPAddress, httpAddressFlag, "localhost", "Address for Prometheus and health check server") - pflags.StringVar(&c.conf.HTTPPort, "http-port", "9090", + pflags.StringVar(&c.conf.HTTPPort, httpPortFlag, "9090", "Port for Prometheus and health check server") pflags.BoolVar(&c.conf.Debug, "debug", false, "Enable pprof on the localhost admin server") @@ -432,7 +509,7 @@ https://cloud.google.com/storage/docs/requester-pays`) pflags.StringVar(&c.conf.ImpersonationChain, "impersonate-service-account", "", `Comma separated list of service accounts to impersonate. Last value is the target account.`) - cmd.PersistentFlags().BoolVar(&c.conf.Quiet, "quiet", false, "Log error messages only") + rootCmd.PersistentFlags().BoolVar(&c.conf.Quiet, "quiet", false, "Log error messages only") pflags.BoolVar(&c.conf.AutoIP, "auto-ip", false, `Supports legacy behavior of v1 and will try to connect to first IP address returned by the SQL Admin API. In most cases, this flag should not be used. diff --git a/cmd/wait_test.go b/cmd/wait_test.go new file mode 100644 index 000000000..94af153fb --- /dev/null +++ b/cmd/wait_test.go @@ -0,0 +1,61 @@ +// Copyright 2023 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "net" + "testing" +) + +func TestWaitCommandFlags(t *testing.T) { + ln, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + host, port, err := net.SplitHostPort(ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n")) + }() + + _, err = invokeProxyCommand([]string{ + "wait", + "--http-address", host, + "--http-port", port, + "--max=1s", + }) + if err != nil { + t.Fatal(err) + } +} + +func TestWaitCommandFails(t *testing.T) { + _, err := invokeProxyCommand([]string{ + "wait", + // assuming default host and port + "--max=500ms", + }) + if err == nil { + t.Fatal("wait should fail when endpoint does not respond") + } +} From 0d230183e3efb9526c81b524160394b363fd7eb6 Mon Sep 17 00:00:00 2001 From: Eno Compton Date: Wed, 22 Nov 2023 10:17:14 -0700 Subject: [PATCH 2/3] Improve documentation wording --- cmd/root.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index f048a83a7..5ad7693d1 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -284,16 +284,16 @@ Waiting for Startup Sometimes it is necessary to wait for the Proxy to start. - To help ensure the Proxy is up and ready, the Proxy includes adds a wait + To help ensure the Proxy is up and ready, the Proxy includes a wait subcommand with an optional --max flag to set the maximum time to wait. - By default when invoking this command: + Invoke the wait command, like this: ./cloud-sql-proxy wait - The Proxy will wait up to the maximum time for the /startup endpoint to - respond. This command requires that the Proxy be started in another - process with the HTTP health check enabled. If an alternate health + By default, the Proxy will wait up to the maximum time for the startup + endpoint to respond. The wait command requires that the Proxy be started in + another process with the HTTP health check enabled. If an alternate health check port or address is used, as in: ./cloud-sql-proxy \ From b61ca0cae086ad91db76754e2ae7254230017068 Mon Sep 17 00:00:00 2001 From: Eno Compton Date: Wed, 22 Nov 2023 11:16:49 -0700 Subject: [PATCH 3/3] Improve test reliability --- cmd/wait_test.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/cmd/wait_test.go b/cmd/wait_test.go index 94af153fb..fcc332b67 100644 --- a/cmd/wait_test.go +++ b/cmd/wait_test.go @@ -15,8 +15,10 @@ package cmd import ( + "io" "net" "testing" + "time" ) func TestWaitCommandFlags(t *testing.T) { @@ -35,6 +37,11 @@ func TestWaitCommandFlags(t *testing.T) { return } defer conn.Close() + // Use a read deadline to produce read error + conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + // Read client request first. + io.ReadAll(conn) + // Write a generic 200 response back. conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n")) }() @@ -53,7 +60,7 @@ func TestWaitCommandFails(t *testing.T) { _, err := invokeProxyCommand([]string{ "wait", // assuming default host and port - "--max=500ms", + "--max=100ms", }) if err == nil { t.Fatal("wait should fail when endpoint does not respond")