1 files changed, 26 insertions, 0 deletions
diff --git a/lib/gitlab/database/load_balancing/action_cable_callbacks.rb b/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
new file mode 100644
index 00000000000..4feba989a0a
--- /dev/null
+++ b/lib/gitlab/database/load_balancing/action_cable_callbacks.rb
@@ -0,0 +1,26 @@
+# 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