use errgroup

cmd/ghcs/logs.go

@@ -11,6 +11,7 @@ import (
 	"github.com/github/ghcs/internal/codespaces"
 	"github.com/github/go-liveshare"
 	"github.com/spf13/cobra"
+	"golang.org/x/sync/errgroup"
 )
 
 func newLogsCmd() *cobra.Command {
@@ -84,27 +85,12 @@ func logs(ctx context.Context, tail bool, codespaceName string) error {
 		ctx, localPort, dst, fmt.Sprintf("%s /workspaces/.codespaces/.persistedshare/creation.log", cmdType),
 	)
 
-	// Error channels are buffered so that neither sending goroutine gets stuck.
-
-	tunnelClosed := make(chan error, 1)
-	go func() {
+	group, ctx := errgroup.WithContext(ctx)
+	group.Go(func() error {
 		fwd := liveshare.NewPortForwarder(session, "sshd", remoteSSHServerPort)
-		tunnelClosed <- fwd.ForwardToListener(ctx, listen) // error is non-nil
-	}()
-
-	cmdDone := make(chan error, 1)
-	go func() {
-		cmdDone <- cmd.Run()
-	}()
-
-	select {
-	case err := <-tunnelClosed:
+		err := fwd.ForwardToListener(ctx, listen) // error is non-nil
 		return fmt.Errorf("connection closed: %v", err)
-
-	case err := <-cmdDone:
-		if err != nil {
-			return fmt.Errorf("error retrieving logs: %v", err)
-		}
-		return nil // success
-	}
+	})
+	group.Go(cmd.Run)
+	return group.Wait()
 }

cmd/ghcs/ports.go

@@ -16,6 +16,7 @@ import (
 	"github.com/github/go-liveshare"
 	"github.com/muhammadmuzzammil1998/jsonc"
 	"github.com/spf13/cobra"
+	"golang.org/x/sync/errgroup"
 )
 
 // portOptions represents the options accepted by the ports command.
@@ -272,27 +273,22 @@ func forwardPorts(log *output.Logger, codespaceName string, ports []string) erro
 
 	// Run forwarding of all ports concurrently, aborting all of
 	// them at the first failure, including cancellation of the context.
-	errc := make(chan error, len(portPairs))
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
+	group, ctx := errgroup.WithContext(ctx)
 	for _, pair := range portPairs {
 		pair := pair
-
-		go func() {
+		group.Go(func() error {
 			listen, err := liveshare.ListenTCP(pair.local)
 			if err != nil {
-				errc <- err
-				return
+				return nil
 			}
 			defer listen.Close()
 			log.Printf("Forwarding ports: remote %d <=> local %d\n", pair.remote, pair.local)
 			name := fmt.Sprintf("share-%d", pair.remote)
 			fwd := liveshare.NewPortForwarder(session, name, pair.remote)
-			errc <- fwd.ForwardToListener(ctx, listen) // error always non-nil
-		}()
+			return fwd.ForwardToListener(ctx, listen) // error always non-nil
+		})
 	}
-
-	return <-errc // first error
+	return group.Wait() // first error
 }
 
 type portPair struct {

cmd/ghcs/ssh.go

@@ -13,6 +13,7 @@ import (
 	"github.com/github/ghcs/internal/codespaces"
 	"github.com/github/go-liveshare"
 	"github.com/spf13/cobra"
+	"golang.org/x/sync/errgroup"
 )
 
 func newSSHCmd() *cobra.Command {
@@ -97,28 +98,20 @@ func ssh(ctx context.Context, sshProfile, codespaceName string, localSSHServerPo
 		connectDestination = fmt.Sprintf("%s@localhost", sshUser)
 	}
 
-	tunnelClosed := make(chan error)
-	go func() {
-		fwd := liveshare.NewPortForwarder(session, "sshd", remoteSSHServerPort)
-		tunnelClosed <- fwd.ForwardToListener(ctx, listen) // error is always non-nil
-	}()
-
-	shellClosed := make(chan error)
-	go func() {
-		shellClosed <- codespaces.Shell(ctx, log, localSSHServerPort, connectDestination, usingCustomPort)
-	}()
-
 	log.Println("Ready...")
-	select {
-	case err := <-tunnelClosed:
+	group, ctx := errgroup.WithContext(ctx)
+	group.Go(func() error {
+		fwd := liveshare.NewPortForwarder(session, "sshd", remoteSSHServerPort)
+		err := fwd.ForwardToListener(ctx, listen) // always non-nil
 		return fmt.Errorf("tunnel closed: %v", err)
-
-	case err := <-shellClosed:
-		if err != nil {
+	})
+	group.Go(func() error {
+		if err := codespaces.Shell(ctx, log, localSSHServerPort, connectDestination, usingCustomPort); err != nil {
 			return fmt.Errorf("shell closed: %v", err)
 		}
 		return nil // success
-	}
+	})
+	return group.Wait()
 }
 
 func getContainerID(ctx context.Context, logger *output.Logger, terminal *liveshare.Terminal) (string, error) {