From cefcc0f66e185b808fc93040b1d7805f95bf929f Mon Sep 17 00:00:00 2001
From: Jon Moss <me@jonathanmoss.me>
Date: Tue, 16 Feb 2016 12:20:03 -0500
Subject: [PATCH] Fix `unsubscribed` server side behavior

Before this commit, the `unsubscribed` callbacks in Action Cable server
side channels were never called. This is because when a WebSocket
"goodbye" message was sent from the client, the Action Cable server
didn't properly clean up after the now closed WebSocket. This means that
memory could possibly skyrocket with this behavior, since part of this
commit is to properly remove closed subscriptions from the global
subscriptions hash. Say you have 10,000 users currently connected, and
then all 10,000 disconnect -- before this patch, Action Cable would
still hold onto information (and Ruby objects!) for all of these now
dead connections.
---
 .../action_cable/connection/client_socket.rb  |  7 ++-----
 .../action_cable/connection/subscriptions.rb  |  2 +-
 actioncable/test/client/echo_channel.rb       |  4 ++++
 actioncable/test/client_test.rb               | 21 +++++++++++++++++++
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/actioncable/lib/action_cable/connection/client_socket.rb b/actioncable/lib/action_cable/connection/client_socket.rb
index ef937d7c16..95e1ac4c16 100644
--- a/actioncable/lib/action_cable/connection/client_socket.rb
+++ b/actioncable/lib/action_cable/connection/client_socket.rb
@@ -132,11 +132,8 @@ module ActionCable
           @ready_state = CLOSING
           @close_params = [reason, code]
 
-          if @stream
-            @stream.shutdown
-          else
-            finalize_close
-          end
+          @stream.shutdown if @stream
+          finalize_close
         end
 
         def finalize_close
diff --git a/actioncable/lib/action_cable/connection/subscriptions.rb b/actioncable/lib/action_cable/connection/subscriptions.rb
index d7f95e6a62..24934e12f2 100644
--- a/actioncable/lib/action_cable/connection/subscriptions.rb
+++ b/actioncable/lib/action_cable/connection/subscriptions.rb
@@ -54,7 +54,7 @@ module ActionCable
       end
 
       def unsubscribe_from_all
-        subscriptions.each { |id, channel| channel.unsubscribe_from_channel }
+        subscriptions.each { |id, channel| remove_subscription(channel) }
       end
 
       protected
diff --git a/actioncable/test/client/echo_channel.rb b/actioncable/test/client/echo_channel.rb
index 63e35f194a..5a7bac25c5 100644
--- a/actioncable/test/client/echo_channel.rb
+++ b/actioncable/test/client/echo_channel.rb
@@ -3,6 +3,10 @@ class EchoChannel < ActionCable::Channel::Base
     stream_from "global"
   end
 
+  def unsubscribed
+    'Goodbye from EchoChannel!'
+  end
+
   def ding(data)
     transmit(dong: data['message'])
   end
diff --git a/actioncable/test/client_test.rb b/actioncable/test/client_test.rb
index d30c381131..626b4e34a2 100644
--- a/actioncable/test/client_test.rb
+++ b/actioncable/test/client_test.rb
@@ -198,4 +198,25 @@ class ClientTest < ActionCable::TestCase
       c.close # disappear before read
     end
   end
+
+  def test_unsubscribe_client
+    with_puma_server do |port|
+      app = ActionCable.server
+      identifier = JSON.dump(channel: 'EchoChannel')
+
+      c = faye_client(port)
+      c.send_message command: 'subscribe', identifier: identifier
+      assert_equal({"identifier"=>"{\"channel\":\"EchoChannel\"}", "type"=>"confirm_subscription"}, c.read_message)
+      assert_equal(1, app.connections.count)
+      assert(app.remote_connections.where(identifier: identifier))
+
+      channel = app.connections.first.subscriptions.send(:subscriptions).first[1]
+      channel.expects(:unsubscribed)
+      c.close
+      sleep 0.1 # Data takes a moment to process
+
+      # All data is removed: No more connection or subscription information!
+      assert_equal(0, app.connections.count)
+    end
+  end
 end