review suggestion

bchuecos · web-flow · commit 96167069845f · 2022-03-04T22:30:50.000Z
diff --git a/pkg/cmd/codespace/ports.go b/pkg/cmd/codespace/ports.go
@@ -275,7 +275,7 @@ func (a *App) UpdatePortVisibility(ctx context.Context, codespaceName string, ar
 	}
 	defer safeClose(session, &err)
 
-	errc := make(chan error, 1)
+	g, ctx := errgroup.WithContext(ctx)
 
 	// TODO: check if port visibility can be updated in parallel instead of sequentially
 	for _, port := range ports {
@@ -285,36 +285,33 @@ func (a *App) UpdatePortVisibility(ctx context.Context, codespaceName string, ar
 		ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
 		defer cancel()
 
-		go func(port portVisibility) {
+		g.Go(func() error {
 			updateNotif, err := session.WaitForPortNotification(ctx, port.number, liveshare.PortChangeKindUpdate)
 			if err != nil {
-				errc <- fmt.Errorf("error waiting for port %d to update: %w", port.number, err)
-				return
+				return fmt.Errorf("error waiting for port %d to update: %w", port.number, err)
+
 			}
 			if !updateNotif.Success {
 				if updateNotif.StatusCode == http.StatusForbidden {
-					errc <- newErrUpdatingPortVisibility(port.number, port.visibility, errUpdatePortVisibilityForbidden)
-					return
+					return newErrUpdatingPortVisibility(port.number, port.visibility, errUpdatePortVisibilityForbidden)
 				}
-				errc <- newErrUpdatingPortVisibility(port.number, port.visibility, errors.New(updateNotif.ErrorDetail))
-				return
-			}
-			errc <- nil // success
-		}(port)
+				return newErrUpdatingPortVisibility(port.number, port.visibility, errors.New(updateNotif.ErrorDetail))
 
-		err := session.UpdateSharedServerPrivacy(ctx, port.number, port.visibility)
-		if err != nil {
-			return fmt.Errorf("error updating port %d to %s: %w", port.number, port.visibility, err)
-		}
+			}
+			return nil // success
+		})
 
-		// wait for success or failure
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case err := <-errc:
+		g.Go(func() error {
+			err := session.UpdateSharedServerPrivacy(ctx, port.number, port.visibility)
 			if err != nil {
-				return err
+				return fmt.Errorf("error updating port %d to %s: %w", port.number, port.visibility, err)
 			}
+			return nil
+		})
+
+		// wait for success or failure
+		if err := g.Wait(); err != nil {
+			return err
 		}
 
 		a.StopProgressIndicator()
diff --git a/pkg/liveshare/ports.go b/pkg/liveshare/ports.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/sourcegraph/jsonrpc2"
+	"golang.org/x/sync/errgroup"
 )
 
 // Port describes a port exposed by the container.
@@ -41,33 +42,27 @@ type PortUpdate struct {
 // It returns an identifier that can be used to open an SSH channel to the remote port.
 func (s *Session) startSharing(ctx context.Context, sessionName string, port int) (channelID, error) {
 	args := []interface{}{port, sessionName, fmt.Sprintf("http://localhost:%d", port)}
-	errc := make(chan error, 1)
+	g, ctx := errgroup.WithContext(ctx)
 
-	go func() {
+	g.Go(func() error {
 		startNotification, err := s.WaitForPortNotification(ctx, port, PortChangeKindStart)
 		if err != nil {
-			errc <- fmt.Errorf("error while waiting for port notification: %w", err)
-			return
+			return fmt.Errorf("error while waiting for port notification: %w", err)
+
 		}
 		if !startNotification.Success {
-			errc <- fmt.Errorf("error while starting port sharing: %s", startNotification.ErrorDetail)
-			return
+			return fmt.Errorf("error while starting port sharing: %s", startNotification.ErrorDetail)
 		}
-		errc <- nil // success
-	}()
+		return nil // success
+	})
 
 	var response Port
-	if err := s.rpc.do(ctx, "serverSharing.startSharing", args, &response); err != nil {
-		return channelID{}, err
-	}
+	g.Go(func() error {
+		return s.rpc.do(ctx, "serverSharing.startSharing", args, &response)
+	})
 
-	select {
-	case <-ctx.Done():
-		return channelID{}, ctx.Err()
-	case err := <-errc:
-		if err != nil {
-			return channelID{}, err
-		}
+	if err := g.Wait(); err != nil {
+		return channelID{}, err
 	}
 
 	return channelID{response.StreamName, response.StreamCondition}, nil
