Convert calls-gitaly graphql tooling to field extension

The calls-gitaly tooling is not available when calling
`GraphqlHelpers#resolve` and `#resolve_field`, because it was
implemented as schema instrumentation which must be called from
`Schema.execute`. This changes to use field extensions that are:

- simpler (less code, easier to read)
- can be applied selectively (here we can avoid the runtime cost in
  prod)
- not deprecated; to use the intepreter, we will need to make this move

Changes:
 - Adds `#may_call_gitaly` method to BaseField

app/graphql/gitlab_schema.rb

@@ -14,7 +14,6 @@ class GitlabSchema < GraphQL::Schema
   use BatchLoader::GraphQL
   use Gitlab::Graphql::Authorize
   use Gitlab::Graphql::Present
-  use Gitlab::Graphql::CallsGitaly
   use Gitlab::Graphql::Pagination::Connections
   use Gitlab::Graphql::GenericTracing
   use Gitlab::Graphql::Timeout, max_seconds: Gitlab.config.gitlab.graphql_timeout

app/graphql/types/base_field.rb

@@ -19,6 +19,15 @@ module Types
       kwargs = gitlab_deprecation(kwargs)
 
       super(**kwargs, &block)
+
+      # We want to avoid the overhead of this in prod
+      unless Rails.env.production?
+        extension ::Gitlab::Graphql::CallsGitaly::FieldExtension
+      end
+    end
+
+    def may_call_gitaly?
+      @constant_complexity || @calls_gitaly
     end
 
     def requires_argument?

lib/gitlab/graphql/calls_gitaly.rb

-# frozen_string_literal: true
-
-module Gitlab
-  module Graphql
-    # Wraps the field resolution to count Gitaly calls before and after.
-    # Raises an error if the field calls Gitaly but hadn't declared such.
-    module CallsGitaly
-      extend ActiveSupport::Concern
-
-      def self.use(schema_definition)
-        schema_definition.instrument(:field, Gitlab::Graphql::CallsGitaly::Instrumentation.new, after_built_ins: true)
-      end
-    end
-  end
-end

lib/gitlab/graphql/calls_gitaly/field_extension.rb

+# 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_type.graphql_name}.#{field.graphql_name}"
+        end
+      end
+    end
+  end
+end

lib/gitlab/graphql/calls_gitaly/instrumentation.rb

-# 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

spec/graphql/gitlab_schema_spec.rb

@@ -22,10 +22,6 @@ RSpec.describe GitlabSchema do
     expect(field_instrumenters).to include(instance_of(::Gitlab::Graphql::Present::Instrumentation))
   end
 
-  it 'enables using gitaly call checker' do
-    expect(field_instrumenters).to include(instance_of(::Gitlab::Graphql::CallsGitaly::Instrumentation))
-  end
-
   it 'has the base mutation' do
     expect(described_class.mutation).to eq(::Types::MutationType)
   end

spec/lib/gitlab/graphql/calls_gitaly/field_extension_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::CallsGitaly::FieldExtension, :request_store do
+  include GraphqlHelpers
+
+  let(:field_args) { {} }
+  let(:owner) { fresh_object_type }
+  let(:field) do
+    ::Types::BaseField.new(name: 'value', type: GraphQL::STRING_TYPE, null: true, owner: owner, **field_args)
+  end
+
+  def resolve_value
+    resolve_field(field, { value: 'foo' }, object_type: owner)
+  end
+
+  context 'when the field calls gitaly' do
+    before do
+      owner.define_method :value do
+        Gitlab::SafeRequestStore['gitaly_call_actual'] = 1
+        'fresh-from-the-gitaly-mines!'
+      end
+    end
+
+    context 'when the field has a constant complexity' do
+      let(:field_args) { { complexity: 100 } }
+
+      it 'allows the call' do
+        expect { resolve_value }.not_to raise_error
+      end
+    end
+
+    context 'when the field declares that it calls gitaly' do
+      let(:field_args) { { calls_gitaly: true } }
+
+      it 'allows the call' do
+        expect { resolve_value }.not_to raise_error
+      end
+    end
+
+    context 'when the field does not have these arguments' do
+      let(:field_args) { {} }
+
+      it 'notices, and raises, mentioning the field' do
+        expect { resolve_value }.to raise_error(include('Object.value'))
+      end
+    end
+  end
+
+  context 'when it does not call gitaly' do
+    let(:field_args) { {} }
+
+    it 'does not raise' do
+      value = resolve_value
+
+      expect(value).to eq 'foo'
+    end
+  end
+
+  context 'when some field calls gitaly while we were waiting' do
+    let(:extension) { described_class.new(field: field, options: {}) }
+
+    it 'is acceptable if all are accounted for' do
+      object = :anything
+      arguments = :any_args
+
+      ::Gitlab::SafeRequestStore['gitaly_call_actual'] = 3
+      ::Gitlab::SafeRequestStore['graphql_gitaly_accounted_for'] = 0
+
+      expect do |b|
+        extension.resolve(object: object, arguments: arguments, &b)
+      end.to yield_with_args(object, arguments, [3, 0])
+
+      ::Gitlab::SafeRequestStore['gitaly_call_actual'] = 13
+      ::Gitlab::SafeRequestStore['graphql_gitaly_accounted_for'] = 10
+
+      expect { extension.after_resolve(value: 'foo', memo: [3, 0]) }.not_to raise_error
+    end
+
+    it 'is unacceptable if some of the calls are unaccounted for' do
+      ::Gitlab::SafeRequestStore['gitaly_call_actual'] = 10
+      ::Gitlab::SafeRequestStore['graphql_gitaly_accounted_for'] = 9
+
+      expect { extension.after_resolve(value: 'foo', memo: [0, 0]) }.to raise_error(include('Object.value'))
+    end
+  end
+end

spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb

-# frozen_string_literal: true
-require 'spec_helper'
-
-RSpec.describe Gitlab::Graphql::CallsGitaly::Instrumentation do
-  subject { described_class.new }
-
-  describe '#calls_gitaly_check' do
-    let(:gitaly_field) { Types::BaseField.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true) }
-    let(:no_gitaly_field) { Types::BaseField.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: false) }
-
-    context 'if there are no Gitaly calls' do
-      it 'does not raise an error if calls_gitaly is false' do
-        expect { subject.send(:calls_gitaly_check, no_gitaly_field, 0) }.not_to raise_error
-      end
-    end
-
-    context 'if there is at least 1 Gitaly call' do
-      it 'raises an error if calls_gitaly: is false or not defined' do
-        expect { subject.send(:calls_gitaly_check, no_gitaly_field, 1) }.to raise_error(/specify a constant complexity or add `calls_gitaly: true`/)
-      end
-    end
-  end
-end