From 23310bb6a292849ec5ddc28d521ad85ea339784e Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Thu, 17 Sep 2026 13:16:26 -0300 Subject: [PATCH 1/2] fix: correct audit-log pagination cursor and sample skip-org warnings - nextAuditLogPage now checks resp.After first (the cursor GHEC and GHES actually return for the org audit-log endpoint's rel="next" Link), falling back to NextPageToken/NextPage for any page-style Link header GHES may still emit. The prior fix (#192) only checked NextPageToken/NextPage, so it silently truncated every GHEC org - and most GHES orgs - to a single page. - Request side now sends the cursor via After or Page depending on which shape was received (usageEventPageToken.AuditLogCursorIsPage), instead of always sending Page. - Orgs that permanently lack audit-log access are now warned via a sampled logger (1st, 10th, 100th, then every 1000th occurrence, with total_occurrences) instead of on every poll pass forever. The counter is intentionally shared across all orgs on the feed, not keyed per org. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/logging.go | 31 +++ pkg/connector/logging_test.go | 76 +++++++ pkg/connector/usage_event_feed.go | 71 ++++-- pkg/connector/usage_event_feed_test.go | 54 ++++- .../zap/zaptest/observer/logged_entry.go | 39 ++++ .../zap/zaptest/observer/observer.go | 203 ++++++++++++++++++ vendor/modules.txt | 1 + 7 files changed, 449 insertions(+), 26 deletions(-) create mode 100644 pkg/connector/logging.go create mode 100644 pkg/connector/logging_test.go create mode 100644 vendor/go.uber.org/zap/zaptest/observer/logged_entry.go create mode 100644 vendor/go.uber.org/zap/zaptest/observer/observer.go diff --git a/pkg/connector/logging.go b/pkg/connector/logging.go new file mode 100644 index 00000000..d98865dc --- /dev/null +++ b/pkg/connector/logging.go @@ -0,0 +1,31 @@ +package connector + +import ( + "context" + "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))...) +} + +// 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 + } +} diff --git a/pkg/connector/logging_test.go b/pkg/connector/logging_test.go new file mode 100644 index 00000000..0a0fed11 --- /dev/null +++ b/pkg/connector/logging_test.go @@ -0,0 +1,76 @@ +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 TestSampledWarn_SharedAcrossDistinctOrgs(t *testing.T) { + // skippedOrgs is intentionally a single shared counter on usageEventFeed, + // not keyed per org: the sampling budget applies to the aggregate rate of + // skip events across every org, not to each org individually. + core, logs := observer.New(zapcore.DebugLevel) + ctx := ctxzap.ToContext(context.Background(), zap.New(core)) + + var s sampledWarn + orgs := []string{"org-a", "org-b", "org-c"} + for i := 0; i < 9; i++ { + s.log(ctx, "org lacks audit-log access, skipping it for this pass", zap.String("org", orgs[i%len(orgs)])) + } + + // Only the very first occurrence (org-a) is logged; org-b's and org-c's + // first occurrences do not each get their own log line under the shared + // counter, since only occurrence 1 falls on the sampling schedule before 10. + require.Equal(t, 1, logs.Len()) + require.Equal(t, "org-a", logs.All()[0].ContextMap()["org"]) +} diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 01be314c..a793debf 100644 --- a/pkg/connector/usage_event_feed.go +++ b/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,9 @@ const maxAuditLogPagesPerCall = 20 type usageEventFeed struct { client *github.Client orgs []string + + // used to handle log Warn prints for skipped orgs. + skippedOrgs sampledWarn } func newUsageEventFeed(client *github.Client, orgs []string) *usageEventFeed { @@ -49,10 +51,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 +89,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 +127,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 +155,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 +188,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, "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,8 +215,9 @@ 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 } } @@ -213,6 +225,7 @@ func (f *usageEventFeed) ListEvents( // 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 +252,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= (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= (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= (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 diff --git a/pkg/connector/usage_event_feed_test.go b/pkg/connector/usage_event_feed_test.go index 06f0d481..e42c3be3 100644 --- a/pkg/connector/usage_event_feed_test.go +++ b/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) } diff --git a/vendor/go.uber.org/zap/zaptest/observer/logged_entry.go b/vendor/go.uber.org/zap/zaptest/observer/logged_entry.go new file mode 100644 index 00000000..ef89e25c --- /dev/null +++ b/vendor/go.uber.org/zap/zaptest/observer/logged_entry.go @@ -0,0 +1,39 @@ +// Copyright (c) 2017 Uber Technologies, Inc. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +package observer + +import "go.uber.org/zap/zapcore" + +// A LoggedEntry is an encoding-agnostic representation of a log message. +// Field availability is context dependent. +type LoggedEntry struct { + zapcore.Entry + Context []zapcore.Field +} + +// ContextMap returns a map for all fields in Context. +func (e LoggedEntry) ContextMap() map[string]interface{} { + encoder := zapcore.NewMapObjectEncoder() + for _, f := range e.Context { + f.AddTo(encoder) + } + return encoder.Fields +} diff --git a/vendor/go.uber.org/zap/zaptest/observer/observer.go b/vendor/go.uber.org/zap/zaptest/observer/observer.go new file mode 100644 index 00000000..4f7ce0ec --- /dev/null +++ b/vendor/go.uber.org/zap/zaptest/observer/observer.go @@ -0,0 +1,203 @@ +// Copyright (c) 2016-2022 Uber Technologies, Inc. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +// Package observer provides a zapcore.Core that keeps an in-memory, +// encoding-agnostic representation of log entries. It's useful for +// applications that want to unit test their log output without tying their +// tests to a particular output encoding. +package observer // import "go.uber.org/zap/zaptest/observer" + +import ( + "strings" + "sync" + "time" + + "go.uber.org/zap/internal" + "go.uber.org/zap/zapcore" +) + +// ObservedLogs is a concurrency-safe, ordered collection of observed logs. +type ObservedLogs struct { + mu sync.RWMutex + logs []LoggedEntry +} + +// Len returns the number of items in the collection. +func (o *ObservedLogs) Len() int { + o.mu.RLock() + n := len(o.logs) + o.mu.RUnlock() + return n +} + +// All returns a copy of all the observed logs. +func (o *ObservedLogs) All() []LoggedEntry { + o.mu.RLock() + ret := make([]LoggedEntry, len(o.logs)) + copy(ret, o.logs) + o.mu.RUnlock() + return ret +} + +// TakeAll returns a copy of all the observed logs, and truncates the observed +// slice. +func (o *ObservedLogs) TakeAll() []LoggedEntry { + o.mu.Lock() + ret := o.logs + o.logs = nil + o.mu.Unlock() + return ret +} + +// AllUntimed returns a copy of all the observed logs, but overwrites the +// observed timestamps with time.Time's zero value. This is useful when making +// assertions in tests. +func (o *ObservedLogs) AllUntimed() []LoggedEntry { + ret := o.All() + for i := range ret { + ret[i].Time = time.Time{} + } + return ret +} + +// FilterLevelExact filters entries to those logged at exactly the given level. +func (o *ObservedLogs) FilterLevelExact(level zapcore.Level) *ObservedLogs { + return o.Filter(func(e LoggedEntry) bool { + return e.Level == level + }) +} + +// FilterMessage filters entries to those that have the specified message. +func (o *ObservedLogs) FilterMessage(msg string) *ObservedLogs { + return o.Filter(func(e LoggedEntry) bool { + return e.Message == msg + }) +} + +// FilterLoggerName filters entries to those logged through logger with the specified logger name. +func (o *ObservedLogs) FilterLoggerName(name string) *ObservedLogs { + return o.Filter(func(e LoggedEntry) bool { + return e.LoggerName == name + }) +} + +// FilterMessageSnippet filters entries to those that have a message containing the specified snippet. +func (o *ObservedLogs) FilterMessageSnippet(snippet string) *ObservedLogs { + return o.Filter(func(e LoggedEntry) bool { + return strings.Contains(e.Message, snippet) + }) +} + +// FilterField filters entries to those that have the specified field. +func (o *ObservedLogs) FilterField(field zapcore.Field) *ObservedLogs { + return o.Filter(func(e LoggedEntry) bool { + for _, ctxField := range e.Context { + if ctxField.Equals(field) { + return true + } + } + return false + }) +} + +// FilterFieldKey filters entries to those that have the specified key. +func (o *ObservedLogs) FilterFieldKey(key string) *ObservedLogs { + return o.Filter(func(e LoggedEntry) bool { + for _, ctxField := range e.Context { + if ctxField.Key == key { + return true + } + } + return false + }) +} + +// Filter returns a copy of this ObservedLogs containing only those entries +// for which the provided function returns true. +func (o *ObservedLogs) Filter(keep func(LoggedEntry) bool) *ObservedLogs { + o.mu.RLock() + defer o.mu.RUnlock() + + var filtered []LoggedEntry + for _, entry := range o.logs { + if keep(entry) { + filtered = append(filtered, entry) + } + } + return &ObservedLogs{logs: filtered} +} + +func (o *ObservedLogs) add(log LoggedEntry) { + o.mu.Lock() + o.logs = append(o.logs, log) + o.mu.Unlock() +} + +// New creates a new Core that buffers logs in memory (without any encoding). +// It's particularly useful in tests. +func New(enab zapcore.LevelEnabler) (zapcore.Core, *ObservedLogs) { + ol := &ObservedLogs{} + return &contextObserver{ + LevelEnabler: enab, + logs: ol, + }, ol +} + +type contextObserver struct { + zapcore.LevelEnabler + logs *ObservedLogs + context []zapcore.Field +} + +var ( + _ zapcore.Core = (*contextObserver)(nil) + _ internal.LeveledEnabler = (*contextObserver)(nil) +) + +func (co *contextObserver) Level() zapcore.Level { + return zapcore.LevelOf(co.LevelEnabler) +} + +func (co *contextObserver) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry { + if co.Enabled(ent.Level) { + return ce.AddCore(ent, co) + } + return ce +} + +func (co *contextObserver) With(fields []zapcore.Field) zapcore.Core { + return &contextObserver{ + LevelEnabler: co.LevelEnabler, + logs: co.logs, + context: append(co.context[:len(co.context):len(co.context)], fields...), + } +} + +func (co *contextObserver) Write(ent zapcore.Entry, fields []zapcore.Field) error { + all := make([]zapcore.Field, 0, len(fields)+len(co.context)) + all = append(all, co.context...) + all = append(all, fields...) + co.logs.add(LoggedEntry{ent, all}) + return nil +} + +func (co *contextObserver) Sync() error { + return nil +} diff --git a/vendor/modules.txt b/vendor/modules.txt index 8192e9fc..061ef043 100644 --- a/vendor/modules.txt +++ b/vendor/modules.txt @@ -784,6 +784,7 @@ go.uber.org/zap/internal/exit go.uber.org/zap/internal/pool go.uber.org/zap/internal/stacktrace go.uber.org/zap/zapcore +go.uber.org/zap/zaptest/observer # golang.org/x/crypto v0.54.0 ## explicit; go 1.25.0 golang.org/x/crypto/blowfish From dc219ce236761395a1ba3978d1cb54bba81e1e0e Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Thu, 17 Sep 2026 13:28:53 -0300 Subject: [PATCH 2/2] fix: key skipped-org warning sampling per org, not shared A single shared counter meant one permanently-inaccessible org could drive the sampling budget high enough that another org's first failure landed on a non-sampled occurrence and was never logged, hiding it from operators entirely (PR #193 review). perKeySampledWarn keys the sampler by org name so every distinct org gets its own guaranteed 1st/10th/100th/every-1000th occurrence, independent of how noisy any other org is. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/logging.go | 10 +++++++ pkg/connector/logging_test.go | 45 +++++++++++++++++++++++-------- pkg/connector/usage_event_feed.go | 7 ++--- 3 files changed, 48 insertions(+), 14 deletions(-) diff --git a/pkg/connector/logging.go b/pkg/connector/logging.go index d98865dc..416ca117 100644 --- a/pkg/connector/logging.go +++ b/pkg/connector/logging.go @@ -2,6 +2,7 @@ package connector import ( "context" + "sync" "sync/atomic" "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" @@ -20,6 +21,15 @@ func (s *sampledWarn) log(ctx context.Context, msg string, fields ...zap.Field) 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 { diff --git a/pkg/connector/logging_test.go b/pkg/connector/logging_test.go index 0a0fed11..12b41677 100644 --- a/pkg/connector/logging_test.go +++ b/pkg/connector/logging_test.go @@ -55,22 +55,45 @@ func TestSampledWarn_LogsOnlyOnSampledOccurrences(t *testing.T) { require.Equal(t, uint64(10), second.ContextMap()["total_occurrences"]) } -func TestSampledWarn_SharedAcrossDistinctOrgs(t *testing.T) { - // skippedOrgs is intentionally a single shared counter on usageEventFeed, - // not keyed per org: the sampling budget applies to the aggregate rate of - // skip events across every org, not to each org individually. +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 s sampledWarn + var p perKeySampledWarn orgs := []string{"org-a", "org-b", "org-c"} for i := 0; i < 9; i++ { - s.log(ctx, "org lacks audit-log access, skipping it for this pass", zap.String("org", orgs[i%len(orgs)])) + 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")) - // Only the very first occurrence (org-a) is logged; org-b's and org-c's - // first occurrences do not each get their own log line under the shared - // counter, since only occurrence 1 falls on the sampling schedule before 10. - require.Equal(t, 1, logs.Len()) - require.Equal(t, "org-a", logs.All()[0].ContextMap()["org"]) + 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"]) } diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index a793debf..f0905c35 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -32,8 +32,9 @@ type usageEventFeed struct { client *github.Client orgs []string - // used to handle log Warn prints for skipped orgs. - skippedOrgs sampledWarn + // 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 { @@ -188,7 +189,7 @@ 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): - f.skippedOrgs.log(ctx, "org lacks audit-log access, skipping it for this pass", + f.skippedOrgs.log(ctx, orgName, "org lacks audit-log access, skipping it for this pass", zap.String("org", orgName), zap.Error(err), )