Prevent the stream_from method throwing RuntimeError(can't add a new key into hash during iteration)

Author: OuYangJinTing

actioncable/lib/action_cable/channel/base.rb

@@ -165,6 +165,7 @@ def initialize(connection, identifier, params = {})
 
         @reject_subscription = nil
         @subscription_confirmation_sent = nil
+        @unsubscribed = false
 
         delegate_connection_identifiers
       end
@@ -200,11 +201,16 @@ def subscribe_to_channel
       # cleanup with callbacks. This method is not intended to be called directly by
       # the user. Instead, override the #unsubscribed callback.
       def unsubscribe_from_channel # :nodoc:
+        @unsubscribed = true
         run_callbacks :unsubscribe do
           unsubscribed
         end
       end
 
+      def unsubscribed? # :nodoc:
+        @unsubscribed
+      end
+
       private
         # Called once a consumer has become a subscriber of the channel. Usually the
         # place to set up any streams you want this channel to be sending to the

actioncable/lib/action_cable/channel/streams.rb

@@ -88,6 +88,8 @@ module Streams
       # callback. Defaults to `coder: nil` which does no decoding, passes raw
       # messages.
       def stream_from(broadcasting, callback = nil, coder: nil, &block)
+        return if unsubscribed?
+
         broadcasting = String(broadcasting)
 
         # Don't send the confirmation until pubsub#subscribe is successful

actioncable/test/channel/base_test.rb

@@ -126,6 +126,16 @@ def error_handler
     assert_not_predicate @channel, :subscribed?
   end
 
+  test "unsubscribed? method returns correct status" do
+    assert_not @channel.unsubscribed?
+
+    @channel.subscribe_to_channel
+    assert_not @channel.unsubscribed?
+
+    @channel.unsubscribe_from_channel
+    assert @channel.unsubscribed?
+  end
+
   test "connection identifiers" do
     assert_equal @user.name, @channel.current_user.name
   end

actioncable/test/channel/stream_test.rb

@@ -4,6 +4,7 @@
 require "minitest/mock"
 require "stubs/test_connection"
 require "stubs/room"
+require "concurrent/atomic/cyclic_barrier"
 
 module ActionCable::StreamTests
   class Connection < ActionCable::Connection::Base
@@ -280,6 +281,59 @@ class StreamTest < ActionCable::TestCase
       end
     end
 
+    test "concurrent unsubscribe_from_channel and stream_from do not raise RuntimeError" do
+      ENV["UNSUBSCRIBE_SLEEP_TIME"] = "0.0001" # Set a delay to increase the chance of concurrent execution
+      run_in_eventmachine do
+        connection = TestConnection.new
+        channel = ChatChannel.new connection, "{id: 1}", id: 1
+        channel.subscribe_to_channel
+
+        # Set up initial streams
+        channel.stream_from "room_one"
+        channel.stream_from "room_two"
+        wait_for_async
+
+        # Create barriers to synchronize thread execution
+        barrier = Concurrent::CyclicBarrier.new(2)
+
+        exception_caught = nil
+
+        # Thread 1: calls unsubscribe_from_channel
+        thread1 = Thread.new do
+          barrier.wait
+          # Add a small delay to increase the chance of concurrent execution
+          sleep 0.001
+          channel.unsubscribe_from_channel
+        rescue => e
+          exception_caught = e
+        ensure
+          barrier.wait
+        end
+
+        # Thread 2: calls stream_from during unsubscribe_from_channel iteration
+        thread2 = Thread.new do
+          barrier.wait
+          # Try to add streams while unsubscribe_from_channel is potentially iterating
+          10.times do |i|
+            channel.stream_from "concurrent_room_#{i}"
+            sleep 0.0001 # Small delay to interleave with unsubscribe_from_channel
+          end
+        rescue => e
+          exception_caught = e
+        ensure
+          barrier.wait
+        end
+
+        thread1.join
+        thread2.join
+
+        # Ensure no RuntimeError was raised during concurrent access
+        assert_nil exception_caught, "Concurrent unsubscribe_from_channel and stream_from should not raise RuntimeError: #{exception_caught}"
+      end
+    ensure
+      ENV.delete("UNSUBSCRIBE_SLEEP_TIME")
+    end
+
     private
       def subscribers_of(connection)
         connection

actioncable/test/stubs/test_adapter.rb

@@ -10,6 +10,7 @@ def subscribe(channel, callback, success_callback = nil)
   end
 
   def unsubscribe(channel, callback)
+    sleep ENV["UNSUBSCRIBE_SLEEP_TIME"].to_f if ENV["UNSUBSCRIBE_SLEEP_TIME"]
     subscriber_map[channel].delete(callback)
     subscriber_map.delete(channel) if subscriber_map[channel].empty?
     @@unsubscribe_called = { channel: channel, callback: callback }