Merge branch...

Merge branch '326338-add-a-shared-context-for-instrumentation-classes-for-queries-optimization' into 'master'

Add caching minimum/maximum ids in new instrumentation classes for batch queries

See merge request gitlab-org/gitlab!63560

lib/gitlab/usage/metrics/instrumentations/database_metric.rb

@@ -16,14 +16,20 @@ module Gitlab
           # end
           class << self
             def start(&block)
+              return @metric_start&.call unless block_given?
+
               @metric_start = block
             end
 
             def finish(&block)
+              return @metric_finish&.call unless block_given?
+
               @metric_finish = block
             end
 
             def relation(&block)
+              return @metric_relation&.call unless block_given?
+
               @metric_relation = block
             end
 
@@ -32,15 +38,21 @@ module Gitlab
               @column = column
             end
 
-            attr_reader :metric_operation, :metric_relation, :metric_start, :metric_finish, :column
+            def cache_start_and_finish_as(cache_key)
+              @cache_key = cache_key
+            end
+
+            attr_reader :metric_operation, :metric_relation, :metric_start, :metric_finish, :column, :cache_key
           end
 
           def value
+            start, finish = get_or_cache_batch_ids
+
             method(self.class.metric_operation)
               .call(relation,
                     self.class.column,
-                    start: self.class.metric_start&.call,
-                    finish: self.class.metric_finish&.call)
+                    start: start,
+                    finish: finish)
           end
 
           def to_sql
@@ -73,6 +85,22 @@ module Gitlab
               raise "Unknown time frame: #{time_frame} for DatabaseMetric"
             end
           end
+
+          def get_or_cache_batch_ids
+            return [self.class.start, self.class.finish] unless self.class.cache_key.present?
+
+            key_name = "metric_instrumentation/#{self.class.cache_key}"
+
+            start = Gitlab::Cache.fetch_once("#{key_name}_minimum_id", expires_in: 1.day) do
+              self.class.start
+            end
+
+            finish = Gitlab::Cache.fetch_once("#{key_name}_maximum_id", expires_in: 1.day) do
+              self.class.finish
+            end
+
+            [start, finish]
+          end
         end
       end
     end

spec/lib/gitlab/usage/metrics/instrumentations/database_metric_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
+  subject do
+    described_class.tap do |m|
+      m.relation { Issue }
+      m.operation :count
+      m.start { m.relation.minimum(:id) }
+      m.finish { m.relation.maximum(:id) }
+    end.new(time_frame: 'all')
+  end
+
+  describe '#value' do
+    let_it_be(:issue_1) { create(:issue) }
+    let_it_be(:issue_2) { create(:issue) }
+    let_it_be(:issue_3) { create(:issue) }
+    let_it_be(:issues) { Issue.all }
+
+    before do
+      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+    end
+
+    it 'calculates a correct result' do
+      expect(subject.value).to eq(3)
+    end
+
+    it 'does not cache the result of start and finish', :request_store, :use_clean_rails_redis_caching do
+      expect(Gitlab::Cache).not_to receive(:fetch_once)
+      expect(subject).to receive(:count).with(any_args, hash_including(start: issues.min_by(&:id).id, finish: issues.max_by(&:id).id)).and_call_original
+
+      subject.value
+
+      expect(Rails.cache.read('metric_instrumentation/special_issue_count_minimum_id')).to eq(nil)
+      expect(Rails.cache.read('metric_instrumentation/special_issue_count_maximum_id')).to eq(nil)
+    end
+
+    context 'with start and finish not called' do
+      subject do
+        described_class.tap do |m|
+          m.relation { Issue }
+          m.operation :count
+        end.new(time_frame: 'all')
+      end
+
+      it 'calculates a correct result' do
+        expect(subject.value).to eq(3)
+      end
+    end
+
+    context 'with cache_start_and_finish_as called' do
+      subject do
+        described_class.tap do |m|
+          m.relation { Issue }
+          m.operation :count
+          m.start { m.relation.minimum(:id) }
+          m.finish { m.relation.maximum(:id) }
+          m.cache_start_and_finish_as :special_issue_count
+        end.new(time_frame: 'all')
+      end
+
+      it 'caches using the key name passed', :request_store, :use_clean_rails_redis_caching do
+        expect(Gitlab::Cache).to receive(:fetch_once).with('metric_instrumentation/special_issue_count_minimum_id', any_args).and_call_original
+        expect(Gitlab::Cache).to receive(:fetch_once).with('metric_instrumentation/special_issue_count_maximum_id', any_args).and_call_original
+        expect(subject).to receive(:count).with(any_args, hash_including(start: issues.min_by(&:id).id, finish: issues.max_by(&:id).id)).and_call_original
+
+        subject.value
+
+        expect(Rails.cache.read('metric_instrumentation/special_issue_count_minimum_id')).to eq(issues.min_by(&:id).id)
+        expect(Rails.cache.read('metric_instrumentation/special_issue_count_maximum_id')).to eq(issues.max_by(&:id).id)
+      end
+    end
+  end
+end