refactor(google): update getGroups function to use context and errgroup for concurrency

Author: codrutpanea

connector/google/google.go

@@ -11,14 +11,14 @@ import (
 	"sort"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
 	"github.com/coreos/go-oidc/v3/oidc"
 	"golang.org/x/exp/slices"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
+	"golang.org/x/sync/errgroup"
 	admin "google.golang.org/api/admin/directory/v1"
 	"google.golang.org/api/impersonate"
 	"google.golang.org/api/option"
@@ -30,7 +30,10 @@ import (
 const (
 	issuerURL                  = "https://accounts.google.com"
 	wildcardDomainToAdminEmail = "*"
-	maxConcurrentGroupLookups  = 10
+
+	// defaultConcurrentGroupLookups is the limit used when Config.MaxConcurrentGroupLookups
+	// is zero or negative.
+	defaultConcurrentGroupLookups = 10
 )
 
 // Config holds configuration options for Google logins.
@@ -65,6 +68,10 @@ type Config struct {
 	// If this field is true, fetch direct group membership and transitive group membership
 	FetchTransitiveGroupMembership bool `json:"fetchTransitiveGroupMembership"`
 
+	// MaxConcurrentGroupLookups limits concurrent Admin Directory API calls when resolving
+	// transitive group membership. If zero or negative, the connector default limit applies.
+	MaxConcurrentGroupLookups int `json:"maxConcurrentGroupLookups"`
+
 	// Optional value for the prompt parameter, defaults to consent when offline_access
 	// scope is requested
 	PromptType *string `json:"promptType"`
@@ -123,6 +130,11 @@ func (c *Config) Open(id string, logger *slog.Logger) (conn connector.Connector,
 	}
 
 	clientID := c.ClientID
+	maxConcurrent := c.MaxConcurrentGroupLookups
+	if maxConcurrent <= 0 {
+		maxConcurrent = defaultConcurrentGroupLookups
+	}
+
 	return &googleConnector{
 		redirectURI: c.RedirectURI,
 		oauth2Config: &oauth2.Config{
@@ -142,6 +154,7 @@ func (c *Config) Open(id string, logger *slog.Logger) (conn connector.Connector,
 		serviceAccountFilePath:         c.ServiceAccountFilePath,
 		domainToAdminEmail:             c.DomainToAdminEmail,
 		fetchTransitiveGroupMembership: c.FetchTransitiveGroupMembership,
+		maxConcurrentGroupLookups:      maxConcurrent,
 		adminSrv:                       adminSrv,
 		promptType:                     promptType,
 	}, nil
@@ -163,6 +176,7 @@ type googleConnector struct {
 	serviceAccountFilePath         string
 	domainToAdminEmail             map[string]string
 	fetchTransitiveGroupMembership bool
+	maxConcurrentGroupLookups      int
 	adminSrv                       map[string]*admin.Service
 	promptType                     string
 }
@@ -276,8 +290,7 @@ func (c *googleConnector) createIdentity(ctx context.Context, identity connector
 
 	var groups []string
 	if s.Groups && len(c.adminSrv) > 0 {
-		checkedGroups := make(map[string]struct{})
-		groups, err = c.getGroups(claims.Email, c.fetchTransitiveGroupMembership, checkedGroups)
+		groups, err = c.getGroups(ctx, claims.Email, c.fetchTransitiveGroupMembership)
 		if err != nil {
 			return identity, fmt.Errorf("google: could not retrieve groups: %v", err)
 		}
@@ -302,26 +315,23 @@ func (c *googleConnector) createIdentity(ctx context.Context, identity connector
 }
 
 // getGroups creates a connection to the admin directory service and lists
-// all groups the user is a member of
-func (c *googleConnector) getGroups(email string, fetchTransitiveGroupMembership bool, checkedGroups map[string]struct{}) ([]string, error) {
-	if checkedGroups == nil {
-		checkedGroups = make(map[string]struct{})
-	}
-
-	directGroups, err := c.listGroupEmails(email)
+// all groups the user is a member of.
+func (c *googleConnector) getGroups(ctx context.Context, email string, fetchTransitiveGroupMembership bool) ([]string, error) {
+	directGroups, err := c.listGroupEmails(ctx, email)
 	if err != nil {
 		return nil, err
 	}
 
-	checkedGroupsMu := sync.Mutex{}
+	var seenMu sync.Mutex
+	seen := make(map[string]struct{})
 	userGroups := make([]string, 0, len(directGroups))
 	addGroup := func(groupEmail string) bool {
-		checkedGroupsMu.Lock()
-		defer checkedGroupsMu.Unlock()
-		if _, exists := checkedGroups[groupEmail]; exists {
+		seenMu.Lock()
+		defer seenMu.Unlock()
+		if _, exists := seen[groupEmail]; exists {
 			return false
 		}
-		checkedGroups[groupEmail] = struct{}{}
+		seen[groupEmail] = struct{}{}
 		// TODO (joelspeed): Make desired group key configurable
 		userGroups = append(userGroups, groupEmail)
 		return true
@@ -338,62 +348,41 @@ func (c *googleConnector) getGroups(email string, fetchTransitiveGroupMembership
 		return userGroups, nil
 	}
 
-	// Limit concurrent Google API calls while traversing transitive membership.
-	concurrencyLimiter := make(chan struct{}, maxConcurrentGroupLookups)
-	var workWG sync.WaitGroup
-	var firstErr error
-	var firstErrOnce sync.Once
-	var hasError atomic.Bool
-
-	setError := func(err error) {
-		firstErrOnce.Do(func() {
-			firstErr = err
-			hasError.Store(true)
-		})
-	}
-
-	var traverse func(string)
-	traverse = func(groupEmail string) {
-		defer workWG.Done()
-		if hasError.Load() {
-			return
-		}
-
-		concurrencyLimiter <- struct{}{}
-		if hasError.Load() {
-			<-concurrencyLimiter
-			return
-		}
-		parentGroups, err := c.listGroupEmails(groupEmail)
-		<-concurrencyLimiter
-		if err != nil {
-			setError(fmt.Errorf("could not list transitive groups: %v", err))
-			return
-		}
+	g, gctx := errgroup.WithContext(ctx)
+	g.SetLimit(c.maxConcurrentGroupLookups)
 
-		for _, parentGroupEmail := range parentGroups {
-			if !addGroup(parentGroupEmail) {
-				continue
+	var enqueue func(string)
+	enqueue = func(groupEmail string) {
+		g.Go(func() error {
+			if err := gctx.Err(); err != nil {
+				return err
 			}
-			workWG.Add(1)
-			go traverse(parentGroupEmail)
-		}
+			parentGroups, err := c.listGroupEmails(gctx, groupEmail)
+			if err != nil {
+				return fmt.Errorf("could not list transitive groups: %w", err)
+			}
+			for _, parent := range parentGroups {
+				if addGroup(parent) {
+					enqueue(parent)
+				}
+			}
+			return nil
+		})
 	}
 
 	for _, groupEmail := range seeds {
-		workWG.Add(1)
-		go traverse(groupEmail)
+		enqueue(groupEmail)
 	}
-	workWG.Wait()
-	if firstErr != nil {
-		return nil, firstErr
+
+	if err := g.Wait(); err != nil {
+		return nil, err
 	}
 
 	sort.Strings(userGroups)
 	return userGroups, nil
 }
 
-func (c *googleConnector) listGroupEmails(userKey string) ([]string, error) {
+func (c *googleConnector) listGroupEmails(ctx context.Context, userKey string) ([]string, error) {
 	domain := c.extractDomainFromEmail(userKey)
 	adminSrv, err := c.findAdminService(domain)
 	if err != nil {
@@ -403,8 +392,11 @@ func (c *googleConnector) listGroupEmails(userKey string) ([]string, error) {
 	groupEmails := []string{}
 	groupsList := &admin.Groups{}
 	for {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
 		groupsList, err = adminSrv.Groups.List().
-			UserKey(userKey).PageToken(groupsList.NextPageToken).Do()
+			UserKey(userKey).PageToken(groupsList.NextPageToken).Context(ctx).Do()
 		if err != nil {
 			return nil, fmt.Errorf("could not list groups: %v", err)
 		}

connector/google/google_test.go

@@ -233,9 +233,8 @@ func TestGetGroups(t *testing.T) {
 		callCounterMu.Unlock()
 		t.Run(name, func(t *testing.T) {
 			assert := assert.New(t)
-			lookup := make(map[string]struct{})
 
-			groups, err := conn.getGroups(testCase.userKey, testCase.fetchTransitiveGroupMembership, lookup)
+			groups, err := conn.getGroups(context.Background(), testCase.userKey, testCase.fetchTransitiveGroupMembership)
 			if testCase.shouldErr {
 				assert.NotNil(err)
 			} else {
@@ -294,9 +293,8 @@ func TestDomainToAdminEmailConfig(t *testing.T) {
 		callCounterMu.Unlock()
 		t.Run(name, func(t *testing.T) {
 			assert := assert.New(t)
-			lookup := make(map[string]struct{})
 
-			_, err := conn.getGroups(testCase.userKey, true, lookup)
+			_, err := conn.getGroups(context.Background(), testCase.userKey, true)
 			if testCase.expectedErr != "" {
 				assert.ErrorContains(err, testCase.expectedErr)
 			} else {
@@ -395,9 +393,8 @@ func TestGCEWorkloadIdentity(t *testing.T) {
 	} {
 		t.Run(name, func(t *testing.T) {
 			assert := assert.New(t)
-			lookup := make(map[string]struct{})
 
-			_, err := conn.getGroups(testCase.userKey, true, lookup)
+			_, err := conn.getGroups(context.Background(), testCase.userKey, true)
 			if testCase.expectedErr != "" {
 				assert.ErrorContains(err, testCase.expectedErr)
 			} else {

go.mod

@@ -42,6 +42,7 @@ require (
 	golang.org/x/exp v0.0.0-20240823005443-9b4947da3948
 	golang.org/x/net v0.53.0
 	golang.org/x/oauth2 v0.36.0
+	golang.org/x/sync v0.20.0
 	google.golang.org/api v0.277.0
 	google.golang.org/grpc v1.80.0
 	google.golang.org/protobuf v1.36.11
@@ -141,7 +142,6 @@ require (
 	go.uber.org/zap v1.27.0 // indirect
 	go.yaml.in/yaml/v2 v2.4.2 // indirect
 	golang.org/x/mod v0.34.0 // indirect
-	golang.org/x/sync v0.20.0 // indirect
 	golang.org/x/sys v0.43.0 // indirect
 	golang.org/x/text v0.36.0 // indirect
 	golang.org/x/time v0.15.0 // indirect