4 files changed, 22 insertions, 3 deletions
diff --git a/lib/action_cable/server/base.rb b/lib/action_cable/server/base.rb
index 43849928b9..9315a48f20 100644
--- a/lib/action_cable/server/base.rb
+++ b/lib/action_cable/server/base.rb
@@ -18,6 +18,7 @@ module ActionCable
 
       # Called by rack to setup the server.
       def call(env)
+        setup_heartbeat_timer
         config.connection_class.new(self, env).process
       end
 
diff --git a/lib/action_cable/server/connections.rb b/lib/action_cable/server/connections.rb
index 15d7c3c8c7..b3d1632cf7 100644
--- a/lib/action_cable/server/connections.rb
+++ b/lib/action_cable/server/connections.rb
@@ -4,6 +4,8 @@ module ActionCable
     # you can't use this collection as an full list of all the connections established against your application. Use RemoteConnections for that.
     # As such, this is primarily for internal use.
     module Connections
+      BEAT_INTERVAL = 3
+
       def connections
         @connections ||= []
       end
@@ -16,6 +18,17 @@ module ActionCable
         connections.delete connection
       end
 
+      # Websocket connection implementations differ on when they'll mark a connection as stale. We basically never want a connection to go stale, as you
+      # then can't rely on being able to receive and send to it. So there's a 3 second heartbeat running on all connections. If the beat fails, we automatically
+      # disconnect.
+      def setup_heartbeat_timer
+        EM.next_tick do
+          @heartbeat_timer ||= EventMachine.add_periodic_timer(BEAT_INTERVAL) do
+            EM.next_tick { connections.map &:beat }
+          end
+        end
+      end
+
       def open_connections_statistics
         connections.map(&:statistics)
       end
diff --git a/lib/action_cable/server/worker.rb b/lib/action_cable/server/worker.rb
index d7823ecf93..91496775b8 100644
--- a/lib/action_cable/server/worker.rb
+++ b/lib/action_cable/server/worker.rb
@@ -5,10 +5,13 @@ module ActionCable
       include ActiveSupport::Callbacks
       include Celluloid
 
+      attr_reader :connection
       define_callbacks :work
-      include ClearDatabaseConnections
+      include ActiveRecordConnectionManagement
 
       def invoke(receiver, method, *args)
+        @connection = receiver
+
         run_callbacks :work do
           receiver.send method, *args
         end
@@ -20,6 +23,8 @@ module ActionCable
       end
 
       def run_periodic_timer(channel, callback)
+        @connection = channel.connection
+
         run_callbacks :work do
           callback.respond_to?(:call) ? channel.instance_exec(&callback) : channel.send(callback)
         end
diff --git a/lib/action_cable/server/worker/clear_database_connections.rb b/lib/action_cable/server/worker/active_record_connection_management.rb
index 722d363a41..1ede0095f8 100644
--- a/lib/action_cable/server/worker/clear_database_connections.rb
+++ b/lib/action_cable/server/worker/active_record_connection_management.rb
@@ -2,7 +2,7 @@ module ActionCable
   module Server
     class Worker
       # Clear active connections between units of work so the long-running channel or connection processes do not hoard connections.
-      module ClearDatabaseConnections
+      module ActiveRecordConnectionManagement
         extend ActiveSupport::Concern
 
         included do
@@ -12,7 +12,7 @@ module ActionCable
         end
 
         def with_database_connections
-          yield
+          ActiveRecord::Base.logger.tagged(*connection.logger.tags) { yield }
         ensure
           ActiveRecord::Base.clear_active_connections!
         end