Simplify code around read consistency and add specs

This commit avoid creating a new instance of a class and additional
memory allocations when ensuring read consistency in CE / EE.

ee/lib/ee/gitlab/database/consistency.rb

@@ -9,8 +9,8 @@ module EE
         # In EE we are disabling the database load balancing for calls that
         # require read consistency after recent writes.
         #
-        override :use_primary
-        def use_primary(&block)
+        override :with_read_consistency
+        def with_read_consistency(&block)
           ::Gitlab::Database::LoadBalancing::Session
             .current.use_primary(&block)
         end

ee/lib/gitlab/database/load_balancing/session.rb

@@ -28,6 +28,8 @@ module Gitlab
           @use_primary
         end
 
+        alias_method :using_primary?, :use_primary?
+
         def use_primary!
           @use_primary = true
         end

ee/spec/lib/gitlab/database/consistency_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Consistency do
+  let(:session) do
+    Gitlab::Database::LoadBalancing::Session.current
+  end
+
+  describe '.with_read_consistency' do
+    it 'sticks to primary database' do
+      expect(session).not_to be_using_primary
+
+      block = -> (&control) do
+        described_class.with_read_consistency do
+          expect(session).to be_using_primary
+
+          control.call
+        end
+      end
+
+      expect { |probe| block.call(&probe) }.to yield_control
+    end
+  end
+end

lib/gitlab/database/consistency.rb

@@ -5,7 +5,7 @@ module Gitlab
     ##
     # This class is used to make it possible to ensure read consistency in
     # GitLab EE without the need of overriding a lot of methods / classes /
-    # modules.
+    # classs.
     #
     # This is a CE class that does nothing in CE, because database load
     # balancing is EE-only feature, but you can still use it in CE. It will
@@ -15,21 +15,17 @@ module Gitlab
     # EE only to force usage of the primary database in EE.
     #
     class Consistency
-      def self.with_read_consistency(&block)
-        self.new.use_primary(&block)
-      end
-
       ##
       # In CE there is no database load balancing, so all reads are expected to
       # be consistent by the ACID guarantees of a single PostgreSQL instance.
       #
       # This method is overridden in EE.
       #
-      def use_primary(&block)
+      def self.with_read_consistency(&block)
         yield
       end
     end
   end
 end
 
-::Gitlab::Database::Consistency.prepend_if_ee('EE::Gitlab::Database::Consistency')
+::Gitlab::Database::Consistency.singleton_class.prepend_if_ee('EE::Gitlab::Database::Consistency')

spec/lib/gitlab/ci/trace/chunked_io_spec.rb

@@ -9,7 +9,7 @@ RSpec.describe Gitlab::Ci::Trace::ChunkedIO, :clean_gitlab_redis_cache do
   let(:chunked_io) { described_class.new(build) }
 
   before do
-    stub_feature_flags(ci_enable_live_trace: true)
+    stub_feature_flags(ci_enable_live_trace: true, gitlab_ci_trace_read_consistency: true)
   end
 
   describe "#initialize" do

spec/models/ci/build_trace_chunk_spec.rb

@@ -17,7 +17,7 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
   it_behaves_like 'having unique enum values'
 
   before do
-    stub_feature_flags(ci_enable_live_trace: true)
+    stub_feature_flags(ci_enable_live_trace: true, gitlab_ci_trace_read_consistency: true)
     stub_artifacts_object_storage
   end