refactor connection to use MessagePacker's internal stream

Author: ggmichaelgo

lib/async/container/supervisor/connection.rb

@@ -234,7 +234,7 @@ def initialize(stream, id = 0, **state)
 					@stream = stream
 					@id = id
 					@state = state
-					@message_wrapper = MessageWrapper.new
+					@message_wrapper = MessageWrapper.new(@stream)
 					@reader = nil
 					@calls = {}
 				end
@@ -259,37 +259,16 @@ def next_id
 				#
 				# @parameter message [Hash] The message to write.
 				def write(**message)
-					data = @message_wrapper.pack(message)
-					# Write 4-byte length prefix
-					data_length = [data.bytesize].pack("N")
-					@stream.write(data_length + data)
-					@stream.flush
+					@message_wrapper.write(message)
 				end
 
 				# Read a message from the connection stream.
 				#
 				# @returns [Hash, nil] The parsed message or nil if stream is closed.
 				def read
-					length_data = @stream&.read(4)
-					return nil unless length_data && length_data.bytesize == 4
-					
-					# Unpack 32-bit integer
-					length = length_data.unpack1("N")
-					
-					# Validate message size
-					if length > MAX_MESSAGE_SIZE
-						Console.error(self, "Message too large: #{length} bytes (max: #{MAX_MESSAGE_SIZE})")
-						return nil
-					end
-					
-					# Read the exact amount of data specified
-					data = @stream.read(length)
-					
-					unless data && data.bytesize == length
-						raise EOFError, "Failed to read complete message"
-					end
-					
-					@message_wrapper.unpack(data)
+					@message_wrapper.read
+				rescue EOFError, IOError
+					nil
 				end
 
 				# Iterate over all messages from the connection.

lib/async/container/supervisor/message_wrapper.rb

@@ -9,12 +9,25 @@ module Async
 	module Container
 		module Supervisor
 			class MessageWrapper
-				def initialize
+				def initialize(stream)
 					@factory = MessagePack::Factory.new
 
 					register_types
 
-					@packer = @factory.packer
+					@packer = @factory.packer(stream)
+					@unpacker = @factory.unpacker(stream)
+				end
+				
+				def write(message)
+					data = pack(message)
+					# Console.logger.info("Sending data: #{message.inspect}")
+					@packer.write(data)
+				end
+				
+				def read
+					data = @unpacker.read
+					# Console.logger.info("Received data: #{data.inspect}")
+					data
 				end
 
 				def pack(message)

test/async/container/connection.rb

@@ -21,12 +21,16 @@ def dispatch(call)
 describe Async::Container::Supervisor::Connection do
 	let(:stream) {StringIO.new}
 	let(:connection) {Async::Container::Supervisor::Connection.new(stream)}
-	let(:message_wrapper) {Async::Container::Supervisor::MessageWrapper.new}
+	let(:message_wrapper) {Async::Container::Supervisor::MessageWrapper.new(stream)}
+	
+	def write_message(message)
+		message_wrapper.write(message)
+		stream.rewind
+	end
 
 	with "dispatch" do
 		it "handles failed writes when dispatching a call" do
-			stream.write(JSON.dump({id: 1, do: :test}) << "\n")
-			stream.rewind
+			write_message({id: 1, do: :test})
 
 			expect(stream).to receive(:write).and_raise(IOError, "Write error")
 
@@ -44,8 +48,7 @@ def dispatch(call)
 		end
 
 		it "closes the queue when the connection fails" do
-			stream.write(JSON.dump({id: 1, do: :test}) << "\n")
-			stream.rewind
+			write_message({id: 1, do: :test})
 
 			expect(stream).to receive(:write).and_raise(IOError, "Write error")
 
@@ -146,48 +149,6 @@ def dispatch(call)
 		end
 	end
 
-	it "writes length-prefixed MessagePack data" do
-		connection.write(id: 1, do: :test)
-		
-		stream.rewind
-		
-		# Read 2-byte length prefix
-		length_data = stream.read(4)
-		expect(length_data.bytesize).to be == 4
-		
-		length = length_data.unpack1("N")
-		expect(length).to be > 0
-		
-		# Read MessagePack data
-		data = stream.read(length)
-		expect(data.bytesize).to be == length
-		
-		# Parse MessagePack
-		parsed = message_wrapper.unpack(data)
-		expect(parsed).to have_keys(
-			id: be == 1,
-			do: be == :test
-		)
-	end
-	
-	it "reads length-prefixed MessagePack data" do
-		# Create MessagePack data
-		message = {id: 1, do: "test"}
-		data = message_wrapper.pack(message)
-		
-		# Write with length prefix
-		stream.string = [data.bytesize].pack("N") + data
-		stream.rewind
-		
-		parsed = connection.read
-		
-		# Keys are symbols
-		expect(parsed).to have_keys(
-			id: be == 1,
-			do: be == "test"
-		)
-	end
-	
 	it "returns nil when stream is closed" do
 		stream.string = ""
 		stream.rewind

test/async/container/server.rb

@@ -10,25 +10,12 @@
 	include Async::Container::Supervisor::AServer
 	include Sus::Fixtures::Console::CapturedLogger
 
-	let(:message_wrapper) {Async::Container::Supervisor::MessageWrapper.new}
-	
-	# Helper to write length-prefixed MessagePack data
 	def write_message(stream, message)
-		data = message_wrapper.pack(message)
-		stream.write([data.bytesize].pack("N") + data)
-		stream.flush
+		Async::Container::Supervisor::MessageWrapper.new(stream).write(message)
 	end
 
-	# Helper to read length-prefixed MessagePack data
 	def read_message(stream)
-		length_data = stream.read(4)
-		return nil unless length_data && length_data.bytesize == 4
-		
-		length = length_data.unpack1("N")
-		data = stream.read(length)
-		return nil unless data && data.bytesize == length
-		
-		message_wrapper.unpack(data)
+		Async::Container::Supervisor::MessageWrapper.new(stream).read
 	end
 
 	it "can handle unexpected failures" do