Introduce Gitlab::Cache::RequestStoreWrap

So that we cache the result of UserAccess#can_push_or_merge_to_branch?
in RequestStore, avoiding querying ProtectedBranch over and over for
the list of pipelines (i.e. in PipelineSerializer)

I don't think this is ideal because I don't like the idea of
RequestStore in general, but this is the easiest way to cache it
without changing the architecture. In the future we should cache
more explicitly rather than this kind of global store.

2 files changed, 68 insertions, 2 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 bb05c474fa2..d8b043f5021 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)
@@ -52,7 +58,7 @@ module Gitlab
       can_push_to_branch?(ref) || can_merge_to_branch?(ref)
     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)
@@ -64,7 +70,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)