diff --git a/apps/api/internal/store/api_key_auth_integration_test.go b/apps/api/internal/store/api_key_auth_integration_test.go index 8247c5e..f20828f 100644 --- a/apps/api/internal/store/api_key_auth_integration_test.go +++ b/apps/api/internal/store/api_key_auth_integration_test.go @@ -65,6 +65,36 @@ func TestVerifyLocalAPIKeyHandlesEightConcurrentRequestsWithFourConnections(t *t } } +func TestVerifyLocalAPIKeyCoalescesUsageUpdates(t *testing.T) { + db, verificationStore, created, _ := newLocalAPIKeyVerificationFixture(t, 1) + ctx := context.Background() + + if _, err := verificationStore.VerifyLocalAPIKey(ctx, created.Secret); err != nil { + t.Fatalf("verify local API key first time: %v", err) + } + var firstLastUsedAt time.Time + if err := db.pool.QueryRow(ctx, ` +SELECT last_used_at +FROM gateway_api_keys +WHERE id = $1::uuid`, created.APIKey.ID).Scan(&firstLastUsedAt); err != nil { + t.Fatalf("read first API key usage timestamp: %v", err) + } + + if _, err := verificationStore.VerifyLocalAPIKey(ctx, created.Secret); err != nil { + t.Fatalf("verify local API key second time: %v", err) + } + var secondLastUsedAt time.Time + if err := db.pool.QueryRow(ctx, ` +SELECT last_used_at +FROM gateway_api_keys +WHERE id = $1::uuid`, created.APIKey.ID).Scan(&secondLastUsedAt); err != nil { + t.Fatalf("read second API key usage timestamp: %v", err) + } + if !secondLastUsedAt.Equal(firstLastUsedAt) { + t.Fatalf("usage timestamp changed inside coalescing window: first=%s second=%s", firstLastUsedAt, secondLastUsedAt) + } +} + func newLocalAPIKeyVerificationFixture(t *testing.T, maxConnections int32) (*Store, *Store, CreatedAPIKey, GatewayUser) { t.Helper() db := newIdentityPairingPostgresTestStore(t) diff --git a/apps/api/internal/store/api_key_auth_test.go b/apps/api/internal/store/api_key_auth_test.go index 5b8e299..87a4ce4 100644 --- a/apps/api/internal/store/api_key_auth_test.go +++ b/apps/api/internal/store/api_key_auth_test.go @@ -3,6 +3,7 @@ package store import ( "context" "errors" + "strings" "testing" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" @@ -50,6 +51,12 @@ func TestVerifyLocalAPIKeyClosesCandidateRowsBeforeUpdatingUsage(t *testing.T) { if database.updatedAPIKeyID != "matching-key" { t.Fatalf("updated API key = %q, want matching-key", database.updatedAPIKeyID) } + if !strings.Contains(database.updateSQL, "set_config('synchronous_commit', 'off', true)") { + t.Fatalf("API key usage update did not disable synchronous commit: %s", database.updateSQL) + } + if !strings.Contains(database.updateSQL, "interval '1 minute'") { + t.Fatalf("API key usage update did not coalesce hot-key writes: %s", database.updateSQL) + } if user.APIKeyID != "matching-key" || user.GatewayUserID != "user-id" { t.Fatalf("verified user = %+v", user) } @@ -82,16 +89,18 @@ func TestVerifyLocalAPIKeyReturnsUnauthorizedAfterClosingCandidateRows(t *testin type fakeLocalAPIKeyDatabase struct { rows *fakeLocalAPIKeyRows updatedAPIKeyID string + updateSQL string } func (database *fakeLocalAPIKeyDatabase) Query(context.Context, string, ...any) (pgx.Rows, error) { return database.rows, nil } -func (database *fakeLocalAPIKeyDatabase) Exec(_ context.Context, _ string, arguments ...any) (pgconn.CommandTag, error) { +func (database *fakeLocalAPIKeyDatabase) Exec(_ context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { if !database.rows.closed { return pgconn.CommandTag{}, errors.New("API key usage update started before candidate rows closed") } + database.updateSQL = sql database.updatedAPIKeyID, _ = arguments[0].(string) return pgconn.NewCommandTag("UPDATE 1"), nil } diff --git a/apps/api/internal/store/postgres.go b/apps/api/internal/store/postgres.go index 9397377..32d36f5 100644 --- a/apps/api/internal/store/postgres.go +++ b/apps/api/internal/store/postgres.go @@ -1709,7 +1709,19 @@ WHERE k.key_prefix = $1 if bcrypt.CompareHashAndPassword([]byte(candidate.hash), []byte(secret)) != nil { continue } - if _, err := database.Exec(ctx, `UPDATE gateway_api_keys SET last_used_at = now(), updated_at = now() WHERE id = $1::uuid`, candidate.apiKeyID); err != nil { + // API key usage timestamps are telemetry, not authentication state. Keep + // their writes off the synchronous-replication critical path and coalesce + // hot-key updates so concurrent requests do not contend on the same row. + if _, err := database.Exec(ctx, ` +WITH commit_policy AS ( + SELECT set_config('synchronous_commit', 'off', true) +) +UPDATE gateway_api_keys +SET last_used_at = now(), + updated_at = now() +FROM commit_policy +WHERE id = $1::uuid + AND (last_used_at IS NULL OR last_used_at < now() - interval '1 minute')`, candidate.apiKeyID); err != nil { return nil, err } return &auth.User{ diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index c8dda15..839b621 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -46,6 +46,12 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra return WorkerAllocation{}, err } defer tx.Rollback(ctx) + // Worker heartbeats and allocations are ephemeral coordination state that is + // refreshed every few seconds. Do not hold the global allocation lock while + // waiting for a synchronous replica to acknowledge the commit. + if _, err := tx.Exec(ctx, `SELECT set_config('synchronous_commit', 'off', true)`); err != nil { + return WorkerAllocation{}, err + } if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-worker-capacity', 0))`); err != nil { return WorkerAllocation{}, err }