Better ownership transfer.

Author: samuel-williams-shopify

guides/getting-started/readme.md

@@ -163,6 +163,16 @@ Fiber.schedule do
 end
 ~~~
 
+Transfer automatically recurses through instance variables:
+
+~~~ ruby
+# If request contains a body that's also being tracked:
+request = Request.new
+request.body = Body.new  # Body is tracked
+
+Async::Safe.transfer(request)  # Transfers both request AND body
+~~~
+
 ## Integration with Tests
 
 Add to your test helper (e.g., `config/sus.rb` or `spec/spec_helper.rb`):

lib/async/safe/builtins.rb

@@ -8,6 +8,34 @@
 # Note: Immutable values (nil, true, false, integers, symbols, etc.) are already
 # handled by the frozen? check in the monitor and don't need to be listed here.
 
+# Arrays contain references to other objects that may need transfer:
+class Array
+	ASYNC_SAFE = false
+	
+	def self.async_safe_traverse(instance, &block)
+		instance.each(&block)
+	end
+end
+
+# Hashes contain keys and values that may need transfer:
+class Hash
+	ASYNC_SAFE = false
+	
+	def self.async_safe_traverse(instance, &block)
+		instance.each_key(&block)
+		instance.each_value(&block)
+	end
+end
+
+# Sets contain elements that may need transfer:
+class Set
+	ASYNC_SAFE = false
+	
+	def self.async_safe_traverse(instance, &block)
+		instance.each(&block)
+	end
+end
+
 module Async
 	module Safe
 		# Automatically transfers ownership of objects when they are removed from a Thread::Queue.

lib/async/safe/monitor.rb

@@ -87,13 +87,53 @@ def disable!
 
 			# Explicitly transfer ownership of objects to the current fiber.
 			#
+			# Also recursively transfers ownership of any tracked instance variables and
+			# objects contained in collections (Array, Hash, Set).
+			#
 			# @parameter objects [Array(Object)] The objects to transfer.
 			def transfer(*objects)
+				current = Fiber.current
+				visited = Set.new
+				
+				# Traverse object graph (outside mutex to avoid deadlock):
+				objects.each do |object|
+					traverse_objects(object, visited)
+				end
+				
+				# Transfer all visited objects (convert to array to avoid triggering TracePoint in sync block):
+				objects_to_transfer = visited.to_a
+				
 				@mutex.synchronize do
-					current = Fiber.current
-					
-					objects.each do |object|
-						@owners[object] = current
+					objects_to_transfer.each do |object|
+						@owners[object] = current if @owners.key?(object)
+					end
+				end
+			end
+			
+			private
+			
+			# Traverse the object graph and collect all reachable objects.
+			#
+			# @parameter object [Object] The object to traverse.
+			# @parameter visited [Set] Set of visited objects (object references, not IDs).
+			def traverse_objects(object, visited)
+				# Avoid circular references:
+				return if visited.include?(object)
+				
+				# Skip objects that don't need traversing:
+				return if object.frozen? or object.is_a?(Module)
+				
+				# Skip async-safe (shareable) objects - they're not owned:
+				klass = object.class
+				return if klass.async_safe?(nil)
+				
+				# Mark as visited:
+				visited << object
+				
+				# Recurse through custom traversal if defined:
+				if klass.respond_to?(:async_safe_traverse)
+					klass.async_safe_traverse(object) do |element|
+						traverse_objects(element, visited)
 					end
 				end
 			end
@@ -105,7 +145,7 @@ def check_access(trace_point)
 				object = trace_point.self
 
 				# Skip tracking class/module methods:
-				return if object.is_a?(Class) || object.is_a?(Module)
+				return if object.is_a?(Module)
 
 				# Skip frozen objects:
 				return if object.frozen?

releases.md

@@ -1,5 +1,9 @@
 # Releases
 
+## Unreleased
+
+  - Improved `Async::Safe.transfer` to recursively transfer ownership of tracked instance variables.
+
 ## v0.3.2
 
   - Better error message.

test/async/safe.rb

@@ -73,6 +73,33 @@ def push(item)
 		end.not.to raise_exception
 	end
 
+	it "recursively transfers owned instance variables" do
+		inner_class = Class.new do
+			async_safe!(false)
+			def read; "data"; end
+		end
+		
+		outer_class = Class.new do
+			async_safe!(false)
+			attr_reader :inner
+			def initialize(inner); @inner = inner; end
+		end
+		
+		inner = inner_class.new
+		outer = outer_class.new(inner)
+		
+		# Both owned by main fiber:
+		inner.read
+		outer.inner
+		
+		Fiber.new do
+			Async::Safe.transfer(outer)
+			
+			# Both outer and inner should be transferred:
+			outer.inner.read
+		end.resume
+	end
+	
 	it "allows access after ownership transfer" do
 		body = body_class.new(["a", "b"])
 		body.read  # Main fiber owns it