Merge branch 'optimise-strong-memoize' into 'master'

Reduce memory allocations for `StrongMemoize`

See merge request gitlab-org/gitlab!81920

Gemfile

@@ -420,6 +420,7 @@ group :test do
   gem 'fuubar', '~> 2.2.0'
   gem 'rspec-retry', '~> 0.6.1'
   gem 'rspec_profiling', '~> 0.0.6'
+  gem 'rspec-benchmark', '~> 0.6.0'
   gem 'rspec-parameterized', require: false
 
   gem 'capybara', '~> 3.35.3'

Gemfile.lock

@@ -139,8 +139,11 @@ GEM
     bcrypt (3.1.16)
     benchmark (0.1.1)
     benchmark-ips (2.3.0)
+    benchmark-malloc (0.2.0)
     benchmark-memory (0.1.2)
       memory_profiler (~> 0.9)
+    benchmark-perf (0.6.0)
+    benchmark-trend (0.4.0)
     better_errors (2.9.1)
       coderay (>= 1.0.0)
       erubi (>= 1.0.0)
@@ -1062,6 +1065,11 @@ GEM
       rspec-core (~> 3.10.0)
       rspec-expectations (~> 3.10.0)
       rspec-mocks (~> 3.10.0)
+    rspec-benchmark (0.6.0)
+      benchmark-malloc (~> 0.2)
+      benchmark-perf (~> 0.6)
+      benchmark-trend (~> 0.4)
+      rspec (>= 3.0)
     rspec-core (3.10.1)
       rspec-support (~> 3.10.0)
     rspec-expectations (3.10.1)
@@ -1600,6 +1608,7 @@ DEPENDENCIES
   rexml (~> 3.2.5)
   rouge (~> 3.27.0)
   rqrcode-rails3 (~> 0.1.7)
+  rspec-benchmark (~> 0.6.0)
   rspec-parameterized
   rspec-rails (~> 5.0.1)
   rspec-retry (~> 0.6.1)

lib/gitlab/utils/strong_memoize.rb

@@ -22,10 +22,12 @@ module Gitlab
       #     end
       #
       def strong_memoize(name)
-        if strong_memoized?(name)
-          instance_variable_get(ivar(name))
+        key = ivar(name)
+
+        if instance_variable_defined?(key)
+          instance_variable_get(key)
         else
-          instance_variable_set(ivar(name), yield)
+          instance_variable_set(key, yield)
         end
       end
 
@@ -34,13 +36,23 @@ module Gitlab
       end
 
       def clear_memoization(name)
-        remove_instance_variable(ivar(name)) if instance_variable_defined?(ivar(name))
+        key = ivar(name)
+        remove_instance_variable(key) if instance_variable_defined?(key)
       end
 
       private
 
+      # Convert `"name"`/`:name` into `:@name`
+      #
+      # Depending on a type ensure that there's a single memory allocation
       def ivar(name)
-        "@#{name}"
+        if name.is_a?(Symbol)
+          name.to_s.prepend("@").to_sym
+        elsif name.is_a?(String)
+          :"@#{name}"
+        else
+          raise ArgumentError, "Invalid type of '#{name}'"
+        end
       end
     end
   end

spec/lib/gitlab/utils/strong_memoize_spec.rb

@@ -48,6 +48,36 @@ RSpec.describe Gitlab::Utils::StrongMemoize do
         let(:value) { value }
 
         it_behaves_like 'caching the value'
+
+        it 'raises exception for invalid key' do
+          expect { object.strong_memoize(10) { 20 } }.to raise_error /Invalid type of '10'/
+        end
+      end
+    end
+
+    context "memory allocation", type: :benchmark do
+      let(:value) { 'aaa' }
+
+      before do
+        object.method_name # warmup
+      end
+
+      [:method_name, "method_name"].each do |argument|
+        context "for #{argument.class}" do
+          it 'does allocate exactly one string when fetching value' do
+            expect do
+              object.strong_memoize(argument) { 10 }
+            end.to perform_allocation(1)
+          end
+
+          it 'does allocate exactly one string when storing value' do
+            object.clear_memoization(:method_name) # clear to force set
+
+            expect do
+              object.strong_memoize(argument) { 10 }
+            end.to perform_allocation(1)
+          end
+        end
       end
     end
   end

spec/spec_helper.rb

@@ -199,6 +199,7 @@ RSpec.configure do |config|
   config.include SidekiqMiddleware
   config.include StubActionCableConnection, type: :channel
   config.include StubSpamServices
+  config.include RSpec::Benchmark::Matchers, type: :benchmark
 
   include StubFeatureFlags