Merge branch 'cache-user-keys-count' into 'master'

Cache the number of user SSH keys

See merge request gitlab-org/gitlab-ce!15401

(cherry picked from commit 66bf283d0af91d916a00f269eab062b19f690c6a)

3e561736 Cache the number of user SSH keys

8 files changed, 81 insertions, 28 deletions

diff --git a/app/models/key.rb b/app/models/key.rb
index f119b15c737..815fd1de909 100644
--- a/app/models/key.rb
+++ b/app/models/key.rb
@@ -27,8 +27,10 @@ class Key < ActiveRecord::Base
 
   after_commit :add_to_shell, on: :create
   after_create :post_create_hook
+  after_create :refresh_user_cache
   after_commit :remove_from_shell, on: :destroy
   after_destroy :post_destroy_hook
+  after_destroy :refresh_user_cache
 
   def key=(value)
     value&.delete!("\n\r")
@@ -76,6 +78,12 @@ class Key < ActiveRecord::Base
     )
   end
 
+  def refresh_user_cache
+    return unless user
+
+    Users::KeysCountService.new(user).refresh_cache
+  end
+
   def post_destroy_hook
     SystemHooksService.new.execute_hooks_for(self, :destroy)
   end
diff --git a/app/models/user.rb b/app/models/user.rb
index f436efd604f..47c78d23591 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -170,6 +170,7 @@ class User < ActiveRecord::Base
   after_save :ensure_namespace_correct
   after_update :username_changed_hook, if: :username_changed?
   after_destroy :post_destroy_hook
+  after_destroy :remove_key_cache
   after_commit :update_emails_with_primary_email, on: :update, if: -> { previous_changes.key?('email') }
   after_commit :update_invalid_gpg_signatures, on: :update, if: -> { previous_changes.key?('email') }
 
@@ -624,7 +625,9 @@ class User < ActiveRecord::Base
   end
 
   def require_ssh_key?
-    keys.count == 0 && Gitlab::ProtocolAccess.allowed?('ssh')
+    count = Users::KeysCountService.new(self).count
+
+    count.zero? && Gitlab::ProtocolAccess.allowed?('ssh')
   end
 
   def require_password_creation?
@@ -886,6 +889,10 @@ class User < ActiveRecord::Base
     system_hook_service.execute_hooks_for(self, :destroy)
   end
 
+  def remove_key_cache
+    Users::KeysCountService.new(self).delete_cache
+  end
+
   def delete_async(deleted_by:, params: {})
     block if params[:hard_delete]
     DeleteUserWorker.perform_async(deleted_by.id, id, params)
diff --git a/app/services/base_count_service.rb b/app/services/base_count_service.rb
new file mode 100644
index 00000000000..99cc9a196e6
--- /dev/null
+++ b/app/services/base_count_service.rb
@@ -0,0 +1,34 @@
+# Base class for services that count a single resource such as the number of
+# issues for a project.
+class BaseCountService
+  def relation_for_count
+    raise(
+      NotImplementedError,
+      '"relation_for_count" must be implemented and return an ActiveRecord::Relation'
+    )
+  end
+
+  def count
+    Rails.cache.fetch(cache_key, raw: raw?) { uncached_count }.to_i
+  end
+
+  def refresh_cache
+    Rails.cache.write(cache_key, uncached_count, raw: raw?)
+  end
+
+  def uncached_count
+    relation_for_count.count
+  end
+
+  def delete_cache
+    Rails.cache.delete(cache_key)
+  end
+
+  def raw?
+    false
+  end
+
+  def cache_key
+    raise NotImplementedError, 'cache_key must be implemented and return a String'
+  end
+end
diff --git a/app/services/projects/count_service.rb b/app/services/projects/count_service.rb
index aa034315280..7e575b2d6f3 100644
--- a/app/services/projects/count_service.rb
+++ b/app/services/projects/count_service.rb
@@ -1,7 +1,7 @@
 module Projects
   # Base class for the various service classes that count project data (e.g.
   # issues or forks).
-  class CountService
+  class CountService < BaseCountService
     # The version of the cache format. This should be bumped whenever the
     # underlying logic changes. This removes the need for explicitly flushing
     # all caches.
@@ -11,29 +11,6 @@ module Projects
       @project = project
     end
 
-    def relation_for_count
-      raise(
-        NotImplementedError,
-        '"relation_for_count" must be implemented and return an ActiveRecord::Relation'
-      )
-    end
-
-    def count
-      Rails.cache.fetch(cache_key) { uncached_count }
-    end
-
-    def refresh_cache
-      Rails.cache.write(cache_key, uncached_count)
-    end
-
-    def uncached_count
-      relation_for_count.count
-    end
-
-    def delete_cache
-      Rails.cache.delete(cache_key)
-    end
-
     def cache_key_name
       raise(
         NotImplementedError,
diff --git a/app/services/projects/forks_count_service.rb b/app/services/projects/forks_count_service.rb
index 3a0fa84b868..d9bdf3a8ad7 100644
--- a/app/services/projects/forks_count_service.rb
+++ b/app/services/projects/forks_count_service.rb
@@ -1,6 +1,6 @@
 module Projects
   # Service class for getting and caching the number of forks of a project.
-  class ForksCountService < CountService
+  class ForksCountService < Projects::CountService
     def relation_for_count
       @project.forks
     end
diff --git a/app/services/projects/open_issues_count_service.rb b/app/services/projects/open_issues_count_service.rb
index 3c0d186a73c..25de97325e2 100644
--- a/app/services/projects/open_issues_count_service.rb
+++ b/app/services/projects/open_issues_count_service.rb
@@ -1,7 +1,7 @@
 module Projects
   # Service class for counting and caching the number of open issues of a
   # project.
-  class OpenIssuesCountService < CountService
+  class OpenIssuesCountService < Projects::CountService
     def relation_for_count
       # We don't include confidential issues in this number since this would
       # expose the number of confidential issues to non project members.
diff --git a/app/services/projects/open_merge_requests_count_service.rb b/app/services/projects/open_merge_requests_count_service.rb
index 2a90f78b90d..77e6448fd5e 100644
--- a/app/services/projects/open_merge_requests_count_service.rb
+++ b/app/services/projects/open_merge_requests_count_service.rb
@@ -1,7 +1,7 @@
 module Projects
   # Service class for counting and caching the number of open merge requests of
   # a project.
-  class OpenMergeRequestsCountService < CountService
+  class OpenMergeRequestsCountService < Projects::CountService
     def relation_for_count
       @project.merge_requests.opened
     end
diff --git a/app/services/users/keys_count_service.rb b/app/services/users/keys_count_service.rb
new file mode 100644
index 00000000000..f82d27eded9
--- /dev/null
+++ b/app/services/users/keys_count_service.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Users
+  # Service class for getting the number of SSH keys that belong to a user.
+  class KeysCountService < BaseCountService
+    attr_reader :user
+
+    # user - The User for which to get the number of SSH keys.
+    def initialize(user)
+      @user = user
+    end
+
+    def relation_for_count
+      user.keys
+    end
+
+    def raw?
+      # Since we're storing simple integers we don't need all of the additional
+      # Marshal data Rails includes by default.
+      true
+    end
+
+    def cache_key
+      "users/key-count-service/#{user.id}"
+    end
+  end
+end