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