fetch remote ssh usernames in parallel

znull · znull · commit 06eb5ad47fbb · 2021-12-17T16:10:44.000-07:00
diff --git a/pkg/cmd/codespace/ssh.go b/pkg/cmd/codespace/ssh.go
@@ -162,7 +162,6 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 }
 
 func (a *App) printOpenSSHConfig(ctx context.Context, opts configOptions) error {
-	// Ensure all child tasks (e.g. port forwarding) terminate before return.
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -181,68 +180,76 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts configOptions) error
 		return fmt.Errorf("error getting codespace info: %w", err)
 	}
 
-	t, err := template.New("ssh_config").Parse(`Host cs.{{.Name}}.{{.EscapedRef}}
-	User {{.SSHUser}}
-	ProxyCommand {{.GHExec}} cs ssh -c {{.Name}} --stdio
-	UserKnownHostsFile=/dev/null
-	StrictHostKeyChecking no
-	LogLevel quiet
-	ControlMaster auto
-
-`)
-	if err != nil {
-		return fmt.Errorf("error formatting template: %w", err)
-	}
-
-	ghexec, err := os.Executable()
-	if err != nil {
-		return err
-	}
-
-	// store a mapping of repository -> remote ssh username. This is
-	// necessary because the username can vary between codespaces, but
-	// since fetching it is slow, we store it here so we at least only do
-	// it once per repository.
-	sshUsers := map[string]string{}
-
+	sshUsers := make(chan sshResult)
+	fetches := 0
 	var status error
 	for _, cs := range codespaces {
-
 		if cs.State != "Available" {
 			fmt.Fprintf(os.Stderr, "skipping unavailable codespace %s: %s\n", cs.Name, cs.State)
 			status = cmdutil.SilentError
 			continue
 		}
 
-		sshUser, ok := sshUsers[cs.Repository.FullName]
-		if !ok {
+		cs := cs
+		fetches += 1
+		go func() {
+			result := sshResult{}
+			defer func() {
+				select {
+				case sshUsers <- result:
+				case <-ctx.Done():
+				}
+			}()
+
 			session, err := openSSHSession(ctx, a, cs, nil)
 			if err != nil {
-				fmt.Fprintf(os.Stderr, "error connecting to codespace: %v\n", err)
-
-				// Move on to the next codespace. We don't want to bail here - just because we're not
-				// able to set up connectivity to one doesn't mean we shouldn't make a best effort to
-				// generate configs for the rest of them.
-				status = cmdutil.SilentError
-				continue
+				result.err = fmt.Errorf("error connecting to codespace: %w", err)
+				return
 			}
 			defer session.Close()
 
-			a.StartProgressIndicatorWithLabel(fmt.Sprintf("Fetching SSH Details for %s", cs.Name))
-			_, sshUser, err = session.StartSSHServer(ctx)
-			a.StopProgressIndicator()
+			//a.StartProgressIndicatorWithLabel(fmt.Sprintf("Fetching SSH Details for %s", cs.Name))
+			_, result.user, err = session.StartSSHServer(ctx)
+			//a.StopProgressIndicator()
 			if err != nil {
-				fmt.Fprintf(os.Stderr, "error getting ssh server details: %v", err)
-				status = cmdutil.SilentError
-				continue // see above
+				result.err = fmt.Errorf("error getting ssh server details: %w", err)
+				return
 			}
-			sshUsers[cs.Repository.FullName] = sshUser
+
+			result.codespace = cs
+		}()
+	}
+
+	ghexec, err := os.Executable()
+	if err != nil {
+		return err
+	}
+
+	t, err := template.New("ssh_config").Parse(`Host cs.{{.Name}}.{{.EscapedRef}}
+	User {{.SSHUser}}
+	ProxyCommand {{.GHExec}} cs ssh -c {{.Name}} --stdio
+	UserKnownHostsFile=/dev/null
+	StrictHostKeyChecking no
+	LogLevel quiet
+	ControlMaster auto
+
+`)
+	if err != nil {
+		return fmt.Errorf("error formatting template: %w", err)
+	}
+
+	for i := 0; i < fetches; i++ {
+		result := <-sshUsers
+		if result.err != nil {
+			fmt.Fprintf(os.Stderr, "%v\n", result.err)
+			status = cmdutil.SilentError
+			continue
 		}
 
 		conf := codespaceSSHConfig{
-			Name:       cs.Name,
-			EscapedRef: strings.ReplaceAll(cs.GitStatus.Ref, "/", "-"),
-			SSHUser:    sshUser,
+			Name:       result.codespace.Name,
+			EscapedRef: strings.ReplaceAll(result.codespace.GitStatus.Ref, "/", "-"),
+			SSHUser:    result.user,
 			GHExec:     ghexec,
 		}
 		if err := t.Execute(a.io.Out, conf); err != nil {
@@ -253,6 +260,12 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts configOptions) error
 	return status
 }
 
+type sshResult struct {
+	codespace *api.Codespace
+	user      string // on success, the remote ssh username; else nil
+	err       error
+}
+
 // codespaceSSHConfig contains values needed to write an OpenSSH host
 // configuration for a single codespace. For example:
 //
