CXP-383 Correct audit-log pagination cursor and sample skip-org warnings

pkg/connector/logging.go

@@ -0,0 +1,41 @@
+package connector
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
+	"go.uber.org/zap"
+)
+
+type sampledWarn struct {
+	n atomic.Uint64
+}
+
+func (s *sampledWarn) log(ctx context.Context, msg string, fields ...zap.Field) {
+	n := s.n.Add(1)
+	if !shouldLogSample(n) {
+		return
+	}
+	ctxzap.Extract(ctx).Warn(msg, append(fields, zap.Uint64("total_occurrences", n))...)
+}
+
+type perKeySampledWarn struct {
+	m sync.Map // key -> *sampledWarn
+}
+
+func (p *perKeySampledWarn) log(ctx context.Context, key, msg string, fields ...zap.Field) {
+	actual, _ := p.m.LoadOrStore(key, &sampledWarn{})
+	actual.(*sampledWarn).log(ctx, msg, fields...)
+}
+
+// shouldLogSample reports whether the nth occurrence should be logged.
+func shouldLogSample(n uint64) bool {
+	switch {
+	case n <= 1, n == 10, n == 100:
+		return true
+	default:
+		return n%1000 == 0
+	}
+}

pkg/connector/logging_test.go

@@ -0,0 +1,99 @@
+package connector
+
+import (
+	"context"
+	"testing"
+
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+	"go.uber.org/zap/zaptest/observer"
+)
+
+func TestShouldLogSample(t *testing.T) {
+	tests := []struct {
+		n    uint64
+		want bool
+	}{
+		{0, true},
+		{1, true},
+		{2, false},
+		{9, false},
+		{10, true},
+		{11, false},
+		{99, false},
+		{100, true},
+		{101, false},
+		{999, false},
+		{1000, true},
+		{1001, false},
+		{1999, false},
+		{2000, true},
+	}
+	for _, tt := range tests {
+		require.Equal(t, tt.want, shouldLogSample(tt.n), "n=%d", tt.n)
+	}
+}
+
+func TestSampledWarn_LogsOnlyOnSampledOccurrences(t *testing.T) {
+	core, logs := observer.New(zapcore.DebugLevel)
+	ctx := ctxzap.ToContext(context.Background(), zap.New(core))
+
+	var s sampledWarn
+	for i := 0; i < 12; i++ {
+		s.log(ctx, "org lacks audit-log access, skipping it for this pass", zap.String("org", "octo-org"))
+	}
+
+	// Occurrences 1 and 10 should be logged; 2-9 and 11-12 should not.
+	require.Equal(t, 2, logs.Len(), "expected exactly 2 sampled log lines out of 12 occurrences")
+
+	first := logs.All()[0]
+	require.Equal(t, uint64(1), first.ContextMap()["total_occurrences"])
+
+	second := logs.All()[1]
+	require.Equal(t, uint64(10), second.ContextMap()["total_occurrences"])
+}
+
+func TestPerKeySampledWarn_EachKeyGetsItsOwnBudget(t *testing.T) {
+	// skippedOrgs is keyed per org so one noisy org's sampling budget can't
+	// starve another org's first occurrence out of the log.
+	core, logs := observer.New(zapcore.DebugLevel)
+	ctx := ctxzap.ToContext(context.Background(), zap.New(core))
+
+	var p perKeySampledWarn
+	orgs := []string{"org-a", "org-b", "org-c"}
+	for i := 0; i < 9; i++ {
+		org := orgs[i%len(orgs)]
+		p.log(ctx, org, "org lacks audit-log access, skipping it for this pass", zap.String("org", org))
+	}
+
+	// Each org's first occurrence is its own occurrence 1, so all three log.
+	require.Equal(t, 3, logs.Len())
+	gotOrgs := make([]string, len(logs.All()))
+	for i, entry := range logs.All() {
+		gotOrgs[i] = entry.ContextMap()["org"].(string)
+	}
+	require.ElementsMatch(t, orgs, gotOrgs)
+}
+
+func TestPerKeySampledWarn_NoisyKeyDoesNotStarveNewKey(t *testing.T) {
+	// Reproduces the diagnosability gap a shared counter has: org-a alone
+	// drives the budget deep into a high sampling gap, then org-b's very
+	// first failure must still surface immediately rather than landing on
+	// org-a's non-sampled occurrence.
+	core, logs := observer.New(zapcore.DebugLevel)
+	ctx := ctxzap.ToContext(context.Background(), zap.New(core))
+
+	var p perKeySampledWarn
+	for i := 0; i < 400; i++ {
+		p.log(ctx, "org-a", "org lacks audit-log access, skipping it for this pass", zap.String("org", "org-a"))
+	}
+	logs.TakeAll() // discard org-a's own sampled lines; only org-b matters here
+
+	p.log(ctx, "org-b", "org lacks audit-log access, skipping it for this pass", zap.String("org", "org-b"))
+
+	require.Equal(t, 1, logs.Len(), "org-b's first failure must be logged even though org-a's counter is at 400")
+	require.Equal(t, "org-b", logs.All()[0].ContextMap()["org"])
+	require.Equal(t, uint64(1), logs.All()[0].ContextMap()["total_occurrences"])
+}

pkg/connector/usage_event_feed.go

@@ -14,7 +14,6 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	"github.com/google/go-github/v69/github"
-	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -32,6 +31,10 @@ const maxAuditLogPagesPerCall = 20
 type usageEventFeed struct {
 	client *github.Client
 	orgs   []string
+
+	// used to handle log Warn prints for skipped orgs, sampled per org so
+	// one noisy org can't starve another org's first occurrence out of the log.
+	skippedOrgs perKeySampledWarn
 }
 
 func newUsageEventFeed(client *github.Client, orgs []string) *usageEventFeed {
@@ -49,10 +52,14 @@ func (f *usageEventFeed) EventFeedMetadata(_ context.Context) *v2.EventFeedMetad
 // org's audit log, walked newest-first until an already-seen entry (at or
 // before Since) is reached.
 type usageEventPageToken struct {
-	Orgs           []string `json:"orgs,omitempty"`
-	OrgIndex       int      `json:"org_index"`
-	AuditLogCursor string   `json:"audit_log_cursor,omitempty"`
-	Since          string   `json:"since,omitempty"`
+	Orgs     []string `json:"orgs,omitempty"`
+	OrgIndex int      `json:"org_index"`
+	// AuditLogCursor is the opaque value to resume from. Whether it goes into
+	// the request's After or Page field depends on AuditLogCursorIsPage - see
+	// nextAuditLogPage.
+	AuditLogCursor       string `json:"audit_log_cursor,omitempty"`
+	AuditLogCursorIsPage bool   `json:"audit_log_cursor_is_page,omitempty"`
+	Since                string `json:"since,omitempty"`
 }
 
 func unmarshalUsageEventPageToken(pToken *pagination.StreamToken) (*usageEventPageToken, error) {
@@ -83,8 +90,6 @@ func (f *usageEventFeed) ListEvents(
 	earliestEvent *timestamppb.Timestamp,
 	pToken *pagination.StreamToken,
 ) ([]*v2.Event, *pagination.StreamState, annotations.Annotations, error) {
-	l := ctxzap.Extract(ctx)
-
 	if f.client == nil {
 		return nil, &pagination.StreamState{HasMore: false}, nil, nil
 	}
@@ -123,6 +128,7 @@ func (f *usageEventFeed) ListEvents(
 	if cursor.OrgIndex < 0 || cursor.OrgIndex >= len(cursor.Orgs) {
 		cursor.OrgIndex = 0
 		cursor.AuditLogCursor = ""
+		cursor.AuditLogCursorIsPage = false
 	}
 
 	since, err := time.Parse(time.RFC3339Nano, cursor.Since)
@@ -150,9 +156,13 @@ func (f *usageEventFeed) ListEvents(
 			Phrase:  github.Ptr(sincePhrase),
 			ListCursorOptions: github.ListCursorOptions{
 				PerPage: maxPageSize,
-				Page:    cursor.AuditLogCursor,
 			},
 		}
+		if cursor.AuditLogCursorIsPage {
+			opts.Page = cursor.AuditLogCursor
+		} else {
+			opts.After = cursor.AuditLogCursor
+		}
 
 		entries, resp, err := f.client.Organizations.GetAuditLog(ctx, orgName, opts)
 		// Read rate-limit headers before the error branch nils resp, since a
@@ -179,8 +189,10 @@ func (f *usageEventFeed) ListEvents(
 				return nil, nil, nil, wrapGitHubError(err, resp,
 					fmt.Sprintf("baton-github: failed to fetch audit log for org %s", orgName))
 			case isNotFoundError(resp) || isPermissionError(resp):
-				l.Warn("org lacks audit-log access, skipping it for this pass",
-					zap.String("org", orgName), zap.Error(err))
+				f.skippedOrgs.log(ctx, orgName, "org lacks audit-log access, skipping it for this pass",
+					zap.String("org", orgName), zap.Error(err),
+				)
+
 				entries, resp = nil, nil
 			default:
 				return nil, nil, nil, wrapGitHubError(err, resp,
@@ -204,15 +216,17 @@ func (f *usageEventFeed) ListEvents(
 		}
 
 		if resp != nil && !reachedBoundary {
-			if nextPage := nextAuditLogPage(resp); nextPage != "" {
+			if nextPage, isPage := nextAuditLogPage(resp); nextPage != "" {
 				cursor.AuditLogCursor = nextPage
+				cursor.AuditLogCursorIsPage = isPage
 				continue
 			}
 		}
 
 		// Done with this org for this pass - advance to the next one.
 		cursor.OrgIndex++
 		cursor.AuditLogCursor = ""
+		cursor.AuditLogCursorIsPage = false
 		if cursor.OrgIndex >= len(cursor.Orgs) {
 			// Pass complete - the next call gets a fresh earliestEvent, so
 			// nothing needs to survive in the cursor.
@@ -239,21 +253,35 @@ func (f *usageEventFeed) ListEvents(
 	return events, &pagination.StreamState{Cursor: tokenStr, HasMore: true}, annos, nil
 }
 
-// nextAuditLogPage returns the token to request the next audit-log page, or
-// "" if there isn't one. GitHub's org audit-log endpoint returns opaque
-// cursor pagination on github.com/GHEC (go-github parses the Link header's
-// non-numeric "page" value into Response.NextPageToken), but GHES-style
-// numeric "page=N" Link headers parse into Response.NextPage (int) instead,
-// leaving NextPageToken empty. Checking only NextPageToken silently truncates
-// GHES audit logs to a single page.
-func nextAuditLogPage(resp *github.Response) string {
+// nextAuditLogPage returns the cursor to request the next audit-log page
+// (empty if there isn't one), and whether that cursor belongs in the
+// request's Page field (true) or its After field (false).
+//
+// The org audit-log endpoint documents three pagination shapes depending on
+// what the server returns in the Link header's rel="next" entry:
+//   - after=<cursor> (GHEC and GHES): the documented, primary mechanism -
+//     go-github parses this into Response.After. Checked first since it's
+//     what both github.com and GHES actually return in practice.
+//   - page=<opaque token> (legacy fallback some GHES versions may still
+//     emit): go-github can't parse a non-numeric page value as an int, so it
+//     lands in Response.NextPageToken instead.
+//   - page=<N> (numeric, classic GHES offset pagination): parses into
+//     Response.NextPage.
+//
+// Checking only NextPageToken/NextPage (as earlier code did) misses the
+// after= case entirely, silently truncating every GHEC org - and most GHES
+// orgs - to a single page.
+func nextAuditLogPage(resp *github.Response) (string, bool) {
+	if resp.After != "" {
+		return resp.After, false
+	}
 	if resp.NextPageToken != "" {
-		return resp.NextPageToken
+		return resp.NextPageToken, true
 	}
 	if resp.NextPage != 0 {
-		return strconv.Itoa(resp.NextPage)
+		return strconv.Itoa(resp.NextPage), true
 	}
-	return ""
+	return "", false
 }
 
 // usageEventFromAuditEntry converts one audit-log entry into a usage event

pkg/connector/usage_event_feed_test.go

@@ -433,22 +433,67 @@ func TestUsageEventFeed_ListEvents_ResumesFromPersistedCursor(t *testing.T) {
 	newer := since.Add(1 * time.Hour)
 
 	// Simulate a previous call that already finished "octo-org-a" and was
-	// mid-page through "octo-org-b" with its own audit-log cursor.
+	// mid-page through "octo-org-b" with its own audit-log cursor, resuming
+	// via classic numeric-page pagination (as GHES may still return).
+	resumeToken := &usageEventPageToken{
+		Orgs:                 []string{"octo-org-a", "octo-org-b"},
+		OrgIndex:             1,
+		AuditLogCursor:       "existing-cursor",
+		AuditLogCursorIsPage: true,
+		Since:                since.Format(time.RFC3339),
+	}
+	cursorStr, err := resumeToken.marshal()
+	require.NoError(t, err)
+
+	var gotOrg, gotPage string
+	httpClient := mock.NewMockedHTTPClient(
+		mock.WithRequestMatchHandler(
+			mock.GetOrgsAuditLogByOrg,
+			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				gotOrg = strings.Split(r.URL.Path, "/")[2]
+				gotPage = r.URL.Query().Get("page")
+				w.Header().Set("Content-Type", "application/json")
+				_, _ = w.Write(mock.MustMarshal([]*github.AuditEntry{
+					{Actor: github.Ptr("bob"), ActorID: github.Ptr(int64(2)), OrgID: github.Ptr(int64(8)), Timestamp: &github.Timestamp{Time: newer}},
+				}))
+			}),
+		),
+	)
+
+	f := newUsageEventFeed(github.NewClient(httpClient), nil)
+
+	events, state, _, err := f.ListEvents(ctx, nil, &pagination.StreamToken{Cursor: cursorStr})
+	require.NoError(t, err)
+	require.Equal(t, "octo-org-b", gotOrg, "should resume at the persisted org, not restart from octo-org-a")
+	require.Equal(t, "existing-cursor", gotPage, "should resume with the persisted audit-log cursor")
+	require.Len(t, events, 1)
+	require.False(t, state.HasMore)
+}
+
+func TestUsageEventFeed_ListEvents_ResumesFromPersistedAfterCursor(t *testing.T) {
+	ctx := context.Background()
+
+	since := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
+	newer := since.Add(1 * time.Hour)
+
+	// Both github.com/GHEC and GHES return an after= cursor (not a numeric
+	// page) in practice, so this is the default/expected resumption path.
 	resumeToken := &usageEventPageToken{
 		Orgs:           []string{"octo-org-a", "octo-org-b"},
 		OrgIndex:       1,
-		AuditLogCursor: "existing-cursor",
+		AuditLogCursor: "existing-after-cursor",
 		Since:          since.Format(time.RFC3339),
 	}
 	cursorStr, err := resumeToken.marshal()
 	require.NoError(t, err)
 
-	var gotOrg, gotPage string
+	var gotOrg, gotAfter, gotPage string
 	httpClient := mock.NewMockedHTTPClient(
 		mock.WithRequestMatchHandler(
 			mock.GetOrgsAuditLogByOrg,
 			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				gotOrg = strings.Split(r.URL.Path, "/")[2]
+				gotAfter = r.URL.Query().Get("after")
 				gotPage = r.URL.Query().Get("page")
 				w.Header().Set("Content-Type", "application/json")
 				_, _ = w.Write(mock.MustMarshal([]*github.AuditEntry{
@@ -463,7 +508,8 @@ func TestUsageEventFeed_ListEvents_ResumesFromPersistedCursor(t *testing.T) {
 	events, state, _, err := f.ListEvents(ctx, nil, &pagination.StreamToken{Cursor: cursorStr})
 	require.NoError(t, err)
 	require.Equal(t, "octo-org-b", gotOrg, "should resume at the persisted org, not restart from octo-org-a")
-	require.Equal(t, "existing-cursor", gotPage, "should resume with the persisted audit-log cursor")
+	require.Equal(t, "existing-after-cursor", gotAfter, "should resume with the persisted cursor via the after= param")
+	require.Empty(t, gotPage, "should not send the cursor as a numeric page param")
 	require.Len(t, events, 1)
 	require.False(t, state.HasMore)
 }