CXP-533 Improve event feeds context handling

pkg/connector/event_feed_common.go

@@ -9,9 +9,8 @@
 // pagination.StreamToken, and never loops internally across directory pages (see
 // ref-antipatterns.md, "Client-Side Pagination Loop").
 //
-// Ordering note: it is not documented whether activities.list returns newest-first when no
-// startTime/orderBy is given, so each lookup fetches a small bounded window (not maxResults=1)
-// and picks the maximum occurredAt client-side, per the plan's acceptance-criteria fallback.
+// Ordering note: activities.list ordering is undocumented, so with startTime=180 days back and
+// maxResults=50, each lookup still picks the maximum occurredAt client-side rather than trusting result order.
 package connector
 
 import (

pkg/connector/google_login_event_feed.go

@@ -2,13 +2,17 @@ package connector
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
+	"time"
 
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	"github.com/conductorone/baton-sdk/pkg/types/resource"
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
+	"go.uber.org/zap"
 	reportsAdmin "google.golang.org/api/admin/reports/v1"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
@@ -20,11 +24,14 @@ type bestActivity struct {
 	occurredAt *timestamppb.Timestamp
 }
 
-// googleLoginLookupMaxResults bounds the per-user Reports API lookup for Google Workspace
-// sign-in events. There is exactly one target app (Google Workspace itself), so this only
-// needs to cover the "does maxResults=1 return newest-first?" ordering uncertainty: fetch a
-// small window and pick the maximum occurredAt client-side.
-const googleLoginLookupMaxResults = 5
+// googleLoginLookupMaxResults bounds the per-user Reports API lookup; the true latest event is
+// picked client-side, so this just needs to be large enough to avoid pagination.
+const googleLoginLookupMaxResults = 50
+
+// googleLoginLookupTimeout caps a single user's lookup, including retries. Kept above the worst
+// case of a hung attempt plus backoff plus a full retry (~51s) so the hung-attempt retry in
+// listActivitiesRateLimitedBounded has room to complete.
+const googleLoginLookupTimeout = 60 * time.Second
 
 // googleLoginEventFeed emits UsageEvents from Google Workspace sign-in activity.
 // Unlike SAML/OAuth feeds, the target resource is always Google Workspace itself.
@@ -48,8 +55,21 @@ func (f *googleLoginEventFeed) EventFeedMetadata(_ context.Context) *v2.EventFee
 }
 
 func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) {
-	r, err := listActivitiesRateLimited(ctx, client, user.Email, reportsAppLogin, "login_success", "", "", googleLoginLookupMaxResults)
+	l := ctxzap.Extract(ctx)
+	startTime := time.Now().Add(-reportsLookback).UTC().Format(time.RFC3339)
+
+	lookupCtx, cancel := context.WithTimeout(ctx, googleLoginLookupTimeout)
+	defer cancel()
+
+	r, err := listActivitiesRateLimitedBounded(lookupCtx, client, user.Email, reportsAppLogin, "login_success", startTime, "", googleLoginLookupMaxResults)
 	if err != nil {
+		if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil {
+			// Our sub-deadline fired, not the caller's context: skip this user instead of
+			// failing the whole batch.
+			l.Debug("google-workspace-connector: timed out listing google login activities, skipping",
+				zap.String("user", user.Email))
+			return nil, nil
+		}
 		return nil, fmt.Errorf("google-workspace-connector: failed to list google login activities for %s: %w", user.Email, err)
 	}
 

pkg/connector/reports_rate_limiter.go

@@ -32,6 +32,14 @@ const (
 	reportsMaxRetries                = 5
 	reportsInitialBackoff            = 500 * time.Millisecond
 	reportsMaxBackoff                = 30 * time.Second
+
+	// reportsPerAttemptTimeout bounds a single ListActivities call, not the retry loop as a
+	// whole, so a genuinely hung request is retried like any other transient error instead of
+	// being confused with the caller's overall lookup deadline expiring.
+	reportsPerAttemptTimeout = 25 * time.Second
+
+	// reportsLookback matches Google's Reports retention window; shared by the event feeds.
+	reportsLookback = 180 * 24 * time.Hour
 )
 
 // reportsRateLimiter is a simple token-bucket limiter built on the standard library only
@@ -92,38 +100,93 @@ func (l *reportsRateLimiter) refillLocked() {
 	l.lastRefill = l.now()
 }
 
+// listActivitiesFunc mirrors GoogleWorkspaceClient.ListActivities so retryListActivities can be
+// tested with a fake implementation.
+type listActivitiesFunc func(ctx context.Context, userKey, applicationName, eventName, startTime, pageToken, filters string, maxResults int64) (*reportsAdmin.Activities, error)
+
 // listActivitiesRateLimited waits for the shared filter-query budget, then calls
 // client.ListActivities, retrying with exponential backoff on 429/503 — both are transient,
-// SDK-retryable conditions, not connector bugs (see patterns-error-handling.md).
+// SDK-retryable conditions, not connector bugs (see patterns-error-handling.md). It imposes no
+// per-attempt timeout: use listActivitiesRateLimitedBounded for callers that set their own
+// lookupCtx sub-deadline and want a hung attempt retried instead of just waited out.
 func listActivitiesRateLimited(
 	ctx context.Context,
 	client *gwclient.GoogleWorkspaceClient,
 	userKey, applicationName, eventName, startTime, pageToken string,
 	maxResults int64,
 ) (*reportsAdmin.Activities, error) {
-	return listActivitiesFilteredRateLimited(ctx, client, userKey, applicationName, eventName, startTime, pageToken, "", maxResults)
+	return retryListActivities(
+		ctx, sharedReportsRateLimiter, client.ListActivities,
+		0, reportsMaxRetries, reportsInitialBackoff, reportsMaxBackoff,
+		userKey, applicationName, eventName, startTime, pageToken, "", maxResults,
+	)
+}
+
+// listActivitiesRateLimitedBounded is listActivitiesRateLimited plus reportsPerAttemptTimeout
+// applied to every attempt. Reserved for the event feed callers (usage/google-login/saml), which
+// always wrap the call in their own bounded lookupCtx, so a hung attempt can be told apart from
+// the caller's own deadline expiring and retried instead of failing the lookup outright.
+func listActivitiesRateLimitedBounded(
+	ctx context.Context,
+	client *gwclient.GoogleWorkspaceClient,
+	userKey, applicationName, eventName, startTime, pageToken string,
+	maxResults int64,
+) (*reportsAdmin.Activities, error) {
+	return listActivitiesFilteredRateLimitedBounded(ctx, client, userKey, applicationName, eventName, startTime, pageToken, "", maxResults)
 }
 
-// listActivitiesFilteredRateLimited is listActivitiesRateLimited plus an optional Reports API
-// `filters` expression (e.g. "client_id==<id>"), for callers that need to scope a lookup to one
-// specific app rather than an entire app-type.
-func listActivitiesFilteredRateLimited(
+// listActivitiesFilteredRateLimitedBounded is listActivitiesRateLimitedBounded plus an optional
+// Reports API `filters` expression (e.g. "client_id==<id>"), for callers that need to scope a
+// lookup to one specific app rather than an entire app-type.
+func listActivitiesFilteredRateLimitedBounded(
 	ctx context.Context,
 	client *gwclient.GoogleWorkspaceClient,
 	userKey, applicationName, eventName, startTime, pageToken, filters string,
 	maxResults int64,
 ) (*reportsAdmin.Activities, error) {
-	backoff := reportsInitialBackoff
+	return retryListActivities(
+		ctx, sharedReportsRateLimiter, client.ListActivities,
+		reportsPerAttemptTimeout, reportsMaxRetries, reportsInitialBackoff, reportsMaxBackoff,
+		userKey, applicationName, eventName, startTime, pageToken, filters, maxResults,
+	)
+}
+
+// retryListActivities holds the retry/backoff/per-attempt-timeout policy, parameterized so tests
+// can drive it with a fake call and short durations. perAttemptTimeout == 0 means "no per-attempt
+// cap" — the caller's own ctx is used as-is and a DeadlineExceeded from it is never retried as
+// a hung attempt.
+func retryListActivities(
+	ctx context.Context,
+	limiter *reportsRateLimiter,
+	call listActivitiesFunc,
+	perAttemptTimeout time.Duration,
+	maxRetries int,
+	initialBackoff, maxBackoff time.Duration,
+	userKey, applicationName, eventName, startTime, pageToken, filters string,
+	maxResults int64,
+) (*reportsAdmin.Activities, error) {
+	applyPerAttemptTimeout := perAttemptTimeout > 0
+	backoff := initialBackoff
 	for attempt := 0; ; attempt++ {
-		if err := sharedReportsRateLimiter.Wait(ctx); err != nil {
+		if err := limiter.Wait(ctx); err != nil {
 			return nil, fmt.Errorf("google-workspace-connector: context cancelled waiting for reports api quota: %w", err)
 		}
 
-		resp, err := client.ListActivities(ctx, userKey, applicationName, eventName, startTime, pageToken, filters, maxResults)
+		attemptCtx := ctx
+		cancel := func() {}
+		if applyPerAttemptTimeout {
+			attemptCtx, cancel = context.WithTimeout(ctx, perAttemptTimeout)
+		}
+		resp, err := call(attemptCtx, userKey, applicationName, eventName, startTime, pageToken, filters, maxResults)
+		cancel()
 		if err == nil {
 			return resp, nil
 		}
-		if attempt >= reportsMaxRetries || !isRetryableReportsError(err) {
+
+		// ctx still being live means attemptCtx's own timeout fired, not the caller's deadline —
+		// treat that like a retryable 429/503 rather than "out of time."
+		hungAttempt := applyPerAttemptTimeout && errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil
+		if attempt >= maxRetries || (!isRetryableReportsError(err) && !hungAttempt) {
 			return nil, err
 		}
 
@@ -134,8 +197,8 @@ func listActivitiesFilteredRateLimited(
 			return nil, ctx.Err()
 		}
 		backoff *= 2
-		if backoff > reportsMaxBackoff {
-			backoff = reportsMaxBackoff
+		if backoff > maxBackoff {
+			backoff = maxBackoff
 		}
 	}
 }