diff options
Diffstat (limited to 'lib/gitlab/graphql/calls_gitaly')
| -rw-r--r-- | lib/gitlab/graphql/calls_gitaly/field_extension.rb | 87 | ||||
| -rw-r--r-- | lib/gitlab/graphql/calls_gitaly/instrumentation.rb | 40 |
2 files changed, 87 insertions, 40 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 diff --git a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb deleted file mode 100644 index 11d3c50e093..00000000000 --- a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb +++ /dev/null @@ -1,40 +0,0 @@ -# frozen_string_literal: true - -module Gitlab - module Graphql - module CallsGitaly - class Instrumentation - # Check if any `calls_gitaly: true` declarations need to be added - # Do nothing if a constant complexity was provided - def instrument(_type, field) - type_object = field.metadata[:type_class] - return field unless type_object.respond_to?(:calls_gitaly?) - return field if type_object.constant_complexity? || type_object.calls_gitaly? - - old_resolver_proc = field.resolve_proc - - gitaly_wrapped_resolve = -> (typed_object, args, ctx) do - previous_gitaly_call_count = Gitlab::GitalyClient.get_request_count - result = old_resolver_proc.call(typed_object, args, ctx) - current_gitaly_call_count = Gitlab::GitalyClient.get_request_count - calls_gitaly_check(type_object, current_gitaly_call_count - previous_gitaly_call_count) - result - end - - field.redefine do - resolve(gitaly_wrapped_resolve) - end - end - - def calls_gitaly_check(type_object, calls) - return if calls < 1 - - # Will inform you if there needs to be `calls_gitaly: true` as a kwarg in the field declaration - # if there is at least 1 Gitaly call involved with the field resolution. - error = RuntimeError.new("Gitaly is called for field '#{type_object.name}' on #{type_object.owner.try(:name)} - please either specify a constant complexity or add `calls_gitaly: true` to the field declaration") - Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error) - end - end - end - end -end |