Fix Terraform resource agent association

coder · kylecarbs · Mar 29, 2022 · Mar 26, 2022 · Mar 26, 2022 · Mar 27, 2022
commit 82f60dbd7813696628183c5f7b66f5cf79ba3c82
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -8,8 +8,10 @@
     "drpcconn",
     "drpcmux",
     "drpcserver",
+    "Dsts",
     "fatih",
     "goarch",
+    "gographviz",
     "goleak",
     "gossh",
     "hashicorp",

diff --git a/cli/cliui/agent.go b/cli/cliui/agent.go
@@ -0,0 +1,73 @@
+package cliui
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/briandowns/spinner"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/codersdk"
+)
+
+type AgentOptions struct {
+	WorkspaceName string
+	Fetch         func(context.Context) (codersdk.WorkspaceResource, error)
+	FetchInterval time.Duration
+	WarnInterval  time.Duration
+}
+
+func Agent(cmd *cobra.Command, opts AgentOptions) error {
+	if opts.FetchInterval == 0 {
+		opts.FetchInterval = 500 * time.Millisecond
+	}
+	if opts.WarnInterval == 0 {
+		opts.WarnInterval = 30 * time.Second
+	}
+	resource, err := opts.Fetch(cmd.Context())
+	if err != nil {
+		return xerrors.Errorf("fetch: %w", err)
+	}
+	if resource.Agent.Status == codersdk.WorkspaceAgentConnected {
+		return nil
+	}
+	if resource.Agent.Status == codersdk.WorkspaceAgentDisconnected {
+		opts.WarnInterval = 0
+	}
+	spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
+	spin.Writer = cmd.OutOrStdout()
+	spin.Suffix = " Waiting for connection from " + Styles.Field.Render(resource.Type+"."+resource.Name) + "..."
+	spin.Start()
+	defer spin.Stop()
+
+	ticker := time.NewTicker(opts.FetchInterval)
+	defer ticker.Stop()
+	timer := time.NewTimer(opts.WarnInterval)
+	defer timer.Stop()
+	for {
+		select {
+		case <-cmd.Context().Done():
+			return cmd.Context().Err()
+		case <-timer.C:
+			message := "Don't panic, your workspace is booting up!"
+			if resource.Agent.Status == codersdk.WorkspaceAgentDisconnected {
+				message = "The workspace agent lost connection! Wait for it to reconnect or run: " + Styles.Code.Render("coder workspaces rebuild "+opts.WorkspaceName)
+			}
+			// This saves the cursor position, then defers clearing from the cursor
+			// position to the end of the screen.
+			fmt.Fprintf(cmd.OutOrStdout(), "\033[s\r\033[2K%s\n\n", Styles.Paragraph.Render(Styles.Prompt.String()+message))
+			defer fmt.Fprintf(cmd.OutOrStdout(), "\033[u\033[J")
+		case <-ticker.C:
+		}
+		resource, err = opts.Fetch(cmd.Context())
+		if err != nil {
+			return xerrors.Errorf("fetch: %w", err)
+		}
+		if resource.Agent.Status != codersdk.WorkspaceAgentConnected {
+			continue
+		}
+		return nil
+	}
+}
diff --git a/cli/cliui/provisionerjob.go b/cli/cliui/provisionerjob.go
@@ -8,13 +8,29 @@ import (
 	"sync"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
 )
 
+func WorkspaceBuild(cmd *cobra.Command, client *codersdk.Client, build uuid.UUID, before time.Time) error {
+	return ProvisionerJob(cmd, ProvisionerJobOptions{
+		Fetch: func() (codersdk.ProvisionerJob, error) {
+			build, err := client.WorkspaceBuild(cmd.Context(), build)
+			return build.Job, err
+		},
+		Cancel: func() error {
+			return client.CancelWorkspaceBuild(cmd.Context(), build)
+		},
+		Logs: func() (<-chan codersdk.ProvisionerJobLog, error) {
+			return client.WorkspaceBuildLogsAfter(cmd.Context(), build, before)
+		},
+	})
+}
+
 type ProvisionerJobOptions struct {
 	Fetch  func() (codersdk.ProvisionerJob, error)
 	Cancel func() error

diff --git a/cli/ssh.go b/cli/ssh.go
@@ -1,10 +1,14 @@
 package cli
 
 import (
+	"context"
+
+	"github.com/pion/webrtc/v3"
 	"github.com/spf13/cobra"
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
 )
@@ -21,6 +25,12 @@ func ssh() *cobra.Command {
 			if err != nil {
 				return err
 			}
+			if workspace.LatestBuild.Job.CompletedAt == nil {
+				err = cliui.WorkspaceBuild(cmd, client, workspace.LatestBuild.ID, workspace.CreatedAt)
+				if err != nil {
+					return err
+				}
+			}
 			if workspace.LatestBuild.Transition == database.WorkspaceTransitionDelete {
 				return xerrors.New("workspace is deleting...")
 			}
@@ -57,11 +67,19 @@ func ssh() *cobra.Command {
 				}
 				return xerrors.Errorf("no sshable agent with address %q: %+v", resourceAddress, resourceKeys)
 			}
-			if resource.Agent.LastConnectedAt == nil {
-				return xerrors.Errorf("agent hasn't connected yet")
+			err = cliui.Agent(cmd, cliui.AgentOptions{
+				WorkspaceName: workspace.Name,
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceResource, error) {
+					return client.WorkspaceResource(ctx, resource.ID)
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("await agent: %w", err)
 			}
 
-			conn, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID, nil, nil)
+			conn, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID, []webrtc.ICEServer{{
+				URLs: []string{"stun:stun.l.google.com:19302"},
+			}}, nil)
 			if err != nil {
 				return err
 			}

diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -24,7 +24,7 @@ func TestSSH(t *testing.T) {
 		t.Parallel()
 		client := coderdtest.New(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		daemonCloser := coderdtest.NewProvisionerDaemon(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
 		agentToken := uuid.NewString()
 		version := coderdtest.CreateProjectVersion(t, client, user.OrganizationID, &echo.Responses{
 			Parse:           echo.ParseComplete,
@@ -49,15 +49,19 @@ func TestSSH(t *testing.T) {
 		coderdtest.AwaitProjectVersionJob(t, client, version.ID)
 		project := coderdtest.CreateProject(t, client, user.OrganizationID, version.ID)
 		workspace := coderdtest.CreateWorkspace(t, client, "", project.ID)
-		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
-		daemonCloser.Close()
-		agentClient := codersdk.New(client.URL)
-		agentClient.SessionToken = agentToken
-		agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
-			Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
-		})
-		defer agentCloser.Close()
-		coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
+		go func() {
+			// Run this async so the SSH command has to wait for
+			// the build and agent to connect!
+			coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+			agentClient := codersdk.New(client.URL)
+			agentClient.SessionToken = agentToken
+			agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
+				Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
+			})
+			t.Cleanup(func() {
+				_ = agentCloser.Close()
+			})
+		}()
 
 		cmd, root := clitest.New(t, "ssh", workspace.Name)
 		clitest.SetupConfig(t, client, root)

diff --git a/cli/workspacecreate.go b/cli/workspacecreate.go
@@ -6,7 +6,6 @@ import (
 	"sort"
 	"time"
 
-	"github.com/briandowns/spinner"
 	"github.com/fatih/color"
 	"github.com/manifoldco/promptui"
 	"github.com/spf13/cobra"
@@ -161,40 +160,9 @@ func workspaceCreate() *cobra.Command {
 			if err != nil {
 				return err
 			}
-			resources, err = client.WorkspaceResourcesByBuild(cmd.Context(), workspace.LatestBuild.ID)
-			if err != nil {
-				return err
-			}
-			spin := spinner.New(spinner.CharSets[5], 100*time.Millisecond, spinner.WithColor("fgGreen"))
-			spin.Writer = cmd.OutOrStdout()
-			spin.Suffix = " Waiting for agent to connect..."
-			spin.Start()
-			defer spin.Stop()
-			for _, resource := range resources {
-				if resource.Agent == nil {
-					continue
-				}
-				ticker := time.NewTicker(1 * time.Second)
-				for {
-					select {
-					case <-cmd.Context().Done():
-						return nil
-					case <-ticker.C:
-					}
-					resource, err := client.WorkspaceResource(cmd.Context(), resource.ID)
-					if err != nil {
-						return err
-					}
-					if resource.Agent.FirstConnectedAt == nil {
-						continue
-					}
-					spin.Stop()
-					_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\nThe %s workspace has been created!\n\n", cliui.Styles.Keyword.Render(workspace.Name))
-					_, _ = fmt.Fprintln(cmd.OutOrStdout(), "  "+cliui.Styles.Code.Render("coder ssh "+workspace.Name))
-					_, _ = fmt.Fprintln(cmd.OutOrStdout())
-					break
-				}
-			}
+			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\nThe %s workspace has been created!\n\n", cliui.Styles.Keyword.Render(workspace.Name))
+			_, _ = fmt.Fprintln(cmd.OutOrStdout(), "  "+cliui.Styles.Code.Render("coder ssh "+workspace.Name))
+			_, _ = fmt.Fprintln(cmd.OutOrStdout())
 
 			return err
 		},

diff --git a/cmd/cliui/main.go b/cmd/cliui/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -155,6 +156,35 @@ func main() {
 		},
 	})
 
+	root.AddCommand(&cobra.Command{
+		Use: "agent",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			resource := codersdk.WorkspaceResource{
+				Type: "google_compute_instance",
+				Name: "dev",
+				Agent: &codersdk.WorkspaceAgent{
+					Status: codersdk.WorkspaceAgentDisconnected,
+				},
+			}
+			go func() {
+				time.Sleep(3 * time.Second)
+				resource.Agent.Status = codersdk.WorkspaceAgentConnected
+			}()
+			err := cliui.Agent(cmd, cliui.AgentOptions{
+				WorkspaceName: "dev",
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceResource, error) {
+					return resource, nil
+				},
+				WarnInterval: 2 * time.Second,
+			})
+			if err != nil {
+				return err
+			}
+			fmt.Printf("Completed!\n")
+			return nil
+		},
+	})
+
 	err := root.Execute()
 	if err != nil {
 		_, _ = fmt.Println(err.Error())