Stick newly issued OAuth tokens and grants to the primary database

What does this MR do and why?

A token issued by POST /oauth/token is normally used by the client on its very next request. On GitLab.com, API reads are served from database replicas. A replica may not have received the new oauth_access_tokens row yet when that request arrives.

Then Gitlab::Auth::AuthFinders#find_oauth_access_token raises UnauthorizedError. The client gets a bare 401 with no auth_fail_reason in api_json.log, because the token was not found rather than rejected. Clients treat a 401 on /api/v4/user as a dead session and discard credentials the server would still accept.

We saw this on GitLab.com with the GitLab Mobile app. A refresh_token grant answered 200 at 13:00:56. Within the same second, GET /api/v4/user returned 401 with meta.auth_fail_reason: null. The user was then signed out. Parallel /user and /user/status requests with the same new token, where only one was refused, is the fingerprint of per-request replica selection.

The read side already anticipated this for the API. find_oauth_access_token calls load_balancer_stick_request(::OauthAccessToken, :oauth_token, token, hash_id: true). But Sticking#find_caught_up_replica only routes to the primary when a write location was recorded for that key. Until now, only the Duo Workflows token services recorded one. Nothing recorded it for tokens issued by the token endpoint. The other OAuth lookup paths never checked for one at all.

The fix

This MR adds a concern, Doorkeeper::Concerns::LoadBalancerSticking. It is included in the two Doorkeeper base classes, Doorkeeper::AccessToken and Doorkeeper::AccessGrant (app/models/doorkeeper/). OauthAccessToken and OauthAccessGrant inherit it. The base classes are the right place: the Doorkeeper initializer sets access_token_class but not access_grant_class, so Doorkeeper creates and looks authorization codes up through Doorkeeper::AccessGrant itself, and Geo's logout token and the observability session use Doorkeeper::AccessToken directly.

On create, after the row is committed, it records the primary's current write location in Redis for every plaintext secret the client receives. These are the access token and the refresh token (namespace oauth_token), and the authorization code (namespace oauth_grant). One bulk_stick call per record does this with a single write-location query on the primary. The Redis key is a SHA256 hash of the secret. The value is written with an EVAL of the atomic Lua stick script. The key expires after 30 seconds.

On lookup, it overrides Doorkeeper's find_by_plaintext_token. This is the one class method behind by_token, by_refresh_token, and the grant's by_token, on both base classes and their subclasses. For a secret shaped like a Doorkeeper secret (64 hex characters; personal access tokens, deploy tokens and Git HTTP passwords are skipped by shape), it calls find_caught_up_replica before the query. The request is then served by any replica that has caught up to the recorded location, or by the primary if none has.

Every lookup of an OAuth access token, refresh token or authorization code by its plaintext goes through those methods. So this covers all of these paths:

1. Gitlab::Auth.oauth_access_token_check (Git over HTTP, container registry login, LFS).
2. Doorkeeper's own lookups behind /oauth/userinfo, /oauth/token/info, and /oauth/introspect.
3. Package registry auth through TokenResolver#resolve_oauth_token.
4. The authorization code exchange. The grant is stuck when /oauth/authorize creates it through Doorkeeper::AccessGrant. Doorkeeper's grant lookup and Oauth::TokensController#organization_for_token_request both read it through a caught-up replica.
5. Refresh. The new refresh token is stuck when issued. The presented refresh token is looked up through a caught-up replica. A second refresh inside the lag window no longer gets invalid_grant.
6. /oauth/revoke and Ai::DuoWorkflows::RevokeTokenService. They now find a token issued a moment ago. Before, they silently succeeded on "not found".
7. CreateAutonomousOauthAccessTokenService, Doorkeeper::AccessToken.create! in the observability session, and any other creation path. The model callback covers them.
8. Geo's LogoutToken, which reads through Doorkeeper::AccessToken.by_token.

What is left out

Device grants (oauth_device_grants, a separate Doorkeeper model) are not covered. The device polls the token endpoint every few seconds. A stale read heals itself on the next poll.

Error handling

Both sides rescue StandardError, report with Gitlab::ErrorTracking.track_and_raise_for_dev_exception, and carry on. In production that means Sentry gets the exception and the request proceeds. In development and test, the error is re-raised.

• Create: the row is already committed, so the secret still reaches the client. The report carries the model name and record id.
• Lookup: the query runs as a normal read on whichever replica the request already has. oauth_access_token_check runs before the personal access token, deploy token and password checks, so a Redis outage on the database-load-balancing instance must not fail Git over HTTP, LFS and registry logins. The report carries the model name and the sticking namespace.

Feature flags

The two sides are behind two flags, both gitlab_com_derisk and default_enabled: false, so they can roll out one after the other. The actor of both is Feature.current_request, so a percentage rollout is a percentage of requests and cannot be targeted at users, groups, or projects.

• stick_new_oauth_tokens_to_primary guards the create-side stick. Rollout issue: #632580. On its own it already fixes the API path from the incident above, because the explicit read-side stick in AuthFinders#find_oauth_access_token honours the recorded location without the second flag.
• oauth_lookups_use_load_balancer_sticking guards the lookup-side routing in the models. Rollout issue: #632808. It is what extends the fix to Git over HTTP, the registry, LFS, the package registry, the OIDC endpoints, revoke, refresh and the code exchange. It goes second: without the first flag it finds no recorded location and only adds the Redis GET.

With one flag for both sides, a percentage rollout would only have helped a request whose issuing request was also in the sample (about 1% at 10%). With two flags rolled out in sequence, each step is linear.

Kept on purpose during rollout

Three explicit sticks stay for now. They are the read-side stick in Gitlab::Auth::AuthFinders#find_oauth_access_token and the write-side sticks in the two Duo Workflows services. The read-side one goes with the lookup flag, the write-side ones with the create flag.

While the flags are on, they duplicate one Redis GET (API requests) or one EVAL (Duo tokens) per request. That is harmless, because the Lua script only ever raises the stored position. Removing them while the flags are off would regress the Duo fix from !211007 (merged). The AuthFinders spec accepts find_caught_up_replica being called at least once; the Duo service specs are unchanged, because the model callback uses bulk_stick and their explicit stick stays the only stick call.

Cost when the flags are on

Create side: one EVAL per record issued, covering the access token and its refresh token together, or the authorization code.

Lookup side: one Redis GET on the database-load-balancing Redis per lookup of a 64-hex secret. The Rack load-balancing middleware's own Redis GET runs only for signed-in sessions, not for token-authenticated API or Git requests, so on Git over HTTP, the registry and LFS this GET is added cost, limited by the shape guard to credentials that can be OAuth secrets. While a recorded location exists (30 seconds after a secret is issued) the lookup also reads the WAL position of each replica once. Within one request a secret can be looked up twice (the code exchange's organization_for_token_request and then Doorkeeper; by_token || by_refresh_token on revoke); both lookups probe.

For scale, Oauth::TokensController#create peaks at just under 200 requests per second on gprd. Rails primary SQL peaks at about 30k requests per second. See the endpoint dashboard and the primary SQL dashboard.

Changelog

There is no changelog entry here because both flags are off by default. The Changelog: fixed entry comes with the MR that removes them.

References

• Prior art: !211007 (merged) ("Fix database load balancer for newly created OauthAccessToken"). It introduced the oauth_token sticking namespace, the API read side, and the Duo write side.
• Background for that MR: #579054 (closed)
• Tracking issue for the generic fix implemented here: #579378
• Its refresh grant duplicate: #599181 (closed)
• Rollout issues: #632580 (create side), #632808 (lookup side)
• Code: app/models/doorkeeper/concerns/load_balancer_sticking.rb, included from app/models/doorkeeper/access_token.rb and app/models/doorkeeper/access_grant.rb
• Read side helper: lib/gitlab/auth/auth_finders.rb (find_oauth_access_token)
• Sticking semantics: lib/gitlab/database/load_balancing/sticking.rb

How to set up and validate locally

The race itself needs a lagging replica. It is hard to reproduce on a single-node GDK. The specs cover the sticking behaviour.

The shared examples a Doorkeeper model stuck to the database load balancer run for both base classes. The request spec for POST /oauth/token covers all four grant types. The authorization code flow spec covers the code from POST /oauth/authorize through the exchange.

bin/rspec spec/models/doorkeeper/access_token_spec.rb spec/models/doorkeeper/access_grant_spec.rb \
  spec/requests/oauth/tokens_controller_spec.rb spec/requests/oauth/flows/authorization_code_spec.rb

To try it in a running GDK, run this in a Rails console:

Feature.enable(:stick_new_oauth_tokens_to_primary)
Feature.enable(:oauth_lookups_use_load_balancer_sticking)

MR acceptance checklist

This MR was evaluated against the acceptance checklist. The footprint is one Redis EVAL per record issued plus one GET per lookup of an OAuth-shaped secret. Keys are a hash of the secret. It uses the same mechanism as the existing Duo token sticking. All of it is behind the two feature flags.

🤖 Generated with Claude Code