feat: add coder ping

coder · coadler · Feb 13, 2023 · Feb 10, 2023 · Feb 10, 2023 · Feb 10, 2023
commit 2f9532f7edad866788bc914ecd85b4f5de0e983a
diff --git a/cli/ping.go b/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
+}
diff --git a/cli/ping_test.go b/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
+	})
+}
diff --git a/cli/root.go b/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(),

diff --git a/cli/speedtest.go b/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
 					}

diff --git a/cli/vscodessh.go b/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
 	}

diff --git a/codersdk/workspaceagentconn.go b/codersdk/workspaceagentconn.go
@@ -17,6 +17,7 @@ import (
 	"github.com/google/uuid"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
+	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/net/speedtest"
 
 	"github.com/coder/coder/coderd/tracing"
@@ -136,7 +137,7 @@ func (c *WorkspaceAgentConn) AwaitReachable(ctx context.Context) bool {
 
 // Ping pings the agent and returns the round-trip time.
 // The bool returns true if the ping was made P2P.
-func (c *WorkspaceAgentConn) Ping(ctx context.Context) (time.Duration, bool, error) {
+func (c *WorkspaceAgentConn) Ping(ctx context.Context) (time.Duration, bool, *ipnstate.PingResult, error) {
 	ctx, span := tracing.StartSpan(ctx)
 	defer span.End()
 

diff --git a/enterprise/coderd/replicas_test.go b/enterprise/coderd/replicas_test.go
@@ -84,7 +84,7 @@ func TestReplicas(t *testing.T) {
 		require.Eventually(t, func() bool {
 			ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitShort)
 			defer cancelFunc()
-			_, _, err = conn.Ping(ctx)
+			_, _, _, err = conn.Ping(ctx)
 			return err == nil
 		}, testutil.WaitLong, testutil.IntervalFast)
 		_ = conn.Close()
@@ -129,7 +129,7 @@ func TestReplicas(t *testing.T) {
 		require.Eventually(t, func() bool {
 			ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.IntervalSlow)
 			defer cancelFunc()
-			_, _, err = conn.Ping(ctx)
+			_, _, _, err = conn.Ping(ctx)
 			return err == nil
 		}, testutil.WaitLong, testutil.IntervalFast)
 		_ = conn.Close()

diff --git a/scaletest/agentconn/run.go b/scaletest/agentconn/run.go
@@ -141,7 +141,7 @@ func waitForDisco(ctx context.Context, logs io.Writer, conn *codersdk.WorkspaceA
 	for i := 0; i < pingAttempts; i++ {
 		_, _ = fmt.Fprintf(logs, "\tDisco ping attempt %d/%d...\n", i+1, pingAttempts)
 		pingCtx, cancel := context.WithTimeout(ctx, defaultRequestTimeout)
-		_, p2p, err := conn.Ping(pingCtx)
+		_, p2p, _, err := conn.Ping(pingCtx)
 		cancel()
 		if err == nil {
 			_, _ = fmt.Fprintf(logs, "\tDisco ping succeeded after %d attempts, p2p = %v\n", i+1, p2p)

diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -432,7 +432,7 @@ func (c *Conn) Status() *ipnstate.Status {
 
 // Ping sends a Disco ping to the Wireguard engine.
 // The bool returned is true if the ping was performed P2P.
-func (c *Conn) Ping(ctx context.Context, ip netip.Addr) (time.Duration, bool, error) {
+func (c *Conn) Ping(ctx context.Context, ip netip.Addr) (time.Duration, bool, *ipnstate.PingResult, error) {
 	errCh := make(chan error, 1)
 	prChan := make(chan *ipnstate.PingResult, 1)
 	go c.wireguardEngine.Ping(ip, tailcfg.PingDisco, func(pr *ipnstate.PingResult) {
@@ -444,11 +444,11 @@ func (c *Conn) Ping(ctx context.Context, ip netip.Addr) (time.Duration, bool, er
 	})
 	select {
 	case err := <-errCh:
-		return 0, false, err
+		return 0, false, nil, err
 	case <-ctx.Done():
-		return 0, false, ctx.Err()
+		return 0, false, nil, ctx.Err()
 	case pr := <-prChan:
-		return time.Duration(pr.LatencySeconds * float64(time.Second)), pr.Endpoint != "", nil
+		return time.Duration(pr.LatencySeconds * float64(time.Second)), pr.Endpoint != "", pr, nil
 	}
 }
 
@@ -477,7 +477,7 @@ func (c *Conn) AwaitReachable(ctx context.Context, ip netip.Addr) bool {
 		ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
 		defer cancel()
 
-		_, _, err := c.Ping(ctx, ip)
+		_, _, _, err := c.Ping(ctx, ip)
 		if err == nil {
 			completed()
 		}