Add "drain" callback and allow to write before started

Author: Jan Henning Thorsen

Changes

@@ -2,6 +2,8 @@ Revision history for Mojo-IOLoop-ReadWriteFork
 
 0.06
        * Fix "Use of uninitialized value in numeric eq (==) at ReadWriteFork.pm line 182."
+       * Add support for write() before child process has started
+       * Add support for "drain" callback to write()
 
 0.05     Wed Feb 19 13:29:54 2014
        * Fix "read" event cannot change ERRNO from sysread()

lib/Mojo/IOLoop/ReadWriteFork.pm

@@ -167,10 +167,11 @@ sub _start {
   elsif($pid) { # parent ===================================================
     warn "[$pid] Child starting ($args->{program} @{$args->{program_args}})\n" if DEBUG;
     $self->{pid} = $pid;
-    $self->{stdin_write} = $stdin_write;
     $self->{stdout_read} = $stdout_read;
     $stdout_read->close_slave if defined $stdout_read and UNIVERSAL::isa($stdout_read, 'IO::Pty');
 
+    $self->_stdin($stdin_write);
+
     Scalar::Util::weaken($self);
     $self->reactor->io($stdout_read => sub {
       $self->{stop} and return;
@@ -254,19 +255,33 @@ sub _setup_recurring_child_alive_check {
 
 =head2 write
 
-  $self->write($buffer);
+  $self = $self->write($chunk);
+  $self = $self->write($chunk, $cb);
+
+Used to write data to the child process STDIN. An optional callback will be
+called once STDIN is drained.
 
-Used to write data to the child process.
+Example:
+
+  $self->write("some data\n", sub {
+    my ($self) = @_;
+    $self->close_gracefully;
+  });
 
 =cut
 
 sub write {
-  my($self, $buffer) = @_;
+  my $self = shift;
+
+  if ($self->{stdin}) {
+    $self->{stdin}->write(@_);
+    warn "[${ \$self->pid }] Wrote buffer (" .url_escape($_[0]) .")\n" if DEBUG;
+  }
+  else {
+    push @{ $self->{stdin_buffer} }, [@_];
+  }
 
-  $self->{stdin_write} or return;
-  print { $self->{stdin_write} } $buffer;
-  $self->{stdin_write}->flush or die "Write buffer (" .url_escape($buffer) .") failed: $!";
-  warn "[${ \$self->pid }] Wrote buffer (" .url_escape($buffer) .")\n" if DEBUG;
+  $self;
 }
 
 =head2 kill
@@ -295,7 +310,8 @@ sub _cleanup {
   $reactor->watch($self->{stdout_read}, 0, 0) if $self->{stdout_read};
   $reactor->remove(delete $self->{stdout_read}) if $self->{stdout_read};
   $reactor->remove(delete $self->{delay}) if $self->{delay};
-  $reactor->remove(delete $self->{stdin_write}) if $self->{stdin_write};
+  $reactor->remove(delete $self->{delay}) if $self->{delay};
+  $self->{stdin}->close if $self->{stdin};
 }
 
 sub _read {
@@ -311,6 +327,19 @@ sub _read {
   $self->emit_safe(read => $buffer);
 }
 
+sub _stdin {
+  my ($self, $handle) = @_;
+  my $stream = Mojo::IOLoop::Stream->new($handle);
+
+  Scalar::Util::weaken($self);
+  $stream->reactor($self->reactor);
+  $stream->timeout(0);
+  $stream->start;
+
+  $self->{stdin} = $stream;
+  $self->write(@$_) for @{ delete $self->{stdin_buffer} || [] };
+}
+
 sub DESTROY { shift->_cleanup }
 
 =head1 AUTHOR

t/drain.t

@@ -0,0 +1,28 @@
+use Mojo::Base -strict;
+use Mojo::IOLoop::ReadWriteFork;
+use Test::More;
+use Test::Memory::Cycle;
+
+{
+  my $run = Mojo::IOLoop::ReadWriteFork->new;
+  my $drain = 0;
+  my $output = '';
+
+  $run->on(close => sub { Mojo::IOLoop->stop; });
+  $run->on(read => sub { $output .= $_[1]; });
+  $run->write("line one\n", sub { $drain++; });
+  $run->start(program => sub {
+    print sysread STDIN, my $buf, 1024;
+    print "\n$buf";
+    print "line two\n";
+  });
+
+  Mojo::IOLoop->timer(3 => sub { Mojo::IOLoop->stop }); # guard
+  Mojo::IOLoop->start;
+  memory_cycle_ok $run, 'no cycle after run';
+
+  like $output, qr/^9\nline one\nline two\n/, 'can write() before start()' or diag $output;
+  is $drain, 1, 'drain callback was called';
+}
+
+done_testing;