Merge branch '340340-fix-actioncable-not-releasing-replica-connections' into 'master'

Fix DB load balancing support for ActionCable

See merge request gitlab-org/gitlab!69889

config/initializers/action_cable.rb

@@ -17,3 +17,6 @@ ActionCable::SubscriptionAdapter::Redis.redis_connector = lambda do |config|
 
   ::Redis.new(args)
 end
+
+Gitlab::ActionCable::RequestStoreCallbacks.install
+Gitlab::Database::LoadBalancing::ActionCableCallbacks.install if Gitlab::Database::LoadBalancing.enable?

lib/gitlab/action_cable/request_store_callbacks.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module ActionCable
+    module RequestStoreCallbacks
+      def self.install
+        ::ActionCable::Server::Worker.set_callback :work, :around, &wrapper
+        ::ActionCable::Channel::Base.set_callback :subscribe, :around, &wrapper
+        ::ActionCable::Channel::Base.set_callback :unsubscribe, :around, &wrapper
+      end
+
+      def self.wrapper
+        lambda do |_, inner|
+          ::Gitlab::WithRequestStore.with_request_store do
+            inner.call
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/load_balancing/action_cable_callbacks.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      module ActionCableCallbacks
+        def self.install
+          ::ActionCable::Server::Worker.set_callback :work, :around, &wrapper
+          ::ActionCable::Channel::Base.set_callback :subscribe, :around, &wrapper
+          ::ActionCable::Channel::Base.set_callback :unsubscribe, :around, &wrapper
+        end
+
+        def self.wrapper
+          lambda do |_, inner|
+            ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
+
+            inner.call
+          ensure
+            ::Gitlab::Database::LoadBalancing.proxy.load_balancer.release_host
+            ::Gitlab::Database::LoadBalancing::Session.clear_session
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/action_cable/request_store_callbacks_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe Gitlab::ActionCable::RequestStoreCallbacks do
+  describe '.wrapper' do
+    it 'enables RequestStore in the inner block' do
+      expect(RequestStore.active?).to eq(false)
+
+      described_class.wrapper.call(
+        nil,
+        lambda do
+          expect(RequestStore.active?).to eq(true)
+        end
+      )
+
+      expect(RequestStore.active?).to eq(false)
+    end
+  end
+end

spec/lib/gitlab/database/load_balancing/action_cable_callbacks_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe Gitlab::Database::LoadBalancing::ActionCableCallbacks, :request_store do
+  describe '.wrapper' do
+    it 'uses primary and then releases the connection and clears the session' do
+      expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+      expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
+
+      described_class.wrapper.call(
+        nil,
+        lambda do
+          expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).to eq(true)
+        end
+      )
+    end
+
+    context 'with an exception' do
+      it 'releases the connection and clears the session' do
+        expect(Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer, :release_host)
+        expect(Gitlab::Database::LoadBalancing::Session).to receive(:clear_session)
+
+        expect do
+          described_class.wrapper.call(nil, lambda { raise 'test_exception' })
+        end.to raise_error('test_exception')
+      end
+    end
+  end
+end