refactor connection to use MessagePack instead of JSON

Author: ggmichaelgo

bake/async/container/supervisor.rb

@@ -44,7 +44,7 @@ def memory_sample(duration: 10, connection_id:)
 		operation = {do: :memory_sample, duration: duration}
 
 		# Use the forward operation to proxy the request to a worker:
-		return connection.call(do: :forward, operation: operation, connection_id: connection_id)
+		connection.call(do: :forward, operation: operation, connection_id: connection_id)
 	end
 end
 

examples/simple/simple.rb

@@ -33,7 +33,7 @@ def setup(container)
 				Console.info(self, "Exiting...")
 			end
 		end
-	end	
+	end
 end
 
 service "sleep" do

lib/async/container/supervisor/connection.rb

@@ -5,6 +5,7 @@
 
 require "json"
 require "async"
+require_relative "message_wrapper"
 
 module Async
 	module Container
@@ -13,6 +14,8 @@ module Supervisor
 			#
 			# Handles message passing, call/response patterns, and connection lifecycle.
 			class Connection
+				MAX_MESSAGE_SIZE = 2 ** 32 - 1
+				
 				# Represents a remote procedure call over a connection.
 				#
 				# Manages the call lifecycle, response queueing, and completion signaling.
@@ -231,7 +234,7 @@ def initialize(stream, id = 0, **state)
 					@stream = stream
 					@id = id
 					@state = state
-					
+					@message_wrapper = MessageWrapper.new
 					@reader = nil
 					@calls = {}
 				end
@@ -251,19 +254,42 @@ def next_id
 
 				# Write a message to the connection stream.
 				#
+				# Uses a length-prefixed protocol: 2-byte length header (big-endian) followed by data.
+				# This allows MessagePack data to contain newlines without breaking message boundaries.
+				#
 				# @parameter message [Hash] The message to write.
 				def write(**message)
-					@stream.write(JSON.dump(message) << "\n")
+					data = @message_wrapper.pack(message)
+					# Write 4-byte length prefix
+					data_length = [data.bytesize].pack("N")
+					@stream.write(data_length + data)
 					@stream.flush
 				end
 
 				# Read a message from the connection stream.
 				#
 				# @returns [Hash, nil] The parsed message or nil if stream is closed.
 				def read
-					if line = @stream&.gets
-						JSON.parse(line, symbolize_names: true)
+					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)
 				end
 
 				# Iterate over all messages from the connection.

lib/async/container/supervisor/message_wrapper.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2025, by Samuel Williams.
+
+require "msgpack"
+
+module Async
+	module Container
+		module Supervisor
+			class MessageWrapper
+				def initialize
+					@factory = MessagePack::Factory.new
+					
+					register_types
+					
+					@packer = @factory.packer
+				end
+				
+				def pack(message)
+					@packer.clear
+					normalized_message = normalize(message)
+					@packer.pack(normalized_message)
+					@packer.full_pack
+				end
+				
+				def unpack(data)
+					@factory.unpack(data)
+				end
+				
+				private
+				
+				def normalize(obj)
+					case obj
+					when Hash
+						obj.transform_values{|v| normalize(v)}
+					when Array
+						obj.map{|v| normalize(v)}
+					else
+						if obj.respond_to?(:as_json)
+							normalize(obj.as_json)
+						else
+							obj
+						end
+					end
+				end
+				
+				def register_types
+					@factory.register_type(0x00, Symbol)
+					
+					@factory.register_type(
+						0x01,
+						Exception,
+						packer: self.method(:pack_exception),
+						unpacker: self.method(:unpack_exception),
+						recursive: true,
+					)
+					
+					@factory.register_type(
+						0x02,
+						Class,
+						packer: ->(klass) {klass.name},
+						unpacker: ->(name) {name},
+					)
+					
+					@factory.register_type(
+						MessagePack::Timestamp::TYPE,
+						Time,
+						packer: MessagePack::Time::Packer,
+						unpacker: MessagePack::Time::Unpacker
+					)
+				end
+				
+				def pack_exception(exception)
+					[exception.class.name, exception.message, exception.backtrace].pack("A*")
+				end
+				
+				def unpack_exception(data)
+					klass, message, backtrace = data.unpack("A*A*A*")
+					klass = Object.const_get(klass)
+					
+					exception = klass.new(message)
+					exception.set_backtrace(backtrace)
+					
+					return exception
+				end
+			end
+		end
+	end
+end

test/async/container/connection.rb

@@ -6,6 +6,7 @@
 require "async/container/supervisor/connection"
 require "sus/fixtures/async/scheduler_context"
 require "stringio"
+require "msgpack"
 
 class TestTarget
 	def initialize(&block)
@@ -20,6 +21,7 @@ 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}
 
 	with "dispatch" do
 		it "handles failed writes when dispatching a call" do
@@ -88,16 +90,16 @@ def dispatch(call)
 			parsed = JSON.parse(json)
 
 			expect(parsed).to have_keys(
-				"do" => be == "test",
-				"data" => be == "value"
-			)
+					"do" => be == "test",
+					"data" => be == "value"
+				)
 		end
 
 		it "can get call message via as_json" do
 			expect(test_call.as_json).to have_keys(
-				do: be == :test,
-				data: be == "value"
-			)
+					do: be == :test,
+					data: be == "value"
+				)
 		end
 
 		it "can iterate over call responses with each" do
@@ -121,11 +123,11 @@ def dispatch(call)
 
 			response = test_call.pop
 			expect(response).to have_keys(
-				id: be == 1,
-				finished: be == true,
-				failed: be == true,
-				error: be == "Something went wrong"
-			)
+					id: be == 1,
+					finished: be == true,
+					failed: be == true,
+					error: be == "Something went wrong"
+				)
 
 			expect(test_call.closed?).to be == true
 		end
@@ -144,30 +146,43 @@ def dispatch(call)
 		end
 	end
 
-	it "writes JSON with newline" do
+	it "writes length-prefixed MessagePack data" do
 		connection.write(id: 1, do: :test)
 
 		stream.rewind
-		output = stream.read
 
-		# Check it's valid JSON with a newline
-		expect(output[-1]).to be == "\n"
+		# Read 2-byte length prefix
+		length_data = stream.read(4)
+		expect(length_data.bytesize).to be == 4
 
-		parsed = JSON.parse(output.chomp)
+		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"
+			id: be == 1,
+			do: be == :test
 		)
 	end
 
-	it "parses JSON lines" do
-		stream.string = JSON.dump({id: 1, do: "test"}) << "\n"
+	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
 
-		message = connection.read
+		parsed = connection.read
 
-		# Connection.read uses symbolize_names: true (keys are symbols, values are as-is)
-		expect(message).to have_keys(
+		# Keys are symbols
+		expect(parsed).to have_keys(
 			id: be == 1,
 			do: be == "test"
 		)