Handle closed streams in #write. (#13)

Author: ggmichaelgo

lib/async/container/supervisor/connection.rb

@@ -253,8 +253,9 @@ def next_id
 				#
 				# @parameter message [Hash] The message to write.
 				def write(**message)
-					@stream.write(JSON.dump(message) << "\n")
-					@stream.flush
+					raise IOError, "Connection is closed!" unless @stream
+					@stream&.write(JSON.dump(message) << "\n")
+					@stream&.flush # it is possible for @stream to become nil after the write call
 				end
 
 				# Read a message from the connection stream.

test/async/container/connection.rb

@@ -17,6 +17,13 @@ def dispatch(call)
 	end
 end
 
+class SlowWriteStringIO < StringIO
+	def write(data)
+		super(data)
+		sleep(1) # Simulate a slow write
+	end
+end
+
 describe Async::Container::Supervisor::Connection do
 	let(:stream) {StringIO.new}
 	let(:connection) {Async::Container::Supervisor::Connection.new(stream)}
@@ -208,5 +215,19 @@ def dispatch(call)
 			connection.close
 			expect(stream).to be(:closed?)
 		end
+		
+		it "closes while writing" do
+			slow_connection = Async::Container::Supervisor::Connection.new(SlowWriteStringIO.new)
+			
+			Async do
+				Async do |task|
+					slow_connection.write(id: 1, do: :test)
+				end
+				
+				Async do |task|
+					slow_connection.close
+				end
+			end
+		end
 	end
 end