close streams when a SocketError occurs

Author: eknapp-inst

lib/net-http2/client.rb

@@ -9,6 +9,7 @@ module NetHttp2
   PROXY_SETTINGS_KEYS = [:proxy_addr, :proxy_port, :proxy_user, :proxy_pass]
 
   AsyncRequestTimeout = Class.new(StandardError)
+  SocketClosedError = Class.new(StandardError)
 
   class Client
 
@@ -39,8 +40,7 @@ def call(method, path, options={})
 
     def call_async(request)
       ensure_open
-      stream = new_monitored_stream_for request
-      stream.async_call_with request
+      new_stream.async_call_with request
     end
 
     def prepare_request(method, path, options={})
@@ -53,6 +53,7 @@ def ssl?
 
     def close
       exit_thread(@socket_thread)
+      @socket_error = SocketClosedError.new
       init_vars
     end
 
@@ -74,10 +75,14 @@ def stream_count
 
     private
 
-    def init_vars
+    def init_vars(error: nil)
       @mutex.synchronize do
         @socket.close if @socket && !@socket.closed?
 
+        (@streams || {}).each do |k, v|
+          v.force_close(@socket_error)
+        end
+
         @h2              = nil
         @socket          = nil
         @socket_thread   = nil
@@ -86,20 +91,16 @@ def init_vars
       end
     end
 
-    def new_stream
-      @mutex.synchronize { NetHttp2::Stream.new(h2_stream: h2.new_stream) }
-    rescue StandardError => e
-      close
-      raise e
-    end
-
-    def new_monitored_stream_for(request)
-      stream = new_stream
+    def new_stream()
+      stream = @mutex.synchronize { NetHttp2::Stream.new(h2_stream: h2.new_stream) }
 
-      @streams[stream.id] = true
-      request.on(:close) { @streams.delete(stream.id) }
+      @streams[stream.id] = stream
+      stream.on(:close) { @streams.delete(stream.id) }
 
       stream
+    rescue StandardError => e
+      close
+      raise e
     end
 
     def ensure_open
@@ -115,13 +116,15 @@ def ensure_open
 
           rescue EOFError
             # socket closed
+            @socket_error = SocketError.new('Socket was remotely closed')
             init_vars
-            callback_or_raise SocketError.new('Socket was remotely closed')
+            callback_or_raise @socket_error
 
           rescue Exception => e
             # error on socket
+            @socket_error = e
             init_vars
-            callback_or_raise e
+            callback_or_raise @socket_error
           end
         end.tap { |t| t.abort_on_exception = true }
       end

lib/net-http2/stream.rb

@@ -10,6 +10,7 @@ def initialize(options={})
       @async     = false
       @completed = false
       @mutex     = Mutex.new
+      @error     = nil
       @cv        = ConditionVariable.new
 
       listen_for_headers
@@ -42,6 +43,15 @@ def async?
       @async
     end
 
+    def wait
+      wait_for_completed
+    end
+
+    def force_close(error = nil)
+      @error = error
+      @mutex.synchronize { @cv.signal }
+    end
+
     private
 
     def listen_for_headers
@@ -98,6 +108,9 @@ def sync_respond
 
     def wait_for_completed
       @mutex.synchronize { @cv.wait(@mutex, @request.timeout) }
+      if @error || !@completed
+        raise @error
+      end
     end
   end
 end