feat: add coder ping (#6161)

cli/ping.go

@@ -0,0 +1,138 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/sloghuman"
+
+	"github.com/coder/coder/cli/cliui"
+	"github.com/coder/coder/codersdk"
+)
+
+func ping() *cobra.Command {
+	var (
+		pingNum     int
+		pingTimeout time.Duration
+		pingWait    time.Duration
+		verbose     bool
+	)
+	cmd := &cobra.Command{
+		Annotations: workspaceCommand,
+		Use:         "ping <workspace>",
+		Short:       "Ping a workspace",
+		Args:        cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx, cancel := context.WithCancel(cmd.Context())
+			defer cancel()
+
+			client, err := CreateClient(cmd)
+			if err != nil {
+				return err
+			}
+
+			workspaceName := args[0]
+			_, workspaceAgent, err := getWorkspaceAndAgent(ctx, cmd, client, codersdk.Me, workspaceName, false)
+			if err != nil {
+				return err
+			}
+
+			var logger slog.Logger
+			if verbose {
+				logger = slog.Make(sloghuman.Sink(cmd.OutOrStdout())).Leveled(slog.LevelDebug)
+			}
+
+			conn, err := client.DialWorkspaceAgent(ctx, workspaceAgent.ID, &codersdk.DialWorkspaceAgentOptions{Logger: logger})
+			if err != nil {
+				return err
+			}
+			defer conn.Close()
+
+			derpMap := conn.DERPMap()
+			_ = derpMap
+
+			n := 0
+			didP2p := false
+			start := time.Now()
+			for {
+				if n > 0 {
+					time.Sleep(time.Second)
+				}
+				n++
+
+				ctx, cancel := context.WithTimeout(ctx, pingTimeout)
+				dur, p2p, pong, err := conn.Ping(ctx)
+				cancel()
+				if err != nil {
+					if xerrors.Is(err, context.DeadlineExceeded) {
+						_, _ = fmt.Fprintf(cmd.OutOrStdout(), "ping to %q timed out \n", workspaceName)
+						if n == pingNum {
+							return nil
+						}
+						continue
+					}
+					if xerrors.Is(err, context.Canceled) {
+						return nil
+					}
+
+					if err.Error() == "no matching peer" {
+						continue
+					}
+
+					_, _ = fmt.Fprintf(cmd.OutOrStdout(), "ping to %q failed %s\n", workspaceName, err.Error())
+					if n == pingNum {
+						return nil
+					}
+					continue
+				}
+
+				dur = dur.Round(time.Millisecond)
+				var via string
+				if p2p {
+					if !didP2p {
+						_, _ = fmt.Fprintln(cmd.OutOrStdout(), "p2p connection established in",
+							cliui.Styles.DateTimeStamp.Render(time.Since(start).Round(time.Millisecond).String()),
+						)
+					}
+					didP2p = true
+
+					via = fmt.Sprintf("%s via %s",
+						cliui.Styles.Fuchsia.Render("p2p"),
+						cliui.Styles.Code.Render(pong.Endpoint),
+					)
+				} else {
+					derpName := "unknown"
+					derpRegion, ok := derpMap.Regions[pong.DERPRegionID]
+					if ok {
+						derpName = derpRegion.RegionName
+					}
+					via = fmt.Sprintf("%s via %s",
+						cliui.Styles.Fuchsia.Render("proxied"),
+						cliui.Styles.Code.Render(fmt.Sprintf("DERP(%s)", derpName)),
+					)
+				}
+
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "pong from %s %s in %s\n",
+					cliui.Styles.Keyword.Render(workspaceName),
+					via,
+					cliui.Styles.DateTimeStamp.Render(dur.String()),
+				)
+
+				if n == pingNum {
+					return nil
+				}
+			}
+		},
+	}
+
+	cmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "Enables verbose logging.")
+	cmd.Flags().DurationVarP(&pingWait, "wait", "", time.Second, "Specifies how long to wait between pings.")
+	cmd.Flags().DurationVarP(&pingTimeout, "timeout", "t", 5*time.Second, "Specifies how long to wait for a ping to complete.")
+	cmd.Flags().IntVarP(&pingNum, "num", "n", 10, "Specifies the number of pings to perform.")
+	return cmd
+}

cli/ping_test.go

@@ -0,0 +1,54 @@
+package cli_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"cdr.dev/slog/sloggers/slogtest"
+
+	"github.com/coder/coder/agent"
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/codersdk/agentsdk"
+	"github.com/coder/coder/pty/ptytest"
+	"github.com/coder/coder/testutil"
+)
+
+func TestPing(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
+		cmd, root := clitest.New(t, "ping", workspace.Name)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetErr(pty.Output())
+		cmd.SetOut(pty.Output())
+
+		agentClient := agentsdk.New(client.URL)
+		agentClient.SetSessionToken(agentToken)
+		agentCloser := agent.New(agent.Options{
+			Client: agentClient,
+			Logger: slogtest.Make(t, nil).Named("agent"),
+		})
+		defer func() {
+			_ = agentCloser.Close()
+		}()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		cmdDone := tGo(t, func() {
+			err := cmd.ExecuteContext(ctx)
+			assert.NoError(t, err)
+		})
+
+		pty.ExpectMatch("pong from " + workspace.Name)
+		cancel()
+		<-cmdDone
+	})
+}

cli/root.go

@@ -85,10 +85,12 @@ func Core() []*cobra.Command {
 		login(),
 		logout(),
 		parameters(),
+		ping(),
 		portForward(),
 		publickey(),
 		rename(),
 		resetPassword(),
+		restart(),
 		scaletest(),
 		schedules(),
 		show(),
@@ -97,7 +99,6 @@ func Core() []*cobra.Command {
 		start(),
 		state(),
 		stop(),
-		restart(),
 		templates(),
 		tokens(),
 		update(),

cli/speedtest.go

@@ -71,7 +71,7 @@ func speedtest() *cobra.Command {
 						return ctx.Err()
 					case <-ticker.C:
 					}
-					dur, p2p, err := conn.Ping(ctx)
+					dur, p2p, _, err := conn.Ping(ctx)
 					if err != nil {
 						continue
 					}

cli/testdata/coder_--help.golden

@@ -36,6 +36,7 @@ Workspace Commands:
   create         Create a workspace
   delete         Delete a workspace
   list           List workspaces
+  ping           Ping a workspace
   rename         Rename a workspace
   restart        Restart a workspace
   schedule       Schedule automated start and stop times for workspaces

cli/testdata/coder_ping_--help.golden

@@ -0,0 +1,26 @@
+Ping a workspace
+
+Usage:
+  coder ping <workspace> [flags]
+
+Flags:
+  -h, --help               help for ping
+  -n, --num int            Specifies the number of pings to perform. (default 10)
+  -t, --timeout duration   Specifies how long to wait for a ping to complete. (default 5s)
+  -v, --verbose            Enables verbose logging.
+      --wait duration      Specifies how long to wait between pings. (default 1s)
+
+Global Flags:
+      --global-config coder   Path to the global coder config directory.
+                              Consumes $CODER_CONFIG_DIR (default "~/.config/coderv2")
+      --header stringArray    HTTP headers added to all requests. Provide as "Key=Value".
+                              Consumes $CODER_HEADER
+      --no-feature-warning    Suppress warnings about unlicensed features.
+                              Consumes $CODER_NO_FEATURE_WARNING
+      --no-version-warning    Suppress warning when client and server versions do not match.
+                              Consumes $CODER_NO_VERSION_WARNING
+      --token string          Specify an authentication token. For security reasons setting
+                              CODER_SESSION_TOKEN is preferred.
+                              Consumes $CODER_SESSION_TOKEN
+      --url string            URL to a deployment.
+                              Consumes $CODER_URL

cli/vscodessh.go

@@ -204,7 +204,7 @@ type sshNetworkStats struct {
 }
 
 func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, start, end time.Time, counts map[netlogtype.Connection]netlogtype.Counts) (*sshNetworkStats, error) {
-	latency, p2p, err := agentConn.Ping(ctx)
+	latency, p2p, _, err := agentConn.Ping(ctx)
 	if err != nil {
 		return nil, err
 	}