From 9657f148237ab1e08c73b141532dc57c68f341c7 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Mon, 17 Aug 2026 13:06:56 -0300 Subject: [PATCH 1/6] Fix event feed DeadlineExceeded from unbounded per-user Reports API fan-out usage_event_feed's OAuth lookup issued one sequential Reports API call per authorized app with no cap, so a batch of users with many authorized apps could force a single ListEvents call to make hundreds of rate-limited calls and blow the SDK's RPC deadline (observed in production). Add a per-call budget (maxLookupCallsPerEventFeedCall) shared across all three "last login" event feeds, with resumable per-user progress via cursor state, so a user needing more calls than the budget allows pauses mid-user instead of stalling or restarting from scratch. Also run a user's per-app lookups concurrently (bounded) so network latency overlaps instead of stacking on top of the rate-limiter wait. Co-Authored-By: Claude Sonnet 5 --- go.mod | 2 +- pkg/connector/event_feed_common.go | 78 +++++--- pkg/connector/event_feed_common_test.go | 168 +++++++++++++++++- pkg/connector/google_login_event_feed.go | 11 +- pkg/connector/saml_event_feed.go | 16 +- pkg/connector/usage_event_feed.go | 112 +++++++++--- vendor/golang.org/x/sync/errgroup/errgroup.go | 151 ++++++++++++++++ vendor/modules.txt | 1 + 8 files changed, 473 insertions(+), 66 deletions(-) create mode 100644 vendor/golang.org/x/sync/errgroup/errgroup.go diff --git a/go.mod b/go.mod index 27ceb3d4..56db3c45 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/stretchr/testify v1.11.1 go.uber.org/zap v1.28.0 golang.org/x/oauth2 v0.36.0 + golang.org/x/sync v0.22.0 google.golang.org/api v0.264.0 google.golang.org/grpc v1.83.0 google.golang.org/protobuf v1.36.11 @@ -137,7 +138,6 @@ require ( golang.org/x/crypto v0.54.0 // indirect golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 // indirect golang.org/x/net v0.57.0 // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/term v0.45.0 // indirect golang.org/x/text v0.40.0 // indirect diff --git a/pkg/connector/event_feed_common.go b/pkg/connector/event_feed_common.go index f16dc7fd..26f613fc 100644 --- a/pkg/connector/event_feed_common.go +++ b/pkg/connector/event_feed_common.go @@ -29,11 +29,16 @@ import ( gwclient "github.com/conductorone/baton-google-workspace/pkg/client" ) -// usersPerEventFeedCall bounds how many users are processed per ListEvents invocation, so a -// single call issues at most this many Reports API filter-queries and returns quickly instead -// of blocking on the shared 250/min quota for an entire directory page (up to 500 users). +// usersPerEventFeedCall bounds how many users are considered per ListEvents invocation. This +// alone does not bound the number of Reports API calls issued (see maxLookupCallsPerEventFeedCall +// below), but keeps the directory-side bookkeeping (ListTokens, cursor size) proportional. const usersPerEventFeedCall = 25 +// maxLookupCallsPerEventFeedCall caps total Reports API calls issued by one ListEvents +// invocation, so a call can't blow its RPC deadline (~16s of rate-limiter time at 220/min). +// A userEventLookup that needs more calls for one user returns a resumeState instead. +const maxLookupCallsPerEventFeedCall = 60 + type pendingUser struct { Email string `json:"email"` ID string `json:"id"` @@ -45,9 +50,13 @@ type pendingUser struct { // When both are empty/exhausted, the walk is complete and the cursor resets to nil so the // next call starts a fresh pass — there is no "since last poll" time window to track, since // each lookup always asks for the current latest login, not a delta. +// +// ResumeState is opaque, feed-defined progress for PendingUsers[0]'s own lookup (e.g. which +// authorized app to resume from), set when a call exhausts its budget mid-user. type userScanCursor struct { PendingUsers []pendingUser `json:"pending_users,omitempty"` DirectoryPageToken string `json:"directory_page_token,omitempty"` + ResumeState string `json:"resume_state,omitempty"` } func unmarshalUserScanCursor(pToken *pagination.StreamToken) (*userScanCursor, error) { @@ -84,8 +93,11 @@ func (c *userScanCursor) marshal() (string, error) { return base64.StdEncoding.EncodeToString(data), nil } -// userEventLookup fetches events for a single user via at most one Reports API call. -type userEventLookup func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) +// userEventLookup fetches events for one user, spending at most `budget` Reports API calls. +// resumeState picks up where a prior call for this user left off ("" = start fresh); if more +// than `budget` calls are needed, it returns a non-empty nextResumeState instead of finishing. +// consumed is how many calls this invocation issued (meaningful only when err == nil). +type userEventLookup func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int) (events []*v2.Event, nextResumeState string, consumed int, err error) // scanUsersForEvents drives one bounded step of the rolling user-directory walk shared by all // three "last login" event feeds. @@ -142,37 +154,63 @@ func scanUsersForEvents( batch = batch[:usersPerEventFeedCall] } + // finish marshals the cursor and returns the call's result; shared by every exit path. + finish := func(events []*v2.Event, hasMore bool) ([]*v2.Event, *pagination.StreamState, error) { + if !hasMore { + cursor = &userScanCursor{} + } + cursorToken, err := cursor.marshal() + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", err) + } + return events, &pagination.StreamState{Cursor: cursorToken, HasMore: hasMore}, nil + } + events := []*v2.Event{} - for _, u := range batch { - userEvents, err := lookup(ctx, client, u) + budget := maxLookupCallsPerEventFeedCall + resumeState := cursor.ResumeState + + for i, u := range batch { + if budget <= 0 { + // Budget spent before starting u; resume here fresh next call. + cursor.PendingUsers = cursor.PendingUsers[i:] + cursor.ResumeState = "" + return finish(events, true) + } + + userEvents, nextResume, consumed, err := lookup(ctx, client, u, resumeState, budget) if err != nil { - // Don't remove `batch` from cursor.PendingUsers until every lookup in it has - // succeeded, so a single user's Reports API blip doesn't lose the remaining - // unprocessed users and restart the walk from the beginning on retry. + // Keep u (and the rest) pending with its resume state, so a retry resumes here + // instead of restarting the whole batch. + cursor.PendingUsers = cursor.PendingUsers[i:] + cursor.ResumeState = resumeState cursorToken, marshalErr := cursor.marshal() if marshalErr != nil { return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", marshalErr) } return nil, &pagination.StreamState{Cursor: cursorToken, HasMore: true}, err } + budget -= consumed + resumeState = "" // only the batch's first (possibly-resumed) user carries one in + for _, e := range userEvents { if earliestEvent != nil && e.GetOccurredAt() != nil && e.GetOccurredAt().AsTime().Before(earliestEvent.AsTime()) { continue } events = append(events, e) } - } - cursor.PendingUsers = cursor.PendingUsers[len(batch):] - hasMore := len(cursor.PendingUsers) > 0 || cursor.DirectoryPageToken != "" - if !hasMore { - cursor = &userScanCursor{} + if nextResume != "" { + // u isn't finished; keep it at the front and stop the batch here. + cursor.PendingUsers = cursor.PendingUsers[i:] + cursor.ResumeState = nextResume + return finish(events, true) + } } - cursorToken, err := cursor.marshal() - if err != nil { - return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", err) - } + cursor.PendingUsers = cursor.PendingUsers[len(batch):] + cursor.ResumeState = "" - return events, &pagination.StreamState{Cursor: cursorToken, HasMore: hasMore}, nil + hasMore := len(cursor.PendingUsers) > 0 || cursor.DirectoryPageToken != "" + return finish(events, hasMore) } diff --git a/pkg/connector/event_feed_common_test.go b/pkg/connector/event_feed_common_test.go index ae7dc5f9..57fa3cf7 100644 --- a/pkg/connector/event_feed_common_test.go +++ b/pkg/connector/event_feed_common_test.go @@ -7,6 +7,8 @@ import ( "net/http" "net/http/httptest" "strconv" + "strings" + "sync" "testing" "time" @@ -128,9 +130,9 @@ func TestScanUsersForEvents_PaginatesAcrossMultipleCallsWithoutLoss(t *testing.T } _, state, err := scanUsersForEvents(context.Background(), client, "customer", "", nil, &pagination.StreamToken{Cursor: cursor}, - func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser) ([]*v2.Event, error) { + func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { visited[u.Email+":lookup"]++ - return nil, nil + return nil, "", 1, nil }) if err != nil { t.Fatalf("scanUsersForEvents: %v", err) @@ -163,11 +165,11 @@ func TestScanUsersForEvents_FiltersEventsBeforeEarliestEvent(t *testing.T) { now := time.Now().UTC() floor := timestamppb.New(now.Add(-1 * time.Hour)) - oldEvent := &v2.Event{Id: "old", OccurredAt: timestamppb.New(now.Add(-2 * time.Hour))} // before floor + oldEvent := &v2.Event{Id: "old", OccurredAt: timestamppb.New(now.Add(-2 * time.Hour))} // before floor newEvent := &v2.Event{Id: "new", OccurredAt: timestamppb.New(now.Add(-30 * time.Minute))} // after floor - lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser) ([]*v2.Event, error) { - return []*v2.Event{oldEvent, newEvent}, nil + lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { + return []*v2.Event{oldEvent, newEvent}, "", 1, nil } events, _, err := scanUsersForEvents(context.Background(), client, "customer", "", floor, &pagination.StreamToken{}, lookup) @@ -325,3 +327,159 @@ func TestUsageEventFeed_DedupesRepeatedClientIDsInTokens(t *testing.T) { t.Fatalf("expected exactly 1 event for the deduped app, got %d", len(events)) } } + +// TestScanUsersForEvents_ResumesWithinUserWhenBudgetExhausted verifies that a user needing more +// Reports calls than the budget allows pauses mid-user, resumes via cursor state across calls, +// and has every unit of work visited exactly once. Uses a synthetic lookup to isolate +// scanUsersForEvents' generic budget/resume mechanics from any one feed's fan-out logic. +func TestScanUsersForEvents_ResumesWithinUserWhenBudgetExhausted(t *testing.T) { + const totalUnits = maxLookupCallsPerEventFeedCall*2 + 7 // spans exactly 3 resumed calls + user := pendingUser{Email: "heavy@example.com", ID: "heavy-user"} + unitLookupCounts := map[int]int{} + + lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, resumeState string, budget int) ([]*v2.Event, string, int, error) { + startIdx := 0 + if resumeState != "" { + idx, err := strconv.Atoi(resumeState) + if err != nil { + t.Fatalf("bad resume state %q: %v", resumeState, err) + } + startIdx = idx + } + consumed := totalUnits - startIdx + if consumed > budget { + consumed = budget + } + events := make([]*v2.Event, 0, consumed) + for i := startIdx; i < startIdx+consumed; i++ { + unitLookupCounts[i]++ + events = append(events, &v2.Event{Id: strconv.Itoa(i)}) + } + nextResume := "" + if startIdx+consumed < totalUnits { + nextResume = strconv.Itoa(startIdx + consumed) + } + return events, nextResume, consumed, nil + } + + server := newDirectoryUsersOnlyServer(t, []*directoryAdmin.User{{Id: user.ID, PrimaryEmail: user.Email}}, 10, + func(string) *reportsAdmin.Activities { return &reportsAdmin.Activities{} }) + defer server.Close() + dir := newTestDirectoryService(t, server.URL, server.Client()) + client := &gwclient.GoogleWorkspaceClient{UserService: dir} + + var cursor string + var allEvents []*v2.Event + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + events, state, err := scanUsersForEvents(context.Background(), client, "customer", "", nil, &pagination.StreamToken{Cursor: cursor}, lookup) + if err != nil { + t.Fatalf("scanUsersForEvents: %v", err) + } + allEvents = append(allEvents, events...) + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 3 { + t.Fatalf("expected exactly 3 calls (budget=%d, units=%d), got %d", maxLookupCallsPerEventFeedCall, totalUnits, calls) + } + if len(allEvents) != totalUnits { + t.Fatalf("expected exactly %d events (one per unit), got %d", totalUnits, len(allEvents)) + } + for i := 0; i < totalUnits; i++ { + if unitLookupCounts[i] != 1 { + t.Fatalf("expected unit %d to be looked up exactly once, got %d", i, unitLookupCounts[i]) + } + } +} + +// TestUsageEventFeed_ResumesAcrossManyAuthorizedApps is the end-to-end version of the test above, +// exercising usage_event_feed's real lookupUser against a user with more authorized apps than +// maxLookupCallsPerEventFeedCall allows per call. +func TestUsageEventFeed_ResumesAcrossManyAuthorizedApps(t *testing.T) { + const userEmail = "heavy@example.com" + const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + tokens = append(tokens, &directoryAdmin.Token{ClientId: fmt.Sprintf("client-%d", i), DisplayText: fmt.Sprintf("App %d", i)}) + } + + var mu sync.Mutex + callCounts := map[string]int{} + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("pageToken") != "" { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{}) + return + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{ + Users: []*directoryAdmin.User{{Id: "profile-heavy", PrimaryEmail: userEmail}}, + }) + }) + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + + var cursor string + var allEvents []*v2.Event + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + events, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: cursor}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + allEvents = append(allEvents, events...) + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 2 { + t.Fatalf("expected exactly 2 ListEvents calls (budget=%d, apps=%d), got %d", maxLookupCallsPerEventFeedCall, numApps, calls) + } + if len(allEvents) != numApps { + t.Fatalf("expected %d events (one per app), got %d", numApps, len(allEvents)) + } + if len(callCounts) != numApps { + t.Fatalf("expected %d distinct apps queried, got %d", numApps, len(callCounts)) + } + for clientID, n := range callCounts { + if n != 1 { + t.Fatalf("expected exactly 1 Reports API call for %s, got %d", clientID, n) + } + } +} diff --git a/pkg/connector/google_login_event_feed.go b/pkg/connector/google_login_event_feed.go index d352abcc..fded81de 100644 --- a/pkg/connector/google_login_event_feed.go +++ b/pkg/connector/google_login_event_feed.go @@ -47,10 +47,11 @@ func (f *googleLoginEventFeed) EventFeedMetadata(_ context.Context) *v2.EventFee } } -func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) { +// lookupUser issues exactly one Reports API call per user, so resumeState/budget are unused. +func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { r, err := listActivitiesRateLimited(ctx, client, user.Email, reportsAppLogin, "login_success", "", "", googleLoginLookupMaxResults) if err != nil { - return nil, fmt.Errorf("google-workspace-connector: failed to list google login activities for %s: %w", user.Email, err) + return nil, "", 0, fmt.Errorf("google-workspace-connector: failed to list google login activities for %s: %w", user.Email, err) } var best *bestActivity @@ -67,14 +68,14 @@ func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient. } } if best == nil { - return nil, nil + return nil, "", 1, nil } userTrait, err := resource.NewUserTrait( resource.WithEmail(best.activity.Actor.Email, true), ) if err != nil { - return nil, fmt.Errorf("google-workspace-connector: failed to create user trait in google login event feed: %w", err) + return nil, "", 0, fmt.Errorf("google-workspace-connector: failed to create user trait in google login event feed: %w", err) } return []*v2.Event{{ @@ -100,7 +101,7 @@ func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient. }, }, }, - }}, nil + }}, "", 1, nil } func (f *googleLoginEventFeed) ListEvents( diff --git a/pkg/connector/saml_event_feed.go b/pkg/connector/saml_event_feed.go index ebae6761..31d03b51 100644 --- a/pkg/connector/saml_event_feed.go +++ b/pkg/connector/saml_event_feed.go @@ -48,11 +48,9 @@ type samlAppActivity struct { appName string } -// lookupUser tracks SAML app usage via Google's "saml" audit log. -// -// Unlike OAuth apps (see usage_event_feed.go), SAML "login_success" fires on every SSO -// authentication, so last login timestamps are accurate. SAML apps are identified by app name -// (no numeric client_id). +// lookupUser tracks SAML app usage via Google's "saml" audit log. Unlike OAuth apps (see +// usage_event_feed.go), one query covers every SAML app for a user, so it always completes in a +// single Reports API call and never needs resume state. func (f *samlEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, samlProfileMap map[string]string, user pendingUser) ([]*v2.Event, error) { r, err := listActivitiesRateLimited(ctx, client, user.Email, reportsAppSAML, "login_success", "", "", samlAppLookupMaxResults) if err != nil { @@ -134,8 +132,12 @@ func (f *samlEventFeed) ListEvents(ctx context.Context, earliestEvent *timestamp } events, streamState, err := scanUsersForEvents(ctx, f.client, f.customerID, f.domain, earliestEvent, pToken, - func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) { - return f.lookupUser(ctx, client, samlProfileMap, user) + func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { + events, err := f.lookupUser(ctx, client, samlProfileMap, user) + if err != nil { + return nil, "", 0, err + } + return events, "", 1, nil }) if err != nil { return nil, streamState, nil, err diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 674c5a35..434a7ee2 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -13,6 +13,8 @@ import ( "github.com/conductorone/baton-sdk/pkg/annotations" "github.com/conductorone/baton-sdk/pkg/pagination" "github.com/conductorone/baton-sdk/pkg/types/resource" + "golang.org/x/sync/errgroup" + directoryAdmin "google.golang.org/api/admin/directory/v1" reportsAdmin "google.golang.org/api/admin/reports/v1" "google.golang.org/api/googleapi" "google.golang.org/protobuf/types/known/timestamppb" @@ -29,6 +31,10 @@ var privateAppIDRegex = regexp.MustCompile("[0-9]{21}") // client-side regardless of Google's actual ordering. const oauthAppLookupMaxResults = 5 +// maxConcurrentAppLookups bounds concurrent per-app Reports lookups for one user. The shared +// reportsRateLimiter still caps overall quota use; this just overlaps network latency. +const maxConcurrentAppLookups = 8 + type usageEventFeed struct { c *gwclient.GoogleWorkspaceClient customerID string @@ -79,29 +85,17 @@ type oauthAppActivity struct { appName string } -// lookupUser enumerates the OAuth apps this user has authorized via Directory API Tokens.list -// (Reports API has no endpoint for "which apps has this user used" — Tokens.list is the cheap, -// Directory-quota source of truth for that), then issues one Reports API lookup per app, -// scoped with filters=client_id==. This trades "1 Reports call per user" (old per-app-type -// window, prone to one app's activity crowding another out of a shared count-bounded window) for -// "1 Directory call + N Reports calls per user" (N = apps that user has authorized) — each app's -// freshness is queried independently, so it can never be crowded out by another app's activity. -func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) { - tokenResp, err := client.ListTokens(ctx, user.ID) - if err != nil { - var gerr *googleapi.Error - if errors.As(err, &gerr) && gerr.Code == http.StatusNotFound { - // Benign: the user was deleted between the directory listing and this lookup. - return nil, nil - } - return nil, fmt.Errorf("google-workspace: failed to list oauth tokens for %s: %w", user.Email, err) - } - // Tokens.list can return multiple Token entries for the same client_id — e.g. a user - // granting a different scope set to the same app at different times. Dedupe here so each - // distinct app only ever triggers one Reports API lookup, never repeated ones. - seenClientIDs := make(map[string]struct{}, len(tokenResp.Items)) +// oauthApp is one distinct app a user has authorized, as discovered via Directory Tokens.list. +type oauthApp struct { + ClientID string + DisplayText string +} - events := make([]*v2.Event, 0, len(tokenResp.Items)) +// distinctAuthorizedApps enumerates this user's authorized OAuth apps via Directory Tokens.list, +// deduping repeated client_ids and filtering out private apps before spending any Reports budget. +func distinctAuthorizedApps(tokenResp *directoryAdmin.Tokens) []oauthApp { + seenClientIDs := make(map[string]struct{}, len(tokenResp.Items)) + apps := make([]oauthApp, 0, len(tokenResp.Items)) for _, t := range tokenResp.Items { if t.ClientId == "" || t.DisplayText == "" { continue @@ -114,17 +108,79 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google continue } seenClientIDs[t.ClientId] = struct{}{} + apps = append(apps, oauthApp{ClientID: t.ClientId, DisplayText: t.DisplayText}) + } + return apps +} + +// lookupUser issues one Reports API lookup per authorized OAuth app, scoped by client_id. A user +// can have more apps than the per-call budget allows: resumeState is the index into the app list +// to resume from, and if apps remain after spending budget, nextResumeState carries it forward. +// Per-app lookups run concurrently (bounded by maxConcurrentAppLookups); the shared rate limiter +// still caps quota use, so this only overlaps network latency. +func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int) ([]*v2.Event, string, int, error) { + tokenResp, err := client.ListTokens(ctx, user.ID) + if err != nil { + var gerr *googleapi.Error + if errors.As(err, &gerr) && gerr.Code == http.StatusNotFound { + // Benign: the user was deleted between the directory listing and this lookup. + return nil, "", 0, nil + } + return nil, "", 0, fmt.Errorf("google-workspace: failed to list oauth tokens for %s: %w", user.Email, err) + } + apps := distinctAuthorizedApps(tokenResp) - event, err := f.lookupAppLogin(ctx, client, user, t.ClientId, t.DisplayText) - if err != nil { - return nil, err + startIdx := 0 + if resumeState != "" { + idx, parseErr := strconv.Atoi(resumeState) + if parseErr != nil { + return nil, "", 0, fmt.Errorf("google-workspace: invalid resume state for %s: %w", user.Email, parseErr) } - if event != nil { - events = append(events, event) + startIdx = idx + } + if startIdx > len(apps) { + startIdx = len(apps) + } + + remaining := apps[startIdx:] + if len(remaining) == 0 { + return nil, "", 0, nil + } + toProcess := remaining + if len(toProcess) > budget { + toProcess = toProcess[:budget] + } + + results := make([]*v2.Event, len(toProcess)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(maxConcurrentAppLookups) + for i, app := range toProcess { + g.Go(func() error { + event, err := f.lookupAppLogin(gctx, client, user, app.ClientID, app.DisplayText) + if err != nil { + return err + } + results[i] = event + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, "", 0, err + } + + events := make([]*v2.Event, 0, len(results)) + for _, e := range results { + if e != nil { + events = append(events, e) } } - return events, nil + consumed := len(toProcess) + nextResume := "" + if startIdx+consumed < len(apps) { + nextResume = strconv.Itoa(startIdx + consumed) + } + return events, nextResume, consumed, nil } // lookupAppLogin fetches this user's most recent "authorize" activity for one specific OAuth diff --git a/vendor/golang.org/x/sync/errgroup/errgroup.go b/vendor/golang.org/x/sync/errgroup/errgroup.go new file mode 100644 index 00000000..c261a8eb --- /dev/null +++ b/vendor/golang.org/x/sync/errgroup/errgroup.go @@ -0,0 +1,151 @@ +// Copyright 2016 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package errgroup provides synchronization, error propagation, and Context +// cancellation for groups of goroutines working on subtasks of a common task. +// +// [errgroup.Group] is related to [sync.WaitGroup] but adds handling of tasks +// returning errors. +package errgroup + +import ( + "context" + "fmt" + "sync" +) + +type token struct{} + +// A Group is a collection of goroutines working on subtasks that are part of +// the same overall task. A Group should not be reused for different tasks. +// +// A zero Group is valid, has no limit on the number of active goroutines, +// and does not cancel on error. +type Group struct { + cancel func(error) + + wg sync.WaitGroup + + sem chan token + + errOnce sync.Once + err error +} + +func (g *Group) done() { + if g.sem != nil { + <-g.sem + } + g.wg.Done() +} + +// WithContext returns a new Group and an associated Context derived from ctx. +// +// The derived Context is canceled the first time a function passed to Go +// returns a non-nil error or the first time Wait returns, whichever occurs +// first. +func WithContext(ctx context.Context) (*Group, context.Context) { + ctx, cancel := context.WithCancelCause(ctx) + return &Group{cancel: cancel}, ctx +} + +// Wait blocks until all function calls from the Go method have returned, then +// returns the first non-nil error (if any) from them. +func (g *Group) Wait() error { + g.wg.Wait() + if g.cancel != nil { + g.cancel(g.err) + } + return g.err +} + +// Go calls the given function in a new goroutine. +// +// The first call to Go must happen before a Wait. +// It blocks until the new goroutine can be added without the number of +// goroutines in the group exceeding the configured limit. +// +// The first goroutine in the group that returns a non-nil error will +// cancel the associated Context, if any. The error will be returned +// by Wait. +func (g *Group) Go(f func() error) { + if g.sem != nil { + g.sem <- token{} + } + + g.wg.Add(1) + go func() { + defer g.done() + + // It is tempting to propagate panics from f() + // up to the goroutine that calls Wait, but + // it creates more problems than it solves: + // - it delays panics arbitrarily, + // making bugs harder to detect; + // - it turns f's panic stack into a mere value, + // hiding it from crash-monitoring tools; + // - it risks deadlocks that hide the panic entirely, + // if f's panic leaves the program in a state + // that prevents the Wait call from being reached. + // See #53757, #74275, #74304, #74306. + + if err := f(); err != nil { + g.errOnce.Do(func() { + g.err = err + if g.cancel != nil { + g.cancel(g.err) + } + }) + } + }() +} + +// TryGo calls the given function in a new goroutine only if the number of +// active goroutines in the group is currently below the configured limit. +// +// The return value reports whether the goroutine was started. +func (g *Group) TryGo(f func() error) bool { + if g.sem != nil { + select { + case g.sem <- token{}: + // Note: this allows barging if and only if channels in general allow barging. + default: + return false + } + } + + g.wg.Add(1) + go func() { + defer g.done() + + if err := f(); err != nil { + g.errOnce.Do(func() { + g.err = err + if g.cancel != nil { + g.cancel(g.err) + } + }) + } + }() + return true +} + +// SetLimit limits the number of active goroutines in this group to at most n. +// A negative value indicates no limit. +// A limit of zero will prevent any new goroutines from being added. +// +// Any subsequent call to the Go method will block until it can add an active +// goroutine without exceeding the configured limit. +// +// The limit must not be modified while any goroutines in the group are active. +func (g *Group) SetLimit(n int) { + if n < 0 { + g.sem = nil + return + } + if active := len(g.sem); active != 0 { + panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", active)) + } + g.sem = make(chan token, n) +} diff --git a/vendor/modules.txt b/vendor/modules.txt index 4008cfd3..7d75824f 100644 --- a/vendor/modules.txt +++ b/vendor/modules.txt @@ -854,6 +854,7 @@ golang.org/x/oauth2/jws golang.org/x/oauth2/jwt # golang.org/x/sync v0.22.0 ## explicit; go 1.25.0 +golang.org/x/sync/errgroup golang.org/x/sync/semaphore golang.org/x/sync/singleflight # golang.org/x/sys v0.47.0 From 1b920628e557093ca55a46abdc5c32dc86cae050 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Mon, 17 Aug 2026 14:47:27 -0300 Subject: [PATCH 2/6] Address PR review findings on event feed budget/resume logic - Clamp resume-state index to 0 to prevent a panic on a negative/corrupted cursor value in usage_event_feed's per-app resume logic. - Stop advancing the cursor in scanUsersForEvents' error branch: the SDK discards the returned StreamState whenever err != nil, so the prior advancement had no effect and only muddied the comment. - Add a wall-clock soft deadline (maxEventFeedCallDuration) alongside the call-count budget, since a call budget alone doesn't bound elapsed time under a shared, contended rate limiter with retry/backoff. - Make the heavy multi-app resume test use an unlimited rate limiter so it doesn't compete for or depend on the shared 220/min quota other tests use. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/event_feed_common.go | 22 +++++++++++++--------- pkg/connector/event_feed_common_test.go | 11 +++++++++++ pkg/connector/usage_event_feed.go | 3 +++ 3 files changed, 27 insertions(+), 9 deletions(-) diff --git a/pkg/connector/event_feed_common.go b/pkg/connector/event_feed_common.go index 26f613fc..b0442e25 100644 --- a/pkg/connector/event_feed_common.go +++ b/pkg/connector/event_feed_common.go @@ -19,6 +19,7 @@ import ( "encoding/base64" "encoding/json" "fmt" + "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/pagination" @@ -34,11 +35,14 @@ import ( // below), but keeps the directory-side bookkeeping (ListTokens, cursor size) proportional. const usersPerEventFeedCall = 25 -// maxLookupCallsPerEventFeedCall caps total Reports API calls issued by one ListEvents -// invocation, so a call can't blow its RPC deadline (~16s of rate-limiter time at 220/min). -// A userEventLookup that needs more calls for one user returns a resumeState instead. +// maxLookupCallsPerEventFeedCall caps Reports API calls per ListEvents invocation. Bounds cost, +// not time (shared rate limiter, retries with backoff) — see maxEventFeedCallDuration for that. const maxLookupCallsPerEventFeedCall = 60 +// maxEventFeedCallDuration is a soft wall-clock budget per ListEvents invocation, checked between +// users, since a call-count budget alone doesn't bound elapsed time. +const maxEventFeedCallDuration = 45 * time.Second + type pendingUser struct { Email string `json:"email"` ID string `json:"id"` @@ -169,10 +173,11 @@ func scanUsersForEvents( events := []*v2.Event{} budget := maxLookupCallsPerEventFeedCall resumeState := cursor.ResumeState + start := time.Now() for i, u := range batch { - if budget <= 0 { - // Budget spent before starting u; resume here fresh next call. + if budget <= 0 || time.Since(start) >= maxEventFeedCallDuration { + // Budget or time spent before starting u; resume here fresh next call. cursor.PendingUsers = cursor.PendingUsers[i:] cursor.ResumeState = "" return finish(events, true) @@ -180,10 +185,9 @@ func scanUsersForEvents( userEvents, nextResume, consumed, err := lookup(ctx, client, u, resumeState, budget) if err != nil { - // Keep u (and the rest) pending with its resume state, so a retry resumes here - // instead of restarting the whole batch. - cursor.PendingUsers = cursor.PendingUsers[i:] - cursor.ResumeState = resumeState + // The SDK drops this StreamState on error, so a retry replays the last successful + // cursor regardless; return it unchanged (it's untouched at this point) rather than + // advancing it for no effect. cursorToken, marshalErr := cursor.marshal() if marshalErr != nil { return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", marshalErr) diff --git a/pkg/connector/event_feed_common_test.go b/pkg/connector/event_feed_common_test.go index 57fa3cf7..c84ac477 100644 --- a/pkg/connector/event_feed_common_test.go +++ b/pkg/connector/event_feed_common_test.go @@ -400,10 +400,21 @@ func TestScanUsersForEvents_ResumesWithinUserWhenBudgetExhausted(t *testing.T) { } } +// withUnlimitedReportsRateLimiter swaps sharedReportsRateLimiter for a large-capacity one for the +// test's duration, so tests issuing many real Reports calls aren't order-dependent on other tests. +func withUnlimitedReportsRateLimiter(t *testing.T) { + t.Helper() + original := sharedReportsRateLimiter + sharedReportsRateLimiter = newReportsRateLimiter(1_000_000) + t.Cleanup(func() { sharedReportsRateLimiter = original }) +} + // TestUsageEventFeed_ResumesAcrossManyAuthorizedApps is the end-to-end version of the test above, // exercising usage_event_feed's real lookupUser against a user with more authorized apps than // maxLookupCallsPerEventFeedCall allows per call. func TestUsageEventFeed_ResumesAcrossManyAuthorizedApps(t *testing.T) { + withUnlimitedReportsRateLimiter(t) + const userEmail = "heavy@example.com" const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 434a7ee2..8b18baab 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -138,6 +138,9 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google } startIdx = idx } + if startIdx < 0 { + startIdx = 0 + } if startIdx > len(apps) { startIdx = len(apps) } From 54e0bceeb30574d0ebec7a2ccc50f7968eb17369 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Mon, 17 Aug 2026 14:56:39 -0300 Subject: [PATCH 3/6] Resume usage event feed's app lookup by client_id, not position Tokens.list is re-fetched fresh on every resumed call and its ordering isn't guaranteed stable, so a positional resume index could silently skip or re-visit apps whenever a user's authorized-app set changed between calls (e.g. an app revoked mid-pagination shifts every later app's index down by one). Resume by the last-processed app's client_id instead, re-locating it in the fresh list each time. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/event_feed_common_test.go | 99 +++++++++++++++++++++++++ pkg/connector/usage_event_feed.go | 28 +++---- 2 files changed, 114 insertions(+), 13 deletions(-) diff --git a/pkg/connector/event_feed_common_test.go b/pkg/connector/event_feed_common_test.go index c84ac477..85e4c50c 100644 --- a/pkg/connector/event_feed_common_test.go +++ b/pkg/connector/event_feed_common_test.go @@ -494,3 +494,102 @@ func TestUsageEventFeed_ResumesAcrossManyAuthorizedApps(t *testing.T) { } } } + +// TestUsageEventFeed_ResumeSurvivesAppListShift verifies that resuming a user's authorized-app +// lookup by client_id (not by positional index) tolerates the app list shifting between the +// first and second ListEvents call. A revoked app that was already processed disappears from +// Tokens.list, shifting every later app's position down by one; a positional-index resume would +// then skip whatever app lands at the old index in the shifted list, silently losing an event. +// Resuming by client_id instead re-locates the last-processed app in the fresh list (wherever it +// now sits) and continues right after it, so every not-yet-processed app is still visited exactly +// once. +func TestUsageEventFeed_ResumeSurvivesAppListShift(t *testing.T) { + withUnlimitedReportsRateLimiter(t) + + const userEmail = "heavy@example.com" + const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls + const revokedAfterFirstCall = "client-5" // processed in call 1, then revoked + + var mu sync.Mutex + callCounts := map[string]int{} + tokensCalls := 0 + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("pageToken") != "" { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{}) + return + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{ + Users: []*directoryAdmin.User{{Id: "profile-heavy", PrimaryEmail: userEmail}}, + }) + }) + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + tokensCalls++ + shifted := tokensCalls > 1 + mu.Unlock() + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if shifted && clientID == revokedAfterFirstCall { + // Revoked between calls: every later app's index shifts down by one. + continue + } + tokens = append(tokens, &directoryAdmin.Token{ClientId: clientID, DisplayText: clientID}) + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + + var cursor string + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + _, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: cursor}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 2 { + t.Fatalf("expected exactly 2 ListEvents calls, got %d", calls) + } + // Every app must be queried exactly once, including the ones after the revoked app whose + // position shifted. A positional-index resume would skip one of them here. + if len(callCounts) != numApps { + t.Fatalf("expected %d distinct apps queried, got %d: %v", numApps, len(callCounts), callCounts) + } + for clientID, n := range callCounts { + if n != 1 { + t.Fatalf("expected exactly 1 Reports API call for %s, got %d", clientID, n) + } + } +} diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 8b18baab..23e7e90e 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -114,8 +114,15 @@ func distinctAuthorizedApps(tokenResp *directoryAdmin.Tokens) []oauthApp { } // lookupUser issues one Reports API lookup per authorized OAuth app, scoped by client_id. A user -// can have more apps than the per-call budget allows: resumeState is the index into the app list -// to resume from, and if apps remain after spending budget, nextResumeState carries it forward. +// can have more apps than the per-call budget allows: resumeState is the client_id of the last +// app processed on a prior call for this user, and if apps remain after spending budget, +// nextResumeState carries the new last-processed client_id forward. A client_id (rather than a +// positional index) is used because Tokens.list is re-fetched fresh on every call and its +// ordering is not guaranteed stable across calls — a positional index could silently skip or +// re-visit apps if the set of authorized apps changes between calls. If resumeState's client_id +// is no longer present (e.g. the app was deauthorized between calls), lookupUser conservatively +// restarts from the beginning rather than guessing a position, which can revisit already-seen +// apps (harmless — the lookup is idempotent) but never skips one. // Per-app lookups run concurrently (bounded by maxConcurrentAppLookups); the shared rate limiter // still caps quota use, so this only overlaps network latency. func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int) ([]*v2.Event, string, int, error) { @@ -132,17 +139,12 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google startIdx := 0 if resumeState != "" { - idx, parseErr := strconv.Atoi(resumeState) - if parseErr != nil { - return nil, "", 0, fmt.Errorf("google-workspace: invalid resume state for %s: %w", user.Email, parseErr) + for i, app := range apps { + if app.ClientID == resumeState { + startIdx = i + 1 + break + } } - startIdx = idx - } - if startIdx < 0 { - startIdx = 0 - } - if startIdx > len(apps) { - startIdx = len(apps) } remaining := apps[startIdx:] @@ -181,7 +183,7 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google consumed := len(toProcess) nextResume := "" if startIdx+consumed < len(apps) { - nextResume = strconv.Itoa(startIdx + consumed) + nextResume = toProcess[len(toProcess)-1].ClientID } return events, nextResume, consumed, nil } From 6f27963afb0a5bb363d142ad54404fa6146e26d8 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Mon, 17 Aug 2026 19:03:50 -0300 Subject: [PATCH 4/6] Bound event feed's per-user fan-out by wall-clock deadline, not just budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit maxEventFeedCallDuration was only checked between users in scanUsersForEvents, never during a single user's own lookup. Since usageEventFeed.lookupUser could fan out a full budget's worth of Reports calls (up to 60) in one go, each retrying up to reportsMaxRetries times against a rate limiter shared across all three event feeds, one heavy contended user could still blow well past the intended wall-clock budget before the check ever ran again — reproducing the DeadlineExceeded this branch exists to fix. Thread a deadline (derived from the soft budget, tightened against ctx's own deadline when set) down into each feed's lookup. usageEventFeed's fan-out now runs in fixed-size chunks and checks the deadline between chunks (never before the first, so a call always makes progress), stopping early and resuming from the last app actually finished once time is up. The two single-call feeds (google login, SAML) just accept and ignore the new parameter. Also reworded a misleading comment: the error-path cursor in scanUsersForEvents isn't "untouched" when a fresh directory page was already fetched earlier in the same call — it's simply moot, since the SDK discards the returned StreamState on error. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/event_feed_common.go | 43 ++++- pkg/connector/event_feed_common_test.go | 193 ++++++++++++++++++++++- pkg/connector/google_login_event_feed.go | 6 +- pkg/connector/saml_event_feed.go | 3 +- pkg/connector/usage_event_feed.go | 75 +++++---- 5 files changed, 279 insertions(+), 41 deletions(-) diff --git a/pkg/connector/event_feed_common.go b/pkg/connector/event_feed_common.go index b0442e25..0f2a1b46 100644 --- a/pkg/connector/event_feed_common.go +++ b/pkg/connector/event_feed_common.go @@ -39,10 +39,19 @@ const usersPerEventFeedCall = 25 // not time (shared rate limiter, retries with backoff) — see maxEventFeedCallDuration for that. const maxLookupCallsPerEventFeedCall = 60 -// maxEventFeedCallDuration is a soft wall-clock budget per ListEvents invocation, checked between -// users, since a call-count budget alone doesn't bound elapsed time. +// maxEventFeedCallDuration is a soft wall-clock budget per ListEvents invocation. It bounds the +// gap between users (see the loop in scanUsersForEvents) and is also handed to each user's own +// lookup as a deadline, so a single user's internal fan-out (e.g. usageEventFeed.lookupUser's +// per-app Reports calls) stops starting new work once it's spent, instead of only being checked +// again after that one user's lookup already returned. const maxEventFeedCallDuration = 45 * time.Second +// eventFeedDeadlineSafetyMargin is subtracted from ctx's deadline (when the caller set one) to +// derive the effective per-call deadline passed to lookup, so this connector stops issuing new +// Reports calls with enough margin to unwind and return a partial result before the RPC itself +// times out — rather than getting DeadlineExceeded mid-call with nothing to show for it. +const eventFeedDeadlineSafetyMargin = 5 * time.Second + type pendingUser struct { Email string `json:"email"` ID string `json:"id"` @@ -100,8 +109,13 @@ func (c *userScanCursor) marshal() (string, error) { // userEventLookup fetches events for one user, spending at most `budget` Reports API calls. // resumeState picks up where a prior call for this user left off ("" = start fresh); if more // than `budget` calls are needed, it returns a non-empty nextResumeState instead of finishing. -// consumed is how many calls this invocation issued (meaningful only when err == nil). -type userEventLookup func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int) (events []*v2.Event, nextResumeState string, consumed int, err error) +// consumed is how many calls this invocation issued (meaningful only when err == nil). deadline +// is a soft wall-clock cutoff: a lookup with internal fan-out (multiple Reports calls per user) +// should stop starting new calls once it's past deadline and return whatever it has, with +// nextResumeState reflecting the last unit of work it actually finished — never a zero-value +// time.Time in practice, but implementations that only ever issue one call per user (there is no +// fan-out to bound) are free to ignore it. +type userEventLookup func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int, deadline time.Time) (events []*v2.Event, nextResumeState string, consumed int, err error) // scanUsersForEvents drives one bounded step of the rolling user-directory walk shared by all // three "last login" event feeds. @@ -175,6 +189,17 @@ func scanUsersForEvents( resumeState := cursor.ResumeState start := time.Now() + // deadline is handed to each user's lookup so a single heavy user's internal fan-out (e.g. + // many Reports calls for one user's authorized apps) also respects the wall-clock budget, + // not just the gap between users below. If the caller's ctx carries its own deadline (the + // real RPC deadline), stop earlier than that with a safety margin instead of racing it. + deadline := start.Add(maxEventFeedCallDuration) + if ctxDeadline, ok := ctx.Deadline(); ok { + if safe := ctxDeadline.Add(-eventFeedDeadlineSafetyMargin); safe.Before(deadline) { + deadline = safe + } + } + for i, u := range batch { if budget <= 0 || time.Since(start) >= maxEventFeedCallDuration { // Budget or time spent before starting u; resume here fresh next call. @@ -183,11 +208,13 @@ func scanUsersForEvents( return finish(events, true) } - userEvents, nextResume, consumed, err := lookup(ctx, client, u, resumeState, budget) + userEvents, nextResume, consumed, err := lookup(ctx, client, u, resumeState, budget, deadline) if err != nil { - // The SDK drops this StreamState on error, so a retry replays the last successful - // cursor regardless; return it unchanged (it's untouched at this point) rather than - // advancing it for no effect. + // The SDK drops this StreamState on error (connectorbuilder.ListEvents returns nil, + // err without ever reading it), so what we marshal here is moot either way; return + // cursor as it currently stands — which may already include a fresh directory-page + // advance from the len(cursor.PendingUsers) == 0 branch above — rather than trying to + // preserve or roll back per-user progress for no effect. cursorToken, marshalErr := cursor.marshal() if marshalErr != nil { return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", marshalErr) diff --git a/pkg/connector/event_feed_common_test.go b/pkg/connector/event_feed_common_test.go index 85e4c50c..b9b55f3e 100644 --- a/pkg/connector/event_feed_common_test.go +++ b/pkg/connector/event_feed_common_test.go @@ -130,7 +130,7 @@ func TestScanUsersForEvents_PaginatesAcrossMultipleCallsWithoutLoss(t *testing.T } _, state, err := scanUsersForEvents(context.Background(), client, "customer", "", nil, &pagination.StreamToken{Cursor: cursor}, - func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { + func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { visited[u.Email+":lookup"]++ return nil, "", 1, nil }) @@ -168,7 +168,7 @@ func TestScanUsersForEvents_FiltersEventsBeforeEarliestEvent(t *testing.T) { oldEvent := &v2.Event{Id: "old", OccurredAt: timestamppb.New(now.Add(-2 * time.Hour))} // before floor newEvent := &v2.Event{Id: "new", OccurredAt: timestamppb.New(now.Add(-30 * time.Minute))} // after floor - lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { + lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { return []*v2.Event{oldEvent, newEvent}, "", 1, nil } @@ -270,6 +270,80 @@ func TestUsageEventFeed_PicksLatestPerAppAndFiltersPrivateApps(t *testing.T) { } } +// TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks verifies that lookupUser's per-app +// fan-out (usage_event_feed.go) respects the deadline passed in from scanUsersForEvents: once +// past deadline, it stops launching new chunks of maxConcurrentAppLookups apps rather than +// draining the whole remaining budget in one call, but it always completes at least the first +// chunk (the deadline is only checked *between* chunks) so a single call still makes forward +// progress even if the deadline was already past when the call started. +func TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks(t *testing.T) { + const userEmail = "heavy@example.com" + const numApps = maxConcurrentAppLookups*2 + 3 // more than one chunk's worth + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + tokens = append(tokens, &directoryAdmin.Token{ClientId: fmt.Sprintf("client-%d", i), DisplayText: fmt.Sprintf("App %d", i)}) + } + + var mu sync.Mutex + callCounts := map[string]int{} + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + user := pendingUser{Email: userEmail, ID: "heavy-user"} + + // A deadline already in the past: only the check between chunks sees it, so the first chunk + // (maxConcurrentAppLookups apps) still runs to completion before lookupUser bails out. + events, nextResume, consumed, err := feed.lookupUser(context.Background(), feed.c, user, "", numApps, time.Now().Add(-time.Hour)) + if err != nil { + t.Fatalf("lookupUser: %v", err) + } + if consumed != maxConcurrentAppLookups { + t.Fatalf("expected exactly one chunk (%d apps) to be consumed before bailing out on the past deadline, got %d", maxConcurrentAppLookups, consumed) + } + if len(events) != maxConcurrentAppLookups { + t.Fatalf("expected exactly %d events (one per app in the completed chunk), got %d", maxConcurrentAppLookups, len(events)) + } + wantNextResume := fmt.Sprintf("client-%d", maxConcurrentAppLookups-1) + if nextResume != wantNextResume { + t.Fatalf("expected nextResume to be the last app finished in the completed chunk (%s), got %q", wantNextResume, nextResume) + } + for i := 0; i < maxConcurrentAppLookups; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] != 1 { + t.Fatalf("expected app %s (in the completed chunk) to be queried exactly once, got %d", clientID, callCounts[clientID]) + } + } + for i := maxConcurrentAppLookups; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] != 0 { + t.Fatalf("expected app %s (beyond the first chunk) not to be queried once the deadline had passed, got %d", clientID, callCounts[clientID]) + } + } +} + // TestUsageEventFeed_DedupesRepeatedClientIDsInTokens verifies that when Tokens.list returns // multiple Token entries for the same client_id (e.g. separate grants for different scope // sets), the feed issues exactly one Reports API lookup for that app and emits exactly one @@ -337,7 +411,7 @@ func TestScanUsersForEvents_ResumesWithinUserWhenBudgetExhausted(t *testing.T) { user := pendingUser{Email: "heavy@example.com", ID: "heavy-user"} unitLookupCounts := map[int]int{} - lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, resumeState string, budget int) ([]*v2.Event, string, int, error) { + lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, resumeState string, budget int, _ time.Time) ([]*v2.Event, string, int, error) { startIdx := 0 if resumeState != "" { idx, err := strconv.Atoi(resumeState) @@ -593,3 +667,116 @@ func TestUsageEventFeed_ResumeSurvivesAppListShift(t *testing.T) { } } } + +// TestUsageEventFeed_ResumeRestartsWhenAnchorAppRevoked locks in lookupUser's "anchor not +// found" fallback (usage_event_feed.go): if the exact app the resume cursor anchors on — the +// last one processed in a prior call — is itself revoked before the next resumed call, the +// fresh app list no longer contains it, so lookupUser can't re-locate a resume position and +// instead restarts this user's app scan from the beginning. This is intentionally conservative +// (it never skips an app), but it isn't free: every app already processed before the anchor gets +// looked up — and its event re-emitted — a second time. TestUsageEventFeed_ResumeSurvivesAppListShift +// revokes an app *before* the anchor, which leaves the anchor itself findable and never exercises +// this restart path. +func TestUsageEventFeed_ResumeRestartsWhenAnchorAppRevoked(t *testing.T) { + withUnlimitedReportsRateLimiter(t) + + const userEmail = "heavy@example.com" + const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls + anchorClientID := fmt.Sprintf("client-%d", maxLookupCallsPerEventFeedCall-1) + + var mu sync.Mutex + callCounts := map[string]int{} + tokensCalls := 0 + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("pageToken") != "" { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{}) + return + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{ + Users: []*directoryAdmin.User{{Id: "profile-heavy", PrimaryEmail: userEmail}}, + }) + }) + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + tokensCalls++ + anchorRevoked := tokensCalls > 1 + mu.Unlock() + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if anchorRevoked && clientID == anchorClientID { + // The exact app the resume cursor anchors on is revoked before the next call. + continue + } + tokens = append(tokens, &directoryAdmin.Token{ClientId: clientID, DisplayText: clientID}) + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + + var cursor string + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + _, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: cursor}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 3 { + t.Fatalf("expected exactly 3 ListEvents calls (call 1, the restart, and the tail), got %d", calls) + } + + // The revoked anchor app was queried exactly once, in call 1, before it disappeared. + if callCounts[anchorClientID] != 1 { + t.Fatalf("expected the revoked anchor app %s to be queried exactly once, got %d", anchorClientID, callCounts[anchorClientID]) + } + + // Every app processed before the anchor in call 1 gets replayed once the restart-from-0 + // fallback kicks in — the bounded, non-skipping cost of the conservative fallback. + for i := 0; i < maxLookupCallsPerEventFeedCall-1; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] != 2 { + t.Fatalf("expected app %s (processed before the revoked anchor) to be replayed once after the restart, got %d calls", clientID, callCounts[clientID]) + } + } + + // No app is ever skipped: every app still present after the revocation is queried at least + // once. + for i := maxLookupCallsPerEventFeedCall; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] < 1 { + t.Fatalf("expected app %s to be queried at least once, got %d", clientID, callCounts[clientID]) + } + } +} diff --git a/pkg/connector/google_login_event_feed.go b/pkg/connector/google_login_event_feed.go index fded81de..76b13780 100644 --- a/pkg/connector/google_login_event_feed.go +++ b/pkg/connector/google_login_event_feed.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strconv" + "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/annotations" @@ -47,8 +48,9 @@ func (f *googleLoginEventFeed) EventFeedMetadata(_ context.Context) *v2.EventFee } } -// lookupUser issues exactly one Reports API call per user, so resumeState/budget are unused. -func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { +// lookupUser issues exactly one Reports API call per user, so resumeState/budget/deadline are +// unused: there's no internal fan-out here to bound. +func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { r, err := listActivitiesRateLimited(ctx, client, user.Email, reportsAppLogin, "login_success", "", "", googleLoginLookupMaxResults) if err != nil { return nil, "", 0, fmt.Errorf("google-workspace-connector: failed to list google login activities for %s: %w", user.Email, err) diff --git a/pkg/connector/saml_event_feed.go b/pkg/connector/saml_event_feed.go index 31d03b51..cc13f180 100644 --- a/pkg/connector/saml_event_feed.go +++ b/pkg/connector/saml_event_feed.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strconv" + "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/annotations" @@ -132,7 +133,7 @@ func (f *samlEventFeed) ListEvents(ctx context.Context, earliestEvent *timestamp } events, streamState, err := scanUsersForEvents(ctx, f.client, f.customerID, f.domain, earliestEvent, pToken, - func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int) ([]*v2.Event, string, int, error) { + func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { events, err := f.lookupUser(ctx, client, samlProfileMap, user) if err != nil { return nil, "", 0, err diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 23e7e90e..5144f732 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -31,8 +31,9 @@ var privateAppIDRegex = regexp.MustCompile("[0-9]{21}") // client-side regardless of Google's actual ordering. const oauthAppLookupMaxResults = 5 -// maxConcurrentAppLookups bounds concurrent per-app Reports lookups for one user. The shared -// reportsRateLimiter still caps overall quota use; this just overlaps network latency. +// maxConcurrentAppLookups bounds concurrent per-app Reports lookups for one user, and also sizes +// the fan-out chunks in lookupUser (see below) between which the deadline budget is checked. The +// shared reportsRateLimiter still caps overall quota use; this just overlaps network latency. const maxConcurrentAppLookups = 8 type usageEventFeed struct { @@ -123,9 +124,15 @@ func distinctAuthorizedApps(tokenResp *directoryAdmin.Tokens) []oauthApp { // is no longer present (e.g. the app was deauthorized between calls), lookupUser conservatively // restarts from the beginning rather than guessing a position, which can revisit already-seen // apps (harmless — the lookup is idempotent) but never skips one. -// Per-app lookups run concurrently (bounded by maxConcurrentAppLookups); the shared rate limiter -// still caps quota use, so this only overlaps network latency. -func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int) ([]*v2.Event, string, int, error) { +// +// Per-app lookups run concurrently in fixed-size chunks (maxConcurrentAppLookups); the shared +// rate limiter still caps quota use, so this only overlaps network latency and retry/backoff +// time. deadline is checked between chunks (never before the first one, so a call always makes +// some progress) — once past it, lookupUser stops starting new chunks and returns what it has +// with nextResumeState set to the last app actually finished. Without this, a single user with a +// full budget's worth of apps could keep launching chunks — each with its own rate-limiter wait +// and up to reportsMaxRetries backoff — well past the caller's soft wall-clock budget. +func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int, deadline time.Time) ([]*v2.Event, string, int, error) { tokenResp, err := client.ListTokens(ctx, user.ID) if err != nil { var gerr *googleapi.Error @@ -156,34 +163,48 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google toProcess = toProcess[:budget] } - results := make([]*v2.Event, len(toProcess)) - g, gctx := errgroup.WithContext(ctx) - g.SetLimit(maxConcurrentAppLookups) - for i, app := range toProcess { - g.Go(func() error { - event, err := f.lookupAppLogin(gctx, client, user, app.ClientID, app.DisplayText) - if err != nil { - return err - } - results[i] = event - return nil - }) - } - if err := g.Wait(); err != nil { - return nil, "", 0, err - } + events := make([]*v2.Event, 0, len(toProcess)) + consumed := 0 + for chunkStart := 0; chunkStart < len(toProcess); chunkStart += maxConcurrentAppLookups { + if chunkStart > 0 && !deadline.IsZero() && time.Now().After(deadline) { + // Past budget: stop starting new chunks. Whatever completed in prior chunks stands; + // nextResume (below) picks up right after the last app actually processed. + break + } + chunkEnd := chunkStart + maxConcurrentAppLookups + if chunkEnd > len(toProcess) { + chunkEnd = len(toProcess) + } + chunk := toProcess[chunkStart:chunkEnd] - events := make([]*v2.Event, 0, len(results)) - for _, e := range results { - if e != nil { - events = append(events, e) + results := make([]*v2.Event, len(chunk)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(maxConcurrentAppLookups) + for i, app := range chunk { + g.Go(func() error { + event, err := f.lookupAppLogin(gctx, client, user, app.ClientID, app.DisplayText) + if err != nil { + return err + } + results[i] = event + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, "", 0, err + } + + for _, e := range results { + if e != nil { + events = append(events, e) + } } + consumed += len(chunk) } - consumed := len(toProcess) nextResume := "" if startIdx+consumed < len(apps) { - nextResume = toProcess[len(toProcess)-1].ClientID + nextResume = toProcess[consumed-1].ClientID } return events, nextResume, consumed, nil } From aed8ca31777be265d7975376a71cd2d5240699b9 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Tue, 18 Aug 2026 13:18:06 -0300 Subject: [PATCH 5/6] Address PR review findings on event feed deadline handling - scanUsersForEvents now checks the ctx-derived deadline instead of the fixed maxEventFeedCallDuration, so a short RPC deadline is respected uniformly by the SAML and Google-login feeds (which never used the deadline internally), not just the usage feed's per-user fan-out. - usageEventFeed.lookupUser drops the fixed-size chunk barrier in favor of a single continuous errgroup, checking the deadline before dispatching each app rather than between chunks, so a slow/rate-limited app no longer idles the other concurrency slots for a whole chunk. - Guard against indexing toProcess[consumed-1] when consumed == 0. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/event_feed_common.go | 2 +- pkg/connector/event_feed_common_test.go | 33 ++++++------ pkg/connector/usage_event_feed.go | 71 +++++++++++-------------- 3 files changed, 48 insertions(+), 58 deletions(-) diff --git a/pkg/connector/event_feed_common.go b/pkg/connector/event_feed_common.go index 0f2a1b46..9427d753 100644 --- a/pkg/connector/event_feed_common.go +++ b/pkg/connector/event_feed_common.go @@ -201,7 +201,7 @@ func scanUsersForEvents( } for i, u := range batch { - if budget <= 0 || time.Since(start) >= maxEventFeedCallDuration { + if budget <= 0 || !time.Now().Before(deadline) { // Budget or time spent before starting u; resume here fresh next call. cursor.PendingUsers = cursor.PendingUsers[i:] cursor.ResumeState = "" diff --git a/pkg/connector/event_feed_common_test.go b/pkg/connector/event_feed_common_test.go index b9b55f3e..56d5695d 100644 --- a/pkg/connector/event_feed_common_test.go +++ b/pkg/connector/event_feed_common_test.go @@ -272,9 +272,9 @@ func TestUsageEventFeed_PicksLatestPerAppAndFiltersPrivateApps(t *testing.T) { // TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks verifies that lookupUser's per-app // fan-out (usage_event_feed.go) respects the deadline passed in from scanUsersForEvents: once -// past deadline, it stops launching new chunks of maxConcurrentAppLookups apps rather than -// draining the whole remaining budget in one call, but it always completes at least the first -// chunk (the deadline is only checked *between* chunks) so a single call still makes forward +// past deadline, it stops dispatching new app lookups rather than draining the whole remaining +// budget in one call, but it always dispatches at least the first app (the deadline is only +// checked *before dispatching the second and later apps*) so a single call still makes forward // progress even if the deadline was already past when the call started. func TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks(t *testing.T) { const userEmail = "heavy@example.com" @@ -314,32 +314,29 @@ func TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks(t *testing.T) { feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") user := pendingUser{Email: userEmail, ID: "heavy-user"} - // A deadline already in the past: only the check between chunks sees it, so the first chunk - // (maxConcurrentAppLookups apps) still runs to completion before lookupUser bails out. + // A deadline already in the past: only the check before dispatching the second (and later) + // app sees it, so the first app is still dispatched and awaited before lookupUser bails out. events, nextResume, consumed, err := feed.lookupUser(context.Background(), feed.c, user, "", numApps, time.Now().Add(-time.Hour)) if err != nil { t.Fatalf("lookupUser: %v", err) } - if consumed != maxConcurrentAppLookups { - t.Fatalf("expected exactly one chunk (%d apps) to be consumed before bailing out on the past deadline, got %d", maxConcurrentAppLookups, consumed) + if consumed != 1 { + t.Fatalf("expected exactly one app to be consumed before bailing out on the past deadline, got %d", consumed) } - if len(events) != maxConcurrentAppLookups { - t.Fatalf("expected exactly %d events (one per app in the completed chunk), got %d", maxConcurrentAppLookups, len(events)) + if len(events) != 1 { + t.Fatalf("expected exactly one event (for the one dispatched app), got %d", len(events)) } - wantNextResume := fmt.Sprintf("client-%d", maxConcurrentAppLookups-1) + wantNextResume := "client-0" if nextResume != wantNextResume { - t.Fatalf("expected nextResume to be the last app finished in the completed chunk (%s), got %q", wantNextResume, nextResume) + t.Fatalf("expected nextResume to be the one app dispatched (%s), got %q", wantNextResume, nextResume) } - for i := 0; i < maxConcurrentAppLookups; i++ { - clientID := fmt.Sprintf("client-%d", i) - if callCounts[clientID] != 1 { - t.Fatalf("expected app %s (in the completed chunk) to be queried exactly once, got %d", clientID, callCounts[clientID]) - } + if callCounts["client-0"] != 1 { + t.Fatalf("expected app client-0 to be queried exactly once, got %d", callCounts["client-0"]) } - for i := maxConcurrentAppLookups; i < numApps; i++ { + for i := 1; i < numApps; i++ { clientID := fmt.Sprintf("client-%d", i) if callCounts[clientID] != 0 { - t.Fatalf("expected app %s (beyond the first chunk) not to be queried once the deadline had passed, got %d", clientID, callCounts[clientID]) + t.Fatalf("expected app %s not to be queried once the deadline had passed, got %d", clientID, callCounts[clientID]) } } } diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 5144f732..5f0d2aba 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -125,13 +125,14 @@ func distinctAuthorizedApps(tokenResp *directoryAdmin.Tokens) []oauthApp { // restarts from the beginning rather than guessing a position, which can revisit already-seen // apps (harmless — the lookup is idempotent) but never skips one. // -// Per-app lookups run concurrently in fixed-size chunks (maxConcurrentAppLookups); the shared +// Per-app lookups run concurrently, up to maxConcurrentAppLookups in flight at once; the shared // rate limiter still caps quota use, so this only overlaps network latency and retry/backoff -// time. deadline is checked between chunks (never before the first one, so a call always makes -// some progress) — once past it, lookupUser stops starting new chunks and returns what it has -// with nextResumeState set to the last app actually finished. Without this, a single user with a -// full budget's worth of apps could keep launching chunks — each with its own rate-limiter wait -// and up to reportsMaxRetries backoff — well past the caller's soft wall-clock budget. +// time. deadline is checked before dispatching each app (never before the first one, so a call +// always makes some progress) — once past it, lookupUser stops dispatching new lookups but still +// waits for the ones already in flight, and returns what it has with nextResumeState set to the +// last app actually dispatched. Without this, a single user with a full budget's worth of apps +// could keep dispatching lookups — each with its own rate-limiter wait and up to +// reportsMaxRetries backoff — well past the caller's soft wall-clock budget. func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int, deadline time.Time) ([]*v2.Event, string, int, error) { tokenResp, err := client.ListTokens(ctx, user.ID) if err != nil { @@ -163,47 +164,39 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google toProcess = toProcess[:budget] } - events := make([]*v2.Event, 0, len(toProcess)) + results := make([]*v2.Event, len(toProcess)) consumed := 0 - for chunkStart := 0; chunkStart < len(toProcess); chunkStart += maxConcurrentAppLookups { - if chunkStart > 0 && !deadline.IsZero() && time.Now().After(deadline) { - // Past budget: stop starting new chunks. Whatever completed in prior chunks stands; - // nextResume (below) picks up right after the last app actually processed. + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(maxConcurrentAppLookups) + for i, app := range toProcess { + if i > 0 && !deadline.IsZero() && time.Now().After(deadline) { + // Past budget: stop dispatching new lookups. Ones already in flight are awaited + // below; nextResume picks up right after the last app actually dispatched. break } - chunkEnd := chunkStart + maxConcurrentAppLookups - if chunkEnd > len(toProcess) { - chunkEnd = len(toProcess) - } - chunk := toProcess[chunkStart:chunkEnd] - - results := make([]*v2.Event, len(chunk)) - g, gctx := errgroup.WithContext(ctx) - g.SetLimit(maxConcurrentAppLookups) - for i, app := range chunk { - g.Go(func() error { - event, err := f.lookupAppLogin(gctx, client, user, app.ClientID, app.DisplayText) - if err != nil { - return err - } - results[i] = event - return nil - }) - } - if err := g.Wait(); err != nil { - return nil, "", 0, err - } - - for _, e := range results { - if e != nil { - events = append(events, e) + consumed = i + 1 + g.Go(func() error { + event, err := f.lookupAppLogin(gctx, client, user, app.ClientID, app.DisplayText) + if err != nil { + return err } + results[i] = event + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, "", 0, err + } + + events := make([]*v2.Event, 0, consumed) + for _, e := range results[:consumed] { + if e != nil { + events = append(events, e) } - consumed += len(chunk) } nextResume := "" - if startIdx+consumed < len(apps) { + if consumed > 0 && startIdx+consumed < len(apps) { nextResume = toProcess[consumed-1].ClientID } return events, nextResume, consumed, nil From 2c381aced7aba20cfda01629f5d8e26875e87a70 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Tue, 18 Aug 2026 13:26:05 -0300 Subject: [PATCH 6/6] chore: fix lint errors --- pkg/connector/event_feed_common.go | 9 ++++++++- pkg/connector/usage_event_feed.go | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/pkg/connector/event_feed_common.go b/pkg/connector/event_feed_common.go index 9427d753..ede73193 100644 --- a/pkg/connector/event_feed_common.go +++ b/pkg/connector/event_feed_common.go @@ -115,7 +115,14 @@ func (c *userScanCursor) marshal() (string, error) { // nextResumeState reflecting the last unit of work it actually finished — never a zero-value // time.Time in practice, but implementations that only ever issue one call per user (there is no // fan-out to bound) are free to ignore it. -type userEventLookup func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int, deadline time.Time) (events []*v2.Event, nextResumeState string, consumed int, err error) +type userEventLookup func( + ctx context.Context, + client *gwclient.GoogleWorkspaceClient, + user pendingUser, + resumeState string, + budget int, + deadline time.Time, +) (events []*v2.Event, nextResumeState string, consumed int, err error) // scanUsersForEvents drives one bounded step of the rolling user-directory walk shared by all // three "last login" event feeds. diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 5f0d2aba..44b11219 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -133,7 +133,14 @@ func distinctAuthorizedApps(tokenResp *directoryAdmin.Tokens) []oauthApp { // last app actually dispatched. Without this, a single user with a full budget's worth of apps // could keep dispatching lookups — each with its own rate-limiter wait and up to // reportsMaxRetries backoff — well past the caller's soft wall-clock budget. -func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, resumeState string, budget int, deadline time.Time) ([]*v2.Event, string, int, error) { +func (f *usageEventFeed) lookupUser( + ctx context.Context, + client *gwclient.GoogleWorkspaceClient, + user pendingUser, + resumeState string, + budget int, + deadline time.Time, +) ([]*v2.Event, string, int, error) { tokenResp, err := client.ListTokens(ctx, user.ID) if err != nil { var gerr *googleapi.Error