Refactor handling of verification state

Shift more responsibility to Gitlab::Geo::VerificationState for managing
its data.

ee/app/models/concerns/geo/verifiable_replicator.rb

@@ -125,17 +125,12 @@ module Geo
       Geo::VerificationWorker.perform_async(replicable_name, model_record.id)
     end
 
-    # Calculates checksum and asks the model/registry to update verification
+    # Calculates checksum and asks the model/registry to manage verification
     # state.
     def verify
-      model_record.verification_started! unless model_record.verification_started?
-
-      calculation_started_at = Time.current
-      checksum = model_record.calculate_checksum
-
-      model_record.verification_succeeded_with_checksum!(checksum, calculation_started_at)
-    rescue => e
-      model_record.verification_failed_with_message!('Error calculating the checksum', e)
+      model_record.track_checksum_attempt! do
+        model_record.calculate_checksum
+      end
     end
 
     # Check if given checksum matches known one

ee/lib/gitlab/geo/verification_state.rb

@@ -177,6 +177,25 @@ module Gitlab
         end
       end
 
+      # Provides a safe and easy way to manage the verification state for a
+      # synchronous checksum calculation.
+      #
+      # @yieldreturn [String] calculated checksum value
+      def track_checksum_attempt!(&block)
+        # This line only applies to Geo::VerificationWorker, not
+        # Geo::VerificationBatchWorker, since the latter sets the whole batch to
+        # "verification_started" in the same DB query that fetches the batch.
+        verification_started! unless verification_started?
+
+        calculation_started_at = Time.current
+
+        checksum = yield
+
+        track_checksum_result!(checksum, calculation_started_at)
+      rescue => e
+        verification_failed_with_message!('Error during verification', e)
+      end
+
       # Convenience method to update checksum and transition to success state.
       #
       # @param [String] checksum value generated by the checksum routine
@@ -193,23 +212,36 @@ module Gitlab
         end
       end
 
-      def resource_updated_during_checksum?(calculation_started_at)
-        self.reset.verification_started_at > calculation_started_at
-      end
-
       # Convenience method to update failure message and transition to failed
       # state.
       #
       # @param [String] message error information
       # @param [StandardError] error exception
       def verification_failed_with_message!(message, error = nil)
-        log_error('Error calculating the checksum', error)
+        log_error(message, error)
 
         self.verification_failure = message
         self.verification_failure += ": #{error.message}" if error.respond_to?(:message)
 
         self.verification_failed!
       end
+
+      private
+
+      # Records the calculated checksum result
+      #
+      # Overridden by ReplicableRegistry so it can also compare with primary
+      # checksum.
+      #
+      # @param [String] calculated checksum value
+      # @param [Time] when checksum calculation was started
+      def track_checksum_result!(checksum, calculation_started_at)
+        verification_succeeded_with_checksum!(checksum, calculation_started_at)
+      end
+
+      def resource_updated_during_checksum?(calculation_started_at)
+        self.reset.verification_started_at > calculation_started_at
+      end
     end
   end
 end

ee/spec/lib/gitlab/geo/verification_state_spec.rb

@@ -157,6 +157,55 @@ RSpec.describe Gitlab::Geo::VerificationState do
     end
   end
 
+  describe '#track_checksum_attempt!' do
+    context 'when verification was not yet started' do
+      it 'starts verification' do
+        expect do
+          subject.track_checksum_attempt! do
+            'a_checksum_value'
+          end
+        end.to change { subject.verification_started_at }.from(nil)
+      end
+
+      it 'sets verification_succeeded' do
+        expect do
+          subject.track_checksum_attempt! do
+            'a_checksum_value'
+          end
+        end.to change { subject.verification_succeeded? }.from(false).to(true)
+      end
+    end
+
+    context 'when verification was started' do
+      it 'does not update verification_started_at' do
+        subject.verification_started!
+        expected = subject.verification_started_at
+
+        subject.track_checksum_attempt! do
+          'a_checksum_value'
+        end
+
+        expect(subject.verification_started_at).to be_within(1.second).of(expected)
+      end
+    end
+
+    it 'yields to the checksum calculation' do
+      expect do |probe|
+        subject.track_checksum_attempt!(&probe)
+      end.to yield_with_no_args
+    end
+
+    context 'when an error occurs while yielding' do
+      it 'sets verification_failed' do
+        subject.track_checksum_attempt! do
+          raise 'an error'
+        end
+
+        expect(subject.reload.verification_failed?).to be_truthy
+      end
+    end
+  end
+
   describe '#verification_succeeded_with_checksum!' do
     before do
       subject.verification_started!

ee/spec/support/shared_examples/models/concerns/verifiable_replicator_shared_examples.rb

@@ -326,32 +326,13 @@ RSpec.shared_examples 'a verifiable replicator' do
   end
 
   describe '#verify' do
-    context 'on a Geo primary' do
-      before do
-        stub_primary_node
-      end
+    it 'wraps the checksum calculation in track_checksum_attempt!' do
+      tracker = double('tracker', calculate_checksum: 'abc123')
+      allow(replicator).to receive(:model_record).and_return(tracker)
 
-      context 'when the checksum succeeds' do
-        it 'delegates checksum calculation and the state change to model_record' do
-          expect(model_record).to receive(:calculate_checksum).and_return('abc123')
-          expect(model_record).to receive(:verification_succeeded_with_checksum!).with('abc123', kind_of(Time))
+      expect(tracker).to receive(:track_checksum_attempt!).and_yield
 
-          replicator.verify
-        end
-      end
-
-      context 'when an error is raised during calculate_checksum' do
-        it 'passes the error message' do
-          error = StandardError.new('Some exception')
-          allow(model_record).to receive(:calculate_checksum) do
-            raise error
-          end
-
-          expect(model_record).to receive(:verification_failed_with_message!).with('Error calculating the checksum', error)
-
-          replicator.verify
-        end
-      end
+      replicator.verify
     end
   end
 end