diff --git a/lib/posix/spawn.rb b/lib/posix/spawn.rb index ac01638..10f55bc 100644 --- a/lib/posix/spawn.rb +++ b/lib/posix/spawn.rb @@ -337,6 +337,10 @@ class MaximumOutputExceeded < StandardError class TimeoutExceeded < StandardError end + # Exception raised when output streaming is aborted early. + class Aborted < StandardError + end + private # Turns the various varargs incantations supported by Process::spawn into a diff --git a/lib/posix/spawn/child.rb b/lib/posix/spawn/child.rb index 98e3ae2..0140583 100644 --- a/lib/posix/spawn/child.rb +++ b/lib/posix/spawn/child.rb @@ -95,6 +95,31 @@ def initialize(*args) @options[:pgroup] = true end @options.delete(:chdir) if @options[:chdir].nil? + + @out, @err = "", "" + + @stdout_stream = Proc.new do |chunk| + @out << chunk + + true + end + + @stderr_stream = Proc.new do |chunk| + @err << chunk + + true + end + + if streams = @options.delete(:streams) + if streams[:stdout] + @stdout_stream = streams[:stdout] + end + + if streams[:stderr] + @stderr_stream = streams[:stderr] + end + end + exec! if !@options.delete(:noexec) end @@ -223,6 +248,10 @@ def read_and_write(input, stdin, stdout, stderr, timeout=nil, max=nil) slice_method = input.respond_to?(:byteslice) ? :byteslice : :slice t = timeout + streams = {stdout => @stdout_stream, stderr => @stderr_stream} + + bytes_seen = 0 + chunk_buffer = "" while readers.any? || writers.any? ready = IO.select(readers, writers, readers + writers, t) raise TimeoutExceeded if ready.nil? @@ -244,9 +273,12 @@ def read_and_write(input, stdin, stdout, stderr, timeout=nil, max=nil) # read from stdout and stderr streams ready[0].each do |fd| - buf = (fd == stdout) ? @out : @err begin - buf << fd.readpartial(BUFSIZE) + fd.readpartial(BUFSIZE, chunk_buffer) + + raise Aborted unless streams[fd].call(chunk_buffer) + + bytes_seen += chunk_buffer.bytesize rescue Errno::EAGAIN, Errno::EINTR rescue EOFError readers.delete(fd) @@ -262,12 +294,10 @@ def read_and_write(input, stdin, stdout, stderr, timeout=nil, max=nil) end # maybe we've hit our max output - if max && ready[0].any? && (@out.size + @err.size) > max + if max && ready[0].any? && bytes_seen > max raise MaximumOutputExceeded end end - - [@out, @err] end # Wait for the child process to exit diff --git a/test/test_child.rb b/test/test_child.rb index e869c64..e09a458 100644 --- a/test/test_child.rb +++ b/test/test_child.rb @@ -131,7 +131,7 @@ def test_max_with_stubborn_child_pgroup_kill def test_max_with_partial_output p = Child.build('yes', :max => 100_000) - assert_nil p.out + assert_empty p.out assert_raises MaximumOutputExceeded do p.exec! end @@ -224,6 +224,79 @@ def test_utf8_input_long assert p.success? end + def test_streaming_stdout + stdout_buff = "" + stdout_stream = Proc.new do |chunk| + stdout_buff << chunk + end + + input = "hello!" + p = Child.new('cat', :input => input, :streams => { + :stdout => stdout_stream + }) + + assert p.success? + assert_equal input, stdout_buff + end + + def test_streaming_stderr + stderr_buff = "" + stderr_stream = Proc.new do |chunk| + stderr_buff << chunk + end + + p = Child.new('ls', '-?', :streams => { + :stderr => stderr_stream + }) + + refute p.success? + assert stderr_buff.size > 0 + end + + def test_streaming_stdout_aborted + stdout_stream = Proc.new do |chunk| + false + end + + input = "hello!" + assert_raises POSIX::Spawn::Aborted do + p = Child.new('cat', :input => input, :streams => { + :stdout => stdout_stream + }) + end + end + + def test_streaming_stderr_aborted + stderr_stream = Proc.new do |chunk| + false + end + + input = "hello!" + assert_raises POSIX::Spawn::Aborted do + p = Child.new('ls', '-?', :streams => { + :stderr => stderr_stream + }) + end + end + + def test_streaming_stdout_over_default_bufsize_boundary + stdout_buff = "" + chunk_count = 0 + stdout_stream = Proc.new do |chunk| + chunk_count += 1 + stdout_buff << chunk + end + + limit = Child::BUFSIZE * 2 + + assert_raises POSIX::Spawn::MaximumOutputExceeded do + Child.new('yes', :streams => {:stdout => stdout_stream}, :max => limit) + end + + assert_equal limit, stdout_buff.bytesize + assert chunk_count > 1 + end + ## # Assertion Helpers