make sure stream exists before flushing

Author: ggmichaelgo

lib/async/container/supervisor/connection.rb

@@ -253,8 +253,8 @@ def next_id
 				#
 				# @parameter message [Hash] The message to write.
 				def write(**message)
-					@stream.write(JSON.dump(message) << "\n")
-					@stream.flush
+					@stream&.write(JSON.dump(message) << "\n")
+					@stream&.flush
 				end
 
 				# Read a message from the connection stream.
@@ -327,9 +327,9 @@ def close
 						@reader = nil
 					end
 
-					if stream = @stream
+					if @stream
+						@stream.close
 						@stream = nil
-						stream.close
 					end
 
 					if @calls

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