Add Gitlab::Utils::StrongMemoize

4 files changed, 93 insertions, 6 deletions

diff --git a/app/services/ci/pipeline_trigger_service.rb b/app/services/ci/pipeline_trigger_service.rb
index 120af8c1e61..a9813d774bb 100644
--- a/app/services/ci/pipeline_trigger_service.rb
+++ b/app/services/ci/pipeline_trigger_service.rb
@@ -1,5 +1,7 @@
 module Ci
   class PipelineTriggerService < BaseService
+    include Gitlab::Utils::StrongMemoize
+
     def execute
       if trigger_from_token
         create_pipeline_from_trigger(trigger_from_token)
@@ -26,9 +28,9 @@ module Ci
     end
 
     def trigger_from_token
-      return @trigger if defined?(@trigger)
-
-      @trigger = Ci::Trigger.find_by_token(params[:token].to_s)
+      strong_memoize(:trigger) do
+        Ci::Trigger.find_by_token(params[:token].to_s)
+      end
     end
 
     def create_pipeline_variables!(pipeline)
diff --git a/lib/api/api_guard.rb b/lib/api/api_guard.rb
index b9c7d443f6c..c1c0d344917 100644
--- a/lib/api/api_guard.rb
+++ b/lib/api/api_guard.rb
@@ -42,6 +42,8 @@ module API
 
     # Helper Methods for Grape Endpoint
     module HelperMethods
+      include Gitlab::Utils::StrongMemoize
+
       def find_current_user!
         user = find_user_from_access_token || find_user_from_warden
         return unless user
@@ -52,9 +54,9 @@ module API
       end
 
       def access_token
-        return @access_token if defined?(@access_token)
-
-        @access_token = find_oauth_access_token || find_personal_access_token
+        strong_memoize(:access_token) do
+          find_oauth_access_token || find_personal_access_token
+        end
       end
 
       def validate_access_token!(scopes: [])
diff --git a/lib/gitlab/utils/strong_memoize.rb b/lib/gitlab/utils/strong_memoize.rb
new file mode 100644
index 00000000000..a2ac9285b56
--- /dev/null
+++ b/lib/gitlab/utils/strong_memoize.rb
@@ -0,0 +1,31 @@
+module Gitlab
+  module Utils
+    module StrongMemoize
+      # Instead of writing patterns like this:
+      #
+      #     def trigger_from_token
+      #       return @trigger if defined?(@trigger)
+      #
+      #       @trigger = Ci::Trigger.find_by_token(params[:token].to_s)
+      #     end
+      #
+      # We could write it like:
+      #
+      #     def trigger_from_token
+      #       strong_memoize(:trigger) do
+      #         Ci::Trigger.find_by_token(params[:token].to_s)
+      #       end
+      #     end
+      #
+      def strong_memoize(name)
+        ivar_name = "@#{name}"
+
+        if instance_variable_defined?(ivar_name)
+          instance_variable_get(ivar_name)
+        else
+          instance_variable_set(ivar_name, yield)
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/utils/strong_memoize_spec.rb b/spec/lib/gitlab/utils/strong_memoize_spec.rb
new file mode 100644
index 00000000000..4a104ab6d97
--- /dev/null
+++ b/spec/lib/gitlab/utils/strong_memoize_spec.rb
@@ -0,0 +1,52 @@
+require 'spec_helper'
+
+describe Gitlab::Utils::StrongMemoize do
+  let(:klass) do
+    struct = Struct.new(:value) do
+      def method_name
+        strong_memoize(:method_name) do
+          trace << value
+          value
+        end
+      end
+
+      def trace
+        @trace ||= []
+      end
+    end
+
+    struct.include(described_class)
+    struct
+  end
+
+  subject(:object) { klass.new(value) }
+
+  shared_examples 'caching the value' do
+    it 'only calls the block once' do
+      value0 = object.method_name
+      value1 = object.method_name
+
+      expect(value0).to eq(value)
+      expect(value1).to eq(value)
+      expect(object.trace).to contain_exactly(value)
+    end
+
+    it 'returns and defines the instance variable for the exact value' do
+      returned_value = object.method_name
+      memoized_value = object.instance_variable_get(:@method_name)
+
+      expect(returned_value).to eql(value)
+      expect(memoized_value).to eql(value)
+    end
+  end
+
+  describe '#strong_memoize' do
+    [nil, false, true, 'value', 0, [0]].each do |value|
+      context "with value #{value}" do
+        let(:value) { value }
+
+        it_behaves_like 'caching the value'
+      end
+    end
+  end
+end