1 files changed, 17 insertions, 6 deletions
diff --git a/src/ruby/end2end/multiple_killed_watching_threads_driver.rb b/src/ruby/end2end/multiple_killed_watching_threads_driver.rb
index 59f6f275e4..8f078cfbed 100755
--- a/src/ruby/end2end/multiple_killed_watching_threads_driver.rb
+++ b/src/ruby/end2end/multiple_killed_watching_threads_driver.rb
@@ -20,7 +20,7 @@ Thread.abort_on_exception = true
 
 include GRPC::Core::ConnectivityStates
 
-def watch_state(ch)
+def watch_state(ch, sleep_time)
   thd = Thread.new do
     state = ch.connectivity_state(false)
     fail "non-idle state: #{state}" unless state == IDLE
@@ -28,23 +28,34 @@ def watch_state(ch)
   end
   # sleep to get the thread into the middle of a
   # "watch connectivity state" call
-  sleep 0.1
+  sleep sleep_time
   thd.kill
 end
 
-def main
+def run_multiple_killed_watches(num_threads, sleep_time)
   channels = []
-  10.times do
+  num_threads.times do
     ch = GRPC::Core::Channel.new('dummy_host',
                                  nil, :this_channel_is_insecure)
-    watch_state(ch)
+    watch_state(ch, sleep_time)
     channels << ch
   end
 
   # checking state should still be safe to call
   channels.each do |c|
-    fail unless c.connectivity_state(false) == FATAL_FAILURE
+    connectivity_state = c.connectivity_state(false)
+    # The state should be FATAL_FAILURE in the case that it was interrupted
+    # while watching connectivity state, and IDLE if it we never started
+    # watching the channel's connectivity state
+    unless [FATAL_FAILURE, IDLE].include?(connectivity_state)
+      fail "unexpected connectivity state: #{connectivity_state}"
+    end
   end
 end
 
+def main
+  run_multiple_killed_watches(10, 0.1)
+  run_multiple_killed_watches(1000, 0.001)
+end
+
 main