Share the Duo seat lookup within a request
What does this MR do and why?
- Every request to the Code Suggestions API checks the Duo entitlement more than once. It happens in the
:access_code_suggestionsability, in abeforeblock ofAPI::CodeSuggestions, then in the project or group policy, then again in the task's unit primitive (complete_codeorgenerate_code). - Each check calls
Ai::UserAuthorizable#allowed_to_use?, which callsduo_entitlement. The/completionsendpoint runs on every keystroke, so this adds up fast. - Master already memoizes
allowed_to_useper request, from a merged change for work item 629337. That memo only helps calls with the exact same arguments. - The three Code Suggestions checks use different unit primitives, and after merge request 256909 also different root namespaces. So the seat lookup (the add-on purchase queries in
check_add_on_purchases) still ran once per check. - Change 1:
check_add_on_purchasesnow keeps its result inGitlab::SafeRequestStore, keyed by the add-on names. The lookup depends only on the add-ons, not on the root namespace or the unit primitive.code_suggestions,complete_code, andgenerate_codeshare the same add-on list, so one lookup can serve all checks in a request. A nil result, for a user without a seat, is cached too. - Change 2: the
allowed_to_usememo key now also includes the project and group fromGitlab::ApplicationContext. When there is noroot_namespace,governing_namespacereads the project or group from that context, and in controllers the context can change during one request. Without this fix, a result from early in the request could be reused after the context changed, and namespace access could be checked against the wrong namespace. - Change 3: the response that
allowed_to_usecaches is now frozen, including itsnamespace_idsandnamespace_ids_from_all_addonsarrays. All callers in a request share one object, so an in-place change by one caller would corrupt it for the others. No caller changes it today. GitLab Duo Code Review suggested this. - We use
Gitlab::SafeRequestStoreand not instance variables, because nothing should be kept between requests, Sidekiq jobs get their own store, and specs that share onelet_it_beuser are not affected: the store is only active with the:request_storetag.
Behavior: no change. A result depends only on the user, the arguments, and the application context, and all of them are now in the cache key. The only case where a cached result could go stale is a seat assigned and checked in the same request or Sidekiq job. No code does that today.
Relation to the merged change: the first version of this MR memoized duo_entitlement itself. While it was in review, the change for work item 629337 merged with its own memo on allowed_to_use. This MR was reworked on top of that change: the duo_entitlement memo was dropped, and the application context fix and the freeze moved onto the merged allowed_to_use memo instead.
In the API request spec, one /completions request used to run the add-on lookup three times. Now it runs once.
References
- Issue: #628396 (closed)
- Review comment that asked for this change: !251601 (comment 3811721385)
- Merged change that added the
allowed_to_usememo: #629337 (closed) - Previous MR that moved the check into policies: !251601 (merged)
- Related follow-up MR: !256909 (merged)
Screenshots or screen recordings
This is a backend change with no UI. There are no screenshots.
How to set up and validate locally
-
Seed a Duo group and a seat for
root. Run this from the Rails root:GITLAB_SIMULATE_SAAS=1 bundle exec 'rake gitlab:duo:setup[duo_pro]' -
Start GDK with SaaS simulation turned on.
-
Open a Rails console and run this. The request store is only active inside a request, so the block turns it on. The first call runs the seat query. The second call uses different arguments, but it does not run a new seat query, because the add-on lookup is shared:
user = User.find_by_username('root') ActiveRecord::Base.logger = Logger.new($stdout) Gitlab::SafeRequestStore.ensure_request_store do user.allowed_to_use?(:code_suggestions) # runs the seat query user.allowed_to_use?(:code_suggestions, unit_primitive_name: :complete_code) # no new seat query end -
Optional: call
/api/v4/code_suggestions/completionswith a personal access token, then checklog/development.log. The add-on purchase query should appear only once per request.
Tests
ee/spec/models/concerns/ai/user_authorizable_spec.rb: new examples check that a different root namespace runs the namespace check again but the seat lookup only once, that a user without a seat triggers the seat lookup only once, that the namespace check runs again when the application context changes, that the cached response and its arrays are frozen, and that nothing is kept without a request store.ee/spec/requests/api/code_suggestions_spec.rb: one new example checks that a completions request looks up the add-on seats once, using the real lookup.- Also run: the
access_code_suggestionsexamples inee/spec/policies/project_policy_spec.rb,ee/spec/policies/ee/group_policy_spec.rb, andee/spec/policies/global_policy_spec.rb, plusee/spec/lib/users/identity_verification/authorize_duo_spec.rb,ee/spec/lib/gitlab/ai_gateway_spec.rb, and the specs added by the merged change for work item 629337. - All passing.
MR acceptance checklist
Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.