Add caching for ES migration_has_finished?

This MR adds a caching layer to the
`Elastic::DataMigrationService.migration_has_finished?` method. We plan
to call it each time we build an Elasticsearch query when we have before
and after migration states, so it needs to be performant.

ee/app/services/elastic/data_migration_service.rb

@@ -20,18 +20,36 @@ module Elastic
         migrations.sort_by(&:version)
       end
 
+      def drop_migration_has_finished_cache!(migration)
+        name = migration.name.underscore
+        Rails.cache.delete cache_key(:migration_has_finished, name)
+      end
+
       def migration_has_finished?(name)
-        migration = migrations.find { |migration| migration.name == name.camelize }
+        Rails.cache.fetch cache_key(:migration_has_finished, name), expires_in: 30.minutes do
+          migration_has_finished_uncached?(name)
+        end
+      end
+
+      def migration_has_finished_uncached?(name)
+        migration = migrations.find { |migration| migration.name == name.to_s.camelize }
 
         !!migration&.load_from_index&.dig('_source', 'completed')
       end
 
       def mark_all_as_completed!
-        migrations.each { |migration| migration.save!(completed: true) }
+        migrations.each do |migration|
+          migration.save!(completed: true)
+          drop_migration_has_finished_cache!(migration)
+        end
       end
 
       private
 
+      def cache_key(method_name, *additional_key)
+        [name, method_name, *additional_key]
+      end
+
       def parse_migration_filename(filename)
         File.basename(filename).scan(MIGRATION_REGEXP).first
       end

ee/app/workers/elastic/migration_worker.rb

@@ -34,6 +34,8 @@ module Elastic
         completed = migration.completed?
         logger.info "MigrationWorker: migration[#{migration.name}] updating with completed: #{completed}"
         migration.save!(completed: completed)
+
+        Elastic::DataMigrationService.drop_migration_has_finished_cache!(migration)
       end
     end
 

ee/spec/services/elastic/data_migration_service_spec.rb

@@ -29,26 +29,44 @@ RSpec.describe Elastic::DataMigrationService, :elastic do
     end
   end
 
-  describe '.migration_has_finished?' do
+  describe '.migration_has_finished_uncached?' do
     let(:migration) { subject.migrations.first }
     let(:migration_name) { migration.name.underscore }
 
     it 'returns true if migration has finished' do
-      expect(subject.migration_has_finished?(migration_name)).to eq(false)
+      expect(subject.migration_has_finished_uncached?(migration_name)).to eq(false)
 
       migration.save!(completed: false)
       refresh_index!
 
-      expect(subject.migration_has_finished?(migration_name)).to eq(false)
+      expect(subject.migration_has_finished_uncached?(migration_name)).to eq(false)
 
       migration.save!(completed: true)
       refresh_index!
 
-      expect(subject.migration_has_finished?(migration_name)).to eq(true)
+      expect(subject.migration_has_finished_uncached?(migration_name)).to eq(true)
+    end
+  end
+
+  describe '.migration_has_finished?' do
+    let(:migration) { subject.migrations.first }
+    let(:migration_name) { migration.name.underscore }
+    let(:finished) { false }
+
+    before do
+      allow(Rails).to receive(:cache).and_return(ActiveSupport::Cache::MemoryStore.new)
+      allow(subject).to receive(:migration_has_finished_uncached?).with(migration_name).and_return(finished)
+    end
+
+    it 'calls the uncached method only once' do
+      expect(subject).to receive(:migration_has_finished_uncached?).once
+
+      expect(subject.migration_has_finished?(migration_name)).to eq(finished)
+      expect(subject.migration_has_finished?(migration_name)).to eq(finished)
     end
   end
 
-  describe 'mark_all_as_completed!' do
+  describe '.mark_all_as_completed!' do
     it 'creates all migration versions' do
       expect(Elastic::MigrationRecord.persisted_versions(completed: true).count).to eq(0)
 
@@ -57,5 +75,41 @@ RSpec.describe Elastic::DataMigrationService, :elastic do
 
       expect(Elastic::MigrationRecord.persisted_versions(completed: true).count).to eq(subject.migrations.count)
     end
+
+    it 'drops all cache keys' do
+      allow(subject).to receive(:migrations).and_return(
+        [
+          Elastic::MigrationRecord.new(version: 100, name: 'SomeMigration', filename: nil),
+          Elastic::MigrationRecord.new(version: 200, name: 'SomeOtherMigration', filename: nil)
+        ]
+      )
+
+      subject.migrations.each do |migration|
+        expect(subject).to receive(:drop_migration_has_finished_cache!).with(migration)
+      end
+
+      subject.mark_all_as_completed!
+    end
+  end
+
+  describe '.drop_migration_has_finished_cache!' do
+    let(:migration) { subject.migrations.first }
+    let(:migration_name) { migration.name.underscore }
+    let(:finished) { true }
+
+    before do
+      allow(Rails).to receive(:cache).and_return(ActiveSupport::Cache::MemoryStore.new)
+      allow(subject).to receive(:migration_has_finished_uncached?).with(migration_name).and_return(finished)
+    end
+
+    it 'drops cache' do
+      expect(subject).to receive(:migration_has_finished_uncached?).twice
+
+      expect(subject.migration_has_finished?(migration_name)).to eq(finished)
+
+      subject.drop_migration_has_finished_cache!(migration)
+
+      expect(subject.migration_has_finished?(migration_name)).to eq(finished)
+    end
   end
 end

ee/spec/workers/elastic/migration_worker_spec.rb

@@ -67,6 +67,7 @@ RSpec.describe Elastic::MigrationWorker, :elastic do
             end
 
             expect(migration).to receive(:save!).with(completed: completed)
+            expect(Elastic::DataMigrationService).to receive(:drop_migration_has_finished_cache!).with(migration)
 
             subject.perform
           end