Use AS::Executor / AS::Reloader to support reloading in ActionCable

1 files changed, 25 insertions, 0 deletions

diff --git a/actioncable/lib/action_cable/engine.rb b/actioncable/lib/action_cable/engine.rb
index ae0c59dccd..c90aadaf2c 100644
--- a/actioncable/lib/action_cable/engine.rb
+++ b/actioncable/lib/action_cable/engine.rb
@@ -51,5 +51,30 @@ module ActionCable
         end
       end
     end
+
+    initializer "action_cable.set_work_hooks" do |app|
+      ActiveSupport.on_load(:action_cable) do
+        ActionCable::Server::Worker.set_callback :work, :around, prepend: true do |_, inner|
+          app.executor.wrap do
+            # If we took a while to get the lock, we may have been halted
+            # in the meantime. As we haven't started doing any real work
+            # yet, we should pretend that we never made it off the queue.
+            unless stopping?
+              inner.call
+            end
+          end
+        end
+
+        wrap = lambda do |_, inner|
+          app.executor.wrap(&inner)
+        end
+        ActionCable::Channel::Base.set_callback :subscribe, :around, prepend: true, &wrap
+        ActionCable::Channel::Base.set_callback :unsubscribe, :around, prepend: true, &wrap
+
+        app.reloader.before_class_unload do
+          ActionCable.server.restart
+        end
+      end
+    end
   end
 end