diff --git a/packages/auth/pkg/auth/internal/service/store.go b/packages/auth/pkg/auth/internal/service/store.go index 43ede9009e..43a9cba974 100644 --- a/packages/auth/pkg/auth/internal/service/store.go +++ b/packages/auth/pkg/auth/internal/service/store.go @@ -3,6 +3,9 @@ package service import ( "context" "fmt" + "sync" + "sync/atomic" + "time" "github.com/google/uuid" "go.opentelemetry.io/otel" @@ -17,6 +20,42 @@ import ( var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/auth/pkg/auth/internal/service") +// last_used is minute-grade observability metadata, but updating it on every +// authenticated request makes team_api_keys one of the highest dead-tuple +// producers in the registry. One write per key per window keeps it fresh +// enough at a fraction of the churn. +const lastUsedWriteWindow = time.Minute + +var ( + lastUsedWrites sync.Map // api key hash -> time.Time of last write + lastUsedCallsSweep atomic.Int64 +) + +func shouldWriteLastUsed(hashedKey string, now time.Time) bool { + // Occasionally drop entries idle for many windows so the map tracks the + // working set of keys, not every key ever seen by the process. + if lastUsedCallsSweep.Add(1)%4096 == 0 { + lastUsedWrites.Range(func(k, v any) bool { + if now.Sub(v.(time.Time)) > 10*lastUsedWriteWindow { + lastUsedWrites.Delete(k) + } + + return true + }) + } + + prev, loaded := lastUsedWrites.LoadOrStore(hashedKey, now) + if !loaded { + return true + } + if now.Sub(prev.(time.Time)) < lastUsedWriteWindow { + return false + } + + // CAS so exactly one concurrent caller wins the expired window. + return lastUsedWrites.CompareAndSwap(hashedKey, prev, now) +} + type authStoreImpl struct { authDB *authdb.Client } @@ -44,14 +83,16 @@ func (s *authStoreImpl) GetTeamByHashedAPIKey(ctx context.Context, hashedKey str return nil, err } - go func() { - // Run the update in a separate context to avoid an extra latency - ctx := context.WithoutCancel(ctx) - updateErr := s.authDB.UpdateLastTimeUsed(ctx, hashedKey) - if updateErr != nil { - logger.L().Error(ctx, "failed to update last time used", zap.Error(updateErr)) - } - }() + if shouldWriteLastUsed(hashedKey, time.Now()) { + go func() { + // Run the update in a separate context to avoid an extra latency + ctx := context.WithoutCancel(ctx) + updateErr := s.authDB.UpdateLastTimeUsed(ctx, hashedKey) + if updateErr != nil { + logger.L().Error(ctx, "failed to update last time used", zap.Error(updateErr)) + } + }() + } team := types.NewTeam(&result.Team, &result.TeamLimit) diff --git a/packages/auth/pkg/auth/internal/service/store_test.go b/packages/auth/pkg/auth/internal/service/store_test.go new file mode 100644 index 0000000000..40acebb3ca --- /dev/null +++ b/packages/auth/pkg/auth/internal/service/store_test.go @@ -0,0 +1,60 @@ +package service + +import ( + "sync" + "testing" + "time" +) + +func TestShouldWriteLastUsed(t *testing.T) { + t.Parallel() + + // Unique keys per invocation: the debounce map is package state shared + // across repeated in-process runs (-count=2). + keyA := "key-a-" + t.Name() + time.Now().String() + keyB := "key-b-" + t.Name() + time.Now().String() + base := time.Now() + + if !shouldWriteLastUsed(keyA, base) { + t.Fatal("first write for a key must pass") + } + if shouldWriteLastUsed(keyA, base.Add(lastUsedWriteWindow/2)) { + t.Fatal("write inside the window must be suppressed") + } + if !shouldWriteLastUsed(keyB, base) { + t.Fatal("independent key must not be suppressed") + } + if !shouldWriteLastUsed(keyA, base.Add(lastUsedWriteWindow+time.Second)) { + t.Fatal("write after the window must pass") + } +} + +func TestShouldWriteLastUsedConcurrent(t *testing.T) { + t.Parallel() + + key := "key-conc-" + time.Now().String() + base := time.Now() + shouldWriteLastUsed(key, base) + + // After the window expires, exactly one concurrent caller wins. + later := base.Add(lastUsedWriteWindow + time.Second) + const n = 16 + wins := make(chan bool, n) + var wg sync.WaitGroup + for range n { + wg.Go(func() { + wins <- shouldWriteLastUsed(key, later) + }) + } + wg.Wait() + close(wins) + won := 0 + for w := range wins { + if w { + won++ + } + } + if won != 1 { + t.Fatalf("expected exactly one winner, got %d", won) + } +}