Add RequestStoreWrap to cache via RequestStore

I don't like the idea of `RequestStore` at all, because it's just a
global state which shouldn't be used at all. But we have a number of
places calling `ProtectedBranch.protected?` and `ProtectedTag.protected?`
in a loop for the same user, project, and ref whenever we're checking
against if the jobs for a given pipeline is accessible for a given user.
This means we're effectively making N queries for the same thing over
and over.

To properly fix this, we need to change how we check the permission,
and that could be a huge work. To solve this quickly, adding a cache
layer for the given request would be quite simple to do.

We're already doing this in Commit#author, and this is extending that
idea and make it generalized.

3 files changed, 183 insertions, 4 deletions

diff --git a/lib/gitlab/cache/request_store_wrap.rb b/lib/gitlab/cache/request_store_wrap.rb
new file mode 100644
index 00000000000..3e0a5f06b53
--- /dev/null
+++ b/lib/gitlab/cache/request_store_wrap.rb
@@ -0,0 +1,60 @@
+module Gitlab
+  module Cache
+    # This module provides a simple way to cache values in RequestStore,
+    # and the cache key would be based on the class name, method name,
+    # customized instance level values, and arguments.
+    #
+    # A simple example:
+    #
+    # class UserAccess
+    #   extend Gitlab::Cache::RequestStoreWrap
+    #
+    #   request_store_wrap_key do
+    #     [user.id, project.id]
+    #   end
+    #
+    #   request_store_wrap def can_push_to_branch?(ref)
+    #     # ...
+    #   end
+    # end
+    #
+    # This way, the result of `can_push_to_branch?` would be cached in
+    # `RequestStore.store` based on the cache key.
+    module RequestStoreWrap
+      def self.extended(klass)
+        return if klass < self
+
+        extension = Module.new
+        klass.const_set(:RequestStoreWrapExtension, extension)
+        klass.prepend(extension)
+      end
+
+      def request_store_wrap_key(&block)
+        if block_given?
+          @request_store_wrap_key = block
+        else
+          @request_store_wrap_key
+        end
+      end
+
+      def request_store_wrap(method_name)
+        const_get(:RequestStoreWrapExtension)
+          .send(:define_method, method_name) do |*args|
+            return super(*args) unless RequestStore.active?
+
+            klass = self.class
+            key = [klass.name,
+                   method_name,
+                   *instance_exec(&klass.request_store_wrap_key),
+                   *args].join(':')
+
+            if RequestStore.store.key?(key)
+              RequestStore.store[key]
+            else
+              RequestStore.store[key] = super(*args)
+            end
+          end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/user_access.rb b/lib/gitlab/user_access.rb
index 3b922da7ced..9c066627011 100644
--- a/lib/gitlab/user_access.rb
+++ b/lib/gitlab/user_access.rb
@@ -1,5 +1,11 @@
 module Gitlab
   class UserAccess
+    extend Gitlab::Cache::RequestStoreWrap
+
+    request_store_wrap_key do
+      [user&.id, project&.id]
+    end
+
     attr_reader :user, :project
 
     def initialize(user, project: nil)
@@ -28,7 +34,7 @@ module Gitlab
       true
     end
 
-    def can_create_tag?(ref)
+    request_store_wrap def can_create_tag?(ref)
       return false unless can_access_git?
 
       if ProtectedTag.protected?(project, ref)
@@ -38,7 +44,7 @@ module Gitlab
       end
     end
 
-    def can_delete_branch?(ref)
+    request_store_wrap def can_delete_branch?(ref)
       return false unless can_access_git?
 
       if ProtectedBranch.protected?(project, ref)
@@ -48,7 +54,7 @@ module Gitlab
       end
     end
 
-    def can_push_to_branch?(ref)
+    request_store_wrap def can_push_to_branch?(ref)
       return false unless can_access_git?
 
       if ProtectedBranch.protected?(project, ref)
@@ -60,7 +66,7 @@ module Gitlab
       end
     end
 
-    def can_merge_to_branch?(ref)
+    request_store_wrap def can_merge_to_branch?(ref)
       return false unless can_access_git?
 
       if ProtectedBranch.protected?(project, ref)
diff --git a/spec/lib/gitlab/cache/request_store_wrap_spec.rb b/spec/lib/gitlab/cache/request_store_wrap_spec.rb
new file mode 100644
index 00000000000..6a95239066b
--- /dev/null
+++ b/spec/lib/gitlab/cache/request_store_wrap_spec.rb
@@ -0,0 +1,113 @@
+require 'spec_helper'
+
+describe Gitlab::Cache::RequestStoreWrap, :request_store do
+  let(:klass) do
+    Class.new do
+      extend Gitlab::Cache::RequestStoreWrap
+
+      attr_accessor :id, :name, :result
+
+      def self.name
+        'ExpensiveAlgorithm'
+      end
+
+      def initialize(id, name, result)
+        self.id = id
+        self.name = name
+        self.result = result
+      end
+
+      request_store_wrap_key do
+        [id, name]
+      end
+
+      request_store_wrap def compute(arg)
+        result << arg
+      end
+
+      request_store_wrap def repute(arg)
+        result << arg
+      end
+    end
+  end
+
+  let(:algorithm) { klass.new('id', 'name', []) }
+
+  context 'when RequestStore is active' do
+    it 'does not compute twice for the same argument' do
+      result = algorithm.compute(true)
+
+      expect(result).to eq([true])
+      expect(algorithm.compute(true)).to eq(result)
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice for the different argument' do
+      algorithm.compute(true)
+      result = algorithm.compute(false)
+
+      expect(result).to eq([true, false])
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice for the different keys, id' do
+      algorithm.compute(true)
+      algorithm.id = 'ad'
+      result = algorithm.compute(true)
+
+      expect(result).to eq([true, true])
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice for the different keys, name' do
+      algorithm.compute(true)
+      algorithm.name = 'same'
+      result = algorithm.compute(true)
+
+      expect(result).to eq([true, true])
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice for the different class name' do
+      algorithm.compute(true)
+      allow(klass).to receive(:name).and_return('CheapAlgo')
+      result = algorithm.compute(true)
+
+      expect(result).to eq([true, true])
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice for the different method' do
+      algorithm.compute(true)
+      result = algorithm.repute(true)
+
+      expect(result).to eq([true, true])
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice if RequestStore starts over' do
+      algorithm.compute(true)
+      RequestStore.end!
+      RequestStore.clear!
+      RequestStore.begin!
+      result = algorithm.compute(true)
+
+      expect(result).to eq([true, true])
+      expect(algorithm.result).to eq(result)
+    end
+  end
+
+  context 'when RequestStore is inactive' do
+    before do
+      RequestStore.end!
+    end
+
+    it 'computes twice even if everything is the same' do
+      algorithm.compute(true)
+      result = algorithm.compute(true)
+
+      expect(result).to eq([true, true])
+      expect(algorithm.result).to eq(result)
+    end
+  end
+end