Add latest changes from gitlab-org/gitlab@master

3 files changed, 49 insertions, 2 deletions

diff --git a/lib/gitlab/action_cable/request_store_callbacks.rb b/lib/gitlab/action_cable/request_store_callbacks.rb
new file mode 100644
index 00000000000..a9f30b0fc10
--- /dev/null
+++ b/lib/gitlab/action_cable/request_store_callbacks.rb
@@ -0,0 +1,21 @@
+# 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
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
diff --git a/lib/gitlab/database/load_balancing/configuration.rb b/lib/gitlab/database/load_balancing/configuration.rb
index 6bd94bad5cf..c874825c98a 100644
--- a/lib/gitlab/database/load_balancing/configuration.rb
+++ b/lib/gitlab/database/load_balancing/configuration.rb
@@ -11,7 +11,7 @@ module Gitlab
 
         # Creates a configuration object for the given ActiveRecord model.
         def self.for_model(model)
-          cfg = model.connection_db_config.configuration_hash
+          cfg = model.connection_db_config.configuration_hash.deep_symbolize_keys
           lb_cfg = cfg[:load_balancing] || {}
           config = new(model)
 
@@ -35,7 +35,7 @@ module Gitlab
             config.hosts = hosts
           end
 
-          discover = (lb_cfg[:discover] || {}).symbolize_keys
+          discover = lb_cfg[:discover] || {}
 
           # We iterate over the known/default keys so we don't end up with
           # random keys in our configuration hash.