diff options
Diffstat (limited to 'lib/gitlab/graphql/calls_gitaly/field_extension.rb')
-rw-r--r-- | lib/gitlab/graphql/calls_gitaly/field_extension.rb | 87 |
1 files changed, 87 insertions, 0 deletions
diff --git a/lib/gitlab/graphql/calls_gitaly/field_extension.rb b/lib/gitlab/graphql/calls_gitaly/field_extension.rb new file mode 100644 index 00000000000..32530b47ce3 --- /dev/null +++ b/lib/gitlab/graphql/calls_gitaly/field_extension.rb @@ -0,0 +1,87 @@ +# frozen_string_literal: true + +module Gitlab + module Graphql + module CallsGitaly + # Check if any `calls_gitaly: true` declarations need to be added + # + # See BaseField: this extension is not applied if the field does not + # need it (i.e. it has a constant complexity or knows that it calls + # gitaly) + class FieldExtension < ::GraphQL::Schema::FieldExtension + include Laziness + + def resolve(object:, arguments:, **rest) + yield(object, arguments, [current_gitaly_call_count, accounted_for]) + end + + def after_resolve(value:, memo:, **rest) + (value, count) = value_with_count(value, memo) + calls_gitaly_check(count) + accounted_for(count) + + value + end + + private + + # Resolutions are not nested nicely (due to laziness), so we have to + # know not just how many calls were made before resolution started, but + # also how many were accounted for by fields with the correct settings + # in between. + # + # e.g. the following is not just plausible, but common: + # + # enter A.user (lazy) + # enter A.x + # leave A.x + # enter A.calls_gitaly + # leave A.calls_gitaly (accounts for 1 call) + # leave A.user + # + # In this circumstance we need to mark the calls made by A.calls_gitaly + # as accounted for, even though they were made after we yielded + # in A.user + def value_with_count(value, (previous_count, previous_accounted_for)) + newly_accounted_for = accounted_for - previous_accounted_for + value = force(value) + count = [current_gitaly_call_count - (previous_count + newly_accounted_for), 0].max + + [value, count] + end + + def current_gitaly_call_count + Gitlab::GitalyClient.get_request_count || 0 + end + + def calls_gitaly_check(calls) + return if calls < 1 || field.may_call_gitaly? + + error = RuntimeError.new(<<~ERROR) + #{field_name} unexpectedly calls Gitaly! + + Please either specify a constant complexity or add `calls_gitaly: true` + to the field declaration + ERROR + Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error) + end + + def accounted_for(count = nil) + return 0 unless Gitlab::SafeRequestStore.active? + + Gitlab::SafeRequestStore["graphql_gitaly_accounted_for"] ||= 0 + + if count.nil? + Gitlab::SafeRequestStore["graphql_gitaly_accounted_for"] + else + Gitlab::SafeRequestStore["graphql_gitaly_accounted_for"] += count + end + end + + def field_name + "#{field.owner.graphql_name}.#{field.graphql_name}" + end + end + end + end +end |