Fix DB load balancing support for ActionCable

ActionCable workers do not run in the context of a request so we have to
add these callbacks that mimic our Rack / Sidekiq middleware.

We force the use of the primary because we currently have no way to
check the WAL location of the event that triggered the subscription and
these subscriptions are triggered in response to a write so up-to-date
information is important.

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