Cache the number of user SSH keys

By caching the number of personal SSH keys we reduce the number of
queries necessary on pages such as ProjectsController#show (which can
end up querying this data multiple times).

The cache is refreshed/flushed whenever an SSH key is added, removed, or
when a user is removed.

13 files changed, 262 insertions, 29 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 ea10e2854d6..be8112749bf 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
diff --git a/changelogs/unreleased/cache-user-keys-count.yml b/changelogs/unreleased/cache-user-keys-count.yml
new file mode 100644
index 00000000000..181be95622c
--- /dev/null
+++ b/changelogs/unreleased/cache-user-keys-count.yml
@@ -0,0 +1,5 @@
+---
+title: Cache the number of user SSH keys
+merge_request:
+author:
+type: performance
diff --git a/spec/models/key_spec.rb b/spec/models/key_spec.rb
index 81c2057e175..4cd9e3f4f1d 100644
--- a/spec/models/key_spec.rb
+++ b/spec/models/key_spec.rb
@@ -166,4 +166,27 @@ describe Key, :mailer do
       expect(key.public_key.key_text).to eq(valid_key)
     end
   end
+
+  describe '#refresh_user_cache', :use_clean_rails_memory_store_caching do
+    context 'when the key belongs to a user' do
+      it 'refreshes the keys count cache for the user' do
+        expect_any_instance_of(Users::KeysCountService)
+          .to receive(:refresh_cache)
+          .and_call_original
+
+        key = create(:personal_key)
+
+        expect(Users::KeysCountService.new(key.user).count).to eq(1)
+      end
+    end
+
+    context 'when the key does not belong to a user' do
+      it 'does nothing' do
+        expect_any_instance_of(Users::KeysCountService)
+          .not_to receive(:refresh_cache)
+
+        create(:key)
+      end
+    end
+  end
 end
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index 88732962071..86647ddf6ce 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -828,7 +828,7 @@ describe User do
       end
     end
 
-    describe '#require_ssh_key?' do
+    describe '#require_ssh_key?', :use_clean_rails_memory_store_caching do
       protocol_and_expectation = {
         'http' => false,
         'ssh' => true,
@@ -843,6 +843,12 @@ describe User do
           expect(user.require_ssh_key?).to eq(expected)
         end
       end
+
+      it 'returns false when the user has 1 or more SSH keys' do
+        key = create(:personal_key)
+
+        expect(key.user.require_ssh_key?).to eq(false)
+      end
     end
   end
 
diff --git a/spec/services/base_count_service_spec.rb b/spec/services/base_count_service_spec.rb
new file mode 100644
index 00000000000..5ec8ed0976d
--- /dev/null
+++ b/spec/services/base_count_service_spec.rb
@@ -0,0 +1,80 @@
+require 'spec_helper'
+
+describe BaseCountService, :use_clean_rails_memory_store_caching do
+  let(:service) { described_class.new }
+
+  describe '#relation_for_count' do
+    it 'raises NotImplementedError' do
+      expect { service.relation_for_count }.to raise_error(NotImplementedError)
+    end
+  end
+
+  describe '#count' do
+    it 'returns the number of values' do
+      expect(service)
+        .to receive(:cache_key)
+        .and_return('foo')
+
+      expect(service)
+        .to receive(:uncached_count)
+        .and_return(5)
+
+      expect(service.count).to eq(5)
+    end
+  end
+
+  describe '#uncached_count' do
+    it 'returns the uncached number of values' do
+      expect(service)
+        .to receive(:relation_for_count)
+        .and_return(double(:relation, count: 5))
+
+      expect(service.uncached_count).to eq(5)
+    end
+  end
+
+  describe '#refresh_cache' do
+    it 'refreshes the cache' do
+      allow(service)
+        .to receive(:cache_key)
+        .and_return('foo')
+
+      allow(service)
+        .to receive(:uncached_count)
+        .and_return(4)
+
+      service.refresh_cache
+
+      expect(Rails.cache.fetch(service.cache_key, raw: service.raw?)).to eq(4)
+    end
+  end
+
+  describe '#delete_cache' do
+    it 'deletes the cache' do
+      allow(service)
+        .to receive(:cache_key)
+        .and_return('foo')
+
+      allow(service)
+        .to receive(:uncached_count)
+        .and_return(4)
+
+      service.refresh_cache
+      service.delete_cache
+
+      expect(Rails.cache.fetch(service.cache_key, raw: service.raw?)).to be_nil
+    end
+  end
+
+  describe '#raw?' do
+    it 'returns false' do
+      expect(service.raw?).to eq(false)
+    end
+  end
+
+  describe '#cache_key' do
+    it 'raises NotImplementedError' do
+      expect { service.cache_key }.to raise_error(NotImplementedError)
+    end
+  end
+end
diff --git a/spec/services/users/keys_count_service_spec.rb b/spec/services/users/keys_count_service_spec.rb
new file mode 100644
index 00000000000..a188cf86772
--- /dev/null
+++ b/spec/services/users/keys_count_service_spec.rb
@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Users::KeysCountService, :use_clean_rails_memory_store_caching do
+  let(:user) { create(:user) }
+  let(:service) { described_class.new(user) }
+
+  describe '#count' do
+    before do
+      create(:personal_key, user: user)
+    end
+
+    it 'returns the number of SSH keys as an Integer' do
+      expect(service.count).to eq(1)
+    end
+
+    it 'caches the number of keys in Redis' do
+      service.delete_cache
+
+      recorder = ActiveRecord::QueryRecorder.new do
+        2.times { service.count }
+      end
+
+      expect(recorder.count).to eq(1)
+    end
+  end
+
+  describe '#refresh_cache' do
+    it 'refreshes the Redis cache' do
+      Rails.cache.write(service.cache_key, 10)
+      service.refresh_cache
+
+      expect(Rails.cache.fetch(service.cache_key, raw: true)).to be_zero
+    end
+  end
+
+  describe '#delete_cache' do
+    it 'removes the cache' do
+      service.count
+      service.delete_cache
+
+      expect(Rails.cache.fetch(service.cache_key, raw: true)).to be_nil
+    end
+  end
+
+  describe '#uncached_count' do
+    it 'returns the number of SSH keys' do
+      expect(service.uncached_count).to be_zero
+    end
+
+    it 'does not cache the number of keys' do
+      recorder = ActiveRecord::QueryRecorder.new do
+        2.times { service.uncached_count }
+      end
+
+      expect(recorder.count).to be > 0
+    end
+  end
+
+  describe '#cache_key' do
+    it 'returns the cache key' do
+      expect(service.cache_key).to eq("users/key-count-service/#{user.id}")
+    end
+  end
+end