diff --git a/lib/gitlab/cache/request_store_wrap.rb b/lib/gitlab/cache/request_store_wrap.rb new file mode 100644 index 0000000000000000000000000000000000000000..3e0a5f06b53788a0723c314e7f082cbde5d18c06 --- /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 3b922da7cedf9c832b45b0995982a2b8f48d12b8..9c066627011ca9014863aaed09527e2089003315 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 0000000000000000000000000000000000000000..6a95239066b32944820f6954e04cd24323c84570 --- /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