From 5a372695969876fc9715709f3a46d03f5375b58f Mon Sep 17 00:00:00 2001 From: Matt Robenolt Date: Wed, 29 Nov 2023 14:20:24 -0800 Subject: [PATCH] Add new `ping` command (#765) --- go.mod | 2 +- go.sum | 4 +- internal/cmd/ping/ping.go | 192 ++++++++++++++++++++++++++++++++++++++ internal/cmd/root.go | 2 + 4 files changed, 197 insertions(+), 3 deletions(-) create mode 100644 internal/cmd/ping/ping.go diff --git a/go.mod b/go.mod index 03e35720..360cc4b4 100644 --- a/go.mod +++ b/go.mod @@ -22,7 +22,7 @@ require ( github.com/mitchellh/go-homedir v1.1.0 github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 github.com/pkg/errors v0.9.1 - github.com/planetscale/planetscale-go v0.93.0 + github.com/planetscale/planetscale-go v0.93.1-0.20231128084905-be5d5eb26e2f github.com/planetscale/sql-proxy v0.13.0 github.com/spf13/cobra v1.8.0 github.com/spf13/pflag v1.0.5 diff --git a/go.sum b/go.sum index ee8803e7..484c1ee4 100644 --- a/go.sum +++ b/go.sum @@ -232,8 +232,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/sftp v1.13.1/go.mod h1:3HaPG6Dq1ILlpPZRO0HVMrsydcdLt6HRDccSgb87qRg= github.com/planetscale/planetscale-go v0.51.0/go.mod h1:+rGpW2u7iQZZx4O/nFj4MZe4xIS22CVegEgl1IkTExQ= -github.com/planetscale/planetscale-go v0.93.0 h1:cogKCS4pEbHqurZWiPROaEAo16ZEPmRTSdd2IZioa8A= -github.com/planetscale/planetscale-go v0.93.0/go.mod h1:MR+LZLhZTIzmi2x91nwLQyIedMFhn6sgQDAKs0wAGNQ= +github.com/planetscale/planetscale-go v0.93.1-0.20231128084905-be5d5eb26e2f h1:BrNKNknUpfXYzb7syB99WGiVwN0yOFTq8c/3pFjxX1Q= +github.com/planetscale/planetscale-go v0.93.1-0.20231128084905-be5d5eb26e2f/go.mod h1:hDSA/dClhuKuW8dNhKN9vQW8E5fo034Rb6qGTf86/yI= github.com/planetscale/sql-proxy v0.13.0 h1:NDjcdqgoNzwbZQTyoIDEoI+K7keC5RRKvdML2roAMn4= github.com/planetscale/sql-proxy v0.13.0/go.mod h1:4Sk6JdoBqQhHv9V4FCOC27YIM3EjU8cLIsw5HqxN8x4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= diff --git a/internal/cmd/ping/ping.go b/internal/cmd/ping/ping.go new file mode 100644 index 00000000..c8902f5f --- /dev/null +++ b/internal/cmd/ping/ping.go @@ -0,0 +1,192 @@ +package ping + +import ( + "cmp" + "context" + "fmt" + "net" + "net/netip" + "slices" + "strings" + "sync" + "time" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/planetscale-go/planetscale" + "github.com/spf13/cobra" +) + +const baseDomain = ".connect.psdb.cloud" + +func PingCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + timeout time.Duration + concurrency uint8 + } + + cmd := &cobra.Command{ + Use: "ping", + Short: "Ping public PlanetScale database endpoints", + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + // XXX: explicitly use a new client that doesn't use authentication + client, err := ps.NewClient( + ps.WithBaseURL(ps.DefaultBaseURL), + ) + if err != nil { + return err + } + + end := ch.Printer.PrintProgress("Fetching regions...") + defer end() + + regions, err := client.Regions.List(ctx, &ps.ListRegionsRequest{}) + if err != nil { + return cmdutil.HandleError(err) + } + end() + + // set of unique providers for the optimized endpoints + providers := make(map[string]struct{}) + for _, r := range regions { + providers[strings.ToLower(r.Provider)] = struct{}{} + } + + endpoints := make([]string, 0, len(providers)+len(regions)) + for p := range providers { + endpoints = append(endpoints, p) + } + for _, r := range regions { + endpoints = append(endpoints, r.Slug) + } + + end = ch.Printer.PrintProgress("Pinging endpoints...") + defer end() + + results := pingEndpoints( + ctx, + endpoints, + flags.concurrency, + flags.timeout, + ) + end() + + return ch.Printer.PrintResource(toResultsTable(results, providers)) + }, + } + + cmd.PersistentFlags().DurationVar(&flags.timeout, "timeout", + 5*time.Second, "Timeout for a ping to succeed.") + cmd.PersistentFlags().Uint8Var(&flags.concurrency, "concurrency", + 8, "Number of concurrent pings.") + + return cmd +} + +type pingResult struct { + key string + d time.Duration + err error +} + +type Result struct { + Endpoint string `header:"endpoint" json:"endpoint"` + Latency string `header:"latency" json:"latency"` + Type string `header:"type" json:"type"` +} + +func directOrOptimized(key string, providers map[string]struct{}) string { + if _, ok := providers[key]; ok { + return "optimized" + } + return "direct" +} + +func toResultsTable(results []pingResult, providers map[string]struct{}) []*Result { + rs := make([]*Result, 0, len(results)) + for _, r := range results { + row := &Result{ + Endpoint: makeHostname(r.key), + Type: directOrOptimized(r.key, providers), + } + if r.err == nil { + row.Latency = r.d.Truncate(100 * time.Microsecond).String() + } else { + row.Latency = "---" + } + + rs = append(rs, row) + } + return rs +} + +func makeHostname(subdomain string) string { + return subdomain + baseDomain +} + +func pingEndpoints(ctx context.Context, eps []string, concurrency uint8, timeout time.Duration) []pingResult { + var ( + wg sync.WaitGroup + mu sync.Mutex + sem = make(chan struct{}, concurrency) + results = make([]pingResult, 0, len(eps)) + ) + + defer close(sem) + + for _, ep := range eps { + ep := ep + wg.Add(1) + select { + case <-ctx.Done(): + return results + case sem <- struct{}{}: + go func() { + d, err := pingEndpoint(ctx, makeHostname(ep), timeout) + // XXX: on failures, set the duration to the timeout so they are sorted last + if err != nil { + d = timeout + } + + mu.Lock() + results = append(results, pingResult{ep, d, err}) + mu.Unlock() + wg.Done() + <-sem + }() + } + } + + wg.Wait() + slices.SortFunc(results, func(a, b pingResult) int { return cmp.Compare(a.d, b.d) }) + return results +} + +func pingEndpoint(ctx context.Context, hostname string, timeout time.Duration) (time.Duration, error) { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + // make hostname a FQDN if not already + if hostname[len(hostname)-1] != '.' { + hostname = hostname + "." + } + + // separately look up DNS so that's separate from our connection time + addrs, err := net.DefaultResolver.LookupNetIP(ctx, "ip", hostname) + if err != nil { + return 0, err + } + if len(addrs) == 0 { + return 0, fmt.Errorf("unable to resolve addr: %s", hostname) + } + + // explicitly time to establish a TCP connection + var d net.Dialer + start := time.Now() + conn, err := d.DialContext(ctx, "tcp", netip.AddrPortFrom(addrs[0], 443).String()) + if err != nil { + return 0, err + } + defer conn.Close() + return time.Since(start), nil +} diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 91560280..409eae20 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -38,6 +38,7 @@ import ( "github.com/planetscale/cli/internal/cmd/deployrequest" "github.com/planetscale/cli/internal/cmd/org" "github.com/planetscale/cli/internal/cmd/password" + "github.com/planetscale/cli/internal/cmd/ping" "github.com/planetscale/cli/internal/cmd/region" "github.com/planetscale/cli/internal/cmd/shell" "github.com/planetscale/cli/internal/cmd/signup" @@ -217,6 +218,7 @@ func runCmd(ctx context.Context, ver, commit, buildDate string, format *printer. rootCmd.AddCommand(deployrequest.DeployRequestCmd(ch)) rootCmd.AddCommand(org.OrgCmd(ch)) rootCmd.AddCommand(password.PasswordCmd(ch)) + rootCmd.AddCommand(ping.PingCmd(ch)) rootCmd.AddCommand(region.RegionCmd(ch)) rootCmd.AddCommand(shell.ShellCmd(ch)) rootCmd.AddCommand(signup.SignupCmd(ch))