Better handling of restart loop.

Author: ioquatix

lib/slack/real_time/client.rb

@@ -106,13 +106,16 @@ def run_loop
 
       # Ensure the server is running, and ping the remote server if no other messages were sent.
       def keep_alive
+        # We can't ping the remote server if we aren't connected.
+        return false if @socket.nil? or !@socket.connected?
+
         time_since_last_message = @socket.time_since_last_message
 
         # If the server responded within the specified time, we are okay:
         return true if time_since_last_message < websocket_ping
 
-        # If the client is not connected or the server has not responded for a while:
-        if !@socket.connected? || time_since_last_message > (websocket_ping * 2)
+        # If the server has not responded for a while:
+        if time_since_last_message > (websocket_ping * 2)
           return false
         end
 

lib/slack/real_time/concurrency/async.rb

@@ -1,4 +1,5 @@
 require 'async/websocket'
+require 'async/notification'
 require 'async/clock'
 
 module Slack
@@ -25,33 +26,42 @@ def start_async(client)
 
           def start_reactor(client)
             Async do |task|
-              self.restart_async(client, @url)
+              @restart = ::Async::Notification.new
 
               if client.run_ping?
-                task.async do |subtask|
+                @ping_task = task.async do |subtask|
                   subtask.annotate 'client keep-alive'
 
-                  while true
+                  # The timer task will naturally exit after the driver is set to nil.
+                  while @restart
                     subtask.sleep client.websocket_ping
-                    run_ping!
+                    client.run_ping! if @restart
                   end
                 end
               end
+
+              while @restart
+                if @client_task
+                  @client_task.stop
+                end
+
+                @client_task = task.async do |subtask|
+                  subtask.annotate 'client run-loop'
+                  client.run_loop
+                end
+
+                @restart.wait
+              end
+
+              @ping_task.stop if @ping_task
             end
           end
 
           def restart_async(client, new_url)
             @url = new_url
             @last_message_at = current_time
 
-            if @client_task
-              @client_task.stop
-            end
-
-            @client_task = task.async do |subtask|
-              subtask.annotate 'client run-loop'
-              client.run_loop
-            end
+            @restart.signal if @restart
           end
 
           def current_time
@@ -63,6 +73,16 @@ def connect!
             run_loop
           end
 
+          # Kill the restart/ping loop.
+          def disconnect!
+            super
+          ensure
+            if restart = @restart
+              @restart = nil
+              restart.signal
+            end
+          end
+
           # Close the socket.
           def close
             super