From 3fd54cbd82374c38c93c5d943883d3a5da83f892 Mon Sep 17 00:00:00 2001
From: Michael Go <michael.go@shopify.com>
Date: Thu, 27 Nov 2025 16:59:05 -0400
Subject: [PATCH 1/2] make sure stream exists before flushing

---
 lib/async/container/supervisor/connection.rb |  8 ++++----
 test/async/container/connection.rb           | 21 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/lib/async/container/supervisor/connection.rb b/lib/async/container/supervisor/connection.rb
index 49e9e28..2874307 100644
--- a/lib/async/container/supervisor/connection.rb
+++ b/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
diff --git a/test/async/container/connection.rb b/test/async/container/connection.rb
index 673b754..0fbd863 100644
--- a/test/async/container/connection.rb
+++ b/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

From a550afe2632a050476f3038b325b204d30a68e4a Mon Sep 17 00:00:00 2001
From: Michael Go <michael.go@shopify.com>
Date: Fri, 28 Nov 2025 18:36:57 -0400
Subject: [PATCH 2/2] raise IOError when connection is closed before a write

---
 lib/async/container/supervisor/connection.rb | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/lib/async/container/supervisor/connection.rb b/lib/async/container/supervisor/connection.rb
index 2874307..b6ca192 100644
--- a/lib/async/container/supervisor/connection.rb
+++ b/lib/async/container/supervisor/connection.rb
@@ -253,8 +253,9 @@ def next_id
 				#
 				# @parameter message [Hash] The message to write.
 				def write(**message)
+					raise IOError, "Connection is closed!" unless @stream
 					@stream&.write(JSON.dump(message) << "\n")
-					@stream&.flush
+					@stream&.flush # it is possible for @stream to become nil after the write call
 				end
 				
 				# Read a message from the connection stream.
@@ -327,9 +328,9 @@ def close
 						@reader = nil
 					end
 					
-					if @stream
-						@stream.close
+					if stream = @stream
 						@stream = nil
+						stream.close
 					end
 					
 					if @calls