1 files changed, 21 insertions, 0 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