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_suggestions ability, in a before block of API::CodeSuggestions, then in the project or group policy, then again in the task's unit primitive (complete_code or generate_code).
  • Each check calls Ai::UserAuthorizable#allowed_to_use?, which calls duo_entitlement. The /completions endpoint runs on every keystroke, so this adds up fast.
  • Master already memoizes allowed_to_use per 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_purchases now keeps its result in Gitlab::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, and generate_code share 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_use memo key now also includes the project and group from Gitlab::ApplicationContext. When there is no root_namespace, governing_namespace reads 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_use caches is now frozen, including its namespace_ids and namespace_ids_from_all_addons arrays. 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::SafeRequestStore and not instance variables, because nothing should be kept between requests, Sidekiq jobs get their own store, and specs that share one let_it_be user are not affected: the store is only active with the :request_store tag.

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

Screenshots or screen recordings

This is a backend change with no UI. There are no screenshots.

How to set up and validate locally

  1. 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]'
  2. Start GDK with SaaS simulation turned on.

  3. 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
  4. Optional: call /api/v4/code_suggestions/completions with a personal access token, then check log/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_suggestions examples in ee/spec/policies/project_policy_spec.rb, ee/spec/policies/ee/group_policy_spec.rb, and ee/spec/policies/global_policy_spec.rb, plus ee/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.

Edited by Vitali Tatarintev

Merge request reports

Loading
Loading