Wrap proc properly in gitaly call counts

- Add `calls_gitaly: true` to some fields missing (hey, it works!)
- Clarify proc wrapping
- Add kwargs argument to `mount_mutation`

app/graphql/types/base_field.rb

@@ -23,12 +23,11 @@ module Types
 
     def calls_gitaly_check(calls)
       return if @calls_gitaly
+      return if calls < 1
 
       # Will inform you if :calls_gitaly should be true or false based on the number of Gitaly calls
       # involved with the request.
-      if calls > 0
+      raise "Gitaly is called for field '#{name}' #{"on type #{owner.name} " if owner}- please add `calls_gitaly: true` to the field declaration"
-        raise "Gitaly is called for field '#{name}' - please add `calls_gitaly: true` to the field declaration"
-      end
     rescue => e
       Gitlab::Sentry.track_exception(e)
     end

app/graphql/types/merge_request_type.rb

@@ -43,7 +43,7 @@ module Types
     field :allow_collaboration, GraphQL::BOOLEAN_TYPE, null: true
     field :should_be_rebased, GraphQL::BOOLEAN_TYPE, method: :should_be_rebased?, null: false
     field :rebase_commit_sha, GraphQL::STRING_TYPE, null: true
-    field :rebase_in_progress, GraphQL::BOOLEAN_TYPE, method: :rebase_in_progress?, null: false
+    field :rebase_in_progress, GraphQL::BOOLEAN_TYPE, method: :rebase_in_progress?, null: false, calls_gitaly: true
     field :merge_commit_message, GraphQL::STRING_TYPE, method: :default_merge_commit_message, null: true, deprecation_reason: "Renamed to defaultMergeCommitMessage"
     field :default_merge_commit_message, GraphQL::STRING_TYPE, null: true
     field :merge_ongoing, GraphQL::BOOLEAN_TYPE, method: :merge_ongoing?, null: false

app/graphql/types/mutation_type.rb

@@ -9,6 +9,6 @@ module Types
     mount_mutation Mutations::AwardEmojis::Add
     mount_mutation Mutations::AwardEmojis::Remove
     mount_mutation Mutations::AwardEmojis::Toggle
-    mount_mutation Mutations::MergeRequests::SetWip
+    mount_mutation Mutations::MergeRequests::SetWip, calls_gitaly: true
   end
 end

app/graphql/types/permission_types/merge_request.rb

@@ -10,8 +10,8 @@ module Types
       abilities :read_merge_request, :admin_merge_request,
                 :update_merge_request, :create_note
 
-      permission_field :push_to_source_branch, method: :can_push_to_source_branch?
+      permission_field :push_to_source_branch, method: :can_push_to_source_branch?, calls_gitaly: true
-      permission_field :remove_source_branch, method: :can_remove_source_branch?
+      permission_field :remove_source_branch, method: :can_remove_source_branch?, calls_gitaly: true
       permission_field :cherry_pick_on_current_merge_request, method: :can_cherry_pick_on_current_merge_request?
       permission_field :revert_on_current_merge_request, method: :can_revert_on_current_merge_request?
     end

app/graphql/types/project_type.rb

@@ -40,7 +40,7 @@ module Types
     field :lfs_enabled, GraphQL::BOOLEAN_TYPE, null: true
     field :merge_requests_ff_only_enabled, GraphQL::BOOLEAN_TYPE, null: true
 
-    field :avatar_url, GraphQL::STRING_TYPE, null: true, resolve: -> (project, args, ctx) do
+    field :avatar_url, GraphQL::STRING_TYPE, null: true, calls_gitaly: true, resolve: -> (project, args, ctx) do
       project.avatar_url(only_path: false)
     end
 

app/graphql/types/repository_type.rb

@@ -6,9 +6,9 @@ module Types
 
     authorize :download_code
 
-    field :root_ref, GraphQL::STRING_TYPE, null: true
+    field :root_ref, GraphQL::STRING_TYPE, null: true, calls_gitaly: true
-    field :empty, GraphQL::BOOLEAN_TYPE, null: false, method: :empty?
+    field :empty, GraphQL::BOOLEAN_TYPE, null: false, method: :empty?, calls_gitaly: true
     field :exists, GraphQL::BOOLEAN_TYPE, null: false, method: :exists?
-    field :tree, Types::Tree::TreeType, null: true, resolver: Resolvers::TreeResolver
+    field :tree, Types::Tree::TreeType, null: true, resolver: Resolvers::TreeResolver, calls_gitaly: true
   end
 end

lib/gitlab/graphql/calls_gitaly/instrumentation.rb

@@ -10,18 +10,19 @@ module Gitlab
           return field unless type_object && type_object.respond_to?(:calls_gitaly_check)
 
           old_resolver_proc = field.resolve_proc
-          wrapped_proc = gitaly_wrapped_resolve(old_resolver_proc, type_object)
-          field.redefine { resolve(wrapped_proc) }
-        end
 
-        def gitaly_wrapped_resolve(old_resolver_proc, type_object)
+          gitaly_wrapped_resolve = -> (typed_object, args, ctx) do
-          proc do |parent_typed_object, args, ctx|
             previous_gitaly_call_count = Gitlab::GitalyClient.get_request_count
-
+            result = old_resolver_proc.call(typed_object, args, ctx)
-            old_resolver_proc.call(parent_typed_object, args, ctx)
-
             current_gitaly_call_count = Gitlab::GitalyClient.get_request_count
             type_object.calls_gitaly_check(current_gitaly_call_count - previous_gitaly_call_count)
+            result
+          rescue => e
+            ap "#{e.message}"
+          end
+
+          field.redefine do
+            resolve(gitaly_wrapped_resolve)
           end
         end
       end

lib/gitlab/graphql/mount_mutation.rb

@@ -6,11 +6,12 @@ module Gitlab
       extend ActiveSupport::Concern
 
       class_methods do
-        def mount_mutation(mutation_class)
+        def mount_mutation(mutation_class, **custom_kwargs)
           # Using an underscored field name symbol will make `graphql-ruby`
           # standardize the field name
           field mutation_class.graphql_name.underscore.to_sym,
-                mutation: mutation_class
+                mutation: mutation_class,
+                **custom_kwargs
         end
       end
     end