add IO.trytee interface

This allows more-efficient use of non-blocking I/O with
tee(2) and will return :EAGAIN instead of generating an
expensive Errno::EAGAIN exception.

2 files changed, 60 insertions, 13 deletions

diff --git a/ext/io_splice/io_splice_ext.c b/ext/io_splice/io_splice_ext.c
index 79d96ee..c97435e 100644
--- a/ext/io_splice/io_splice_ext.c
+++ b/ext/io_splice/io_splice_ext.c
@@ -222,6 +222,20 @@ static VALUE nogvl_tee(void *ptr)
         return (VALUE)tee(a->fd_in, a->fd_out, a->len, a->flags);
 }
 
+static long do_tee(int argc, VALUE *argv, unsigned dflags)
+{
+        VALUE fd_in, fd_out, len, flags;
+        struct tee_args a;
+
+        rb_scan_args(argc, argv, "31", &fd_in, &fd_out, &len, &flags);
+        a.fd_in = my_fileno(fd_in);
+        a.fd_out = my_fileno(fd_out);
+        a.len = (size_t)NUM2ULONG(len);
+        a.flags = NIL_P(flags) ? dflags : NUM2UINT(flags) | dflags;
+
+        return (long)nb_io_run(nogvl_tee, &a, a.flags);
+}
+
 /*
  * call-seq:
  *   IO.tee(fd_in, fd_out, len, flags) => integer
@@ -241,19 +255,10 @@ static VALUE nogvl_tee(void *ptr)
  * See manpage for full documentation:
  * http://kernel.org/doc/man-pages/online/pages/man2/tee.2.html
  */
-static VALUE my_tee(VALUE self,
-        VALUE fd_in, VALUE fd_out,
-        VALUE len, VALUE flags)
+static VALUE my_tee(int argc, VALUE *argv, VALUE self)
 {
-        long n;
-        struct tee_args a = {
-                .fd_in = my_fileno(fd_in),
-                .fd_out = my_fileno(fd_out),
-                .len = (size_t)NUM2ULONG(len),
-                .flags = NUM2UINT(flags),
-        };
-
-        n = (long)nb_io_run(nogvl_tee, &a, a.flags);
+        long n = do_tee(argc, argv, 0);
+
         if (n == 0)
                 rb_eof_error();
         if (n < 0)
@@ -262,6 +267,21 @@ static VALUE my_tee(VALUE self,
         return LONG2NUM(n);
 }
 
+static VALUE trytee(int argc, VALUE *argv, VALUE self)
+{
+        long n = do_tee(argc, argv, SPLICE_F_NONBLOCK);
+
+        if (n == 0)
+                return Qnil;
+        if (n < 0) {
+                if (errno == EAGAIN)
+                        return sym_EAGAIN;
+                rb_sys_fail("tee");
+        }
+
+        return LONG2NUM(n);
+}
+
 struct vmsplice_args {
         int fd;
         struct iovec *iov;
@@ -472,7 +492,8 @@ void Init_io_splice_ext(void)
 
         rb_define_singleton_method(rb_cIO, "splice", my_splice, -1);
         rb_define_singleton_method(rb_cIO, "trysplice", trysplice, -1);
-        rb_define_singleton_method(rb_cIO, "tee", my_tee, 4);
+        rb_define_singleton_method(rb_cIO, "tee", my_tee, -1);
+        rb_define_singleton_method(rb_cIO, "trytee", trytee, -1);
         rb_define_singleton_method(rb_cIO, "vmsplice", my_vmsplice, 3);
 
         /*
diff --git a/test/test_io_splice.rb b/test/test_io_splice.rb
index f5bce6a..1d73f70 100644
--- a/test/test_io_splice.rb
+++ b/test/test_io_splice.rb
@@ -193,6 +193,19 @@ class Test_IO_Splice < Test::Unit::TestCase
     assert_equal str, rda.sysread(5)
   end
 
+  def test_trytee
+    str = 'abcde'
+    size = 5
+    rda, wra = IO.pipe
+    rdb, wrb = IO.pipe
+
+    assert_nothing_raised { wra.syswrite(str) }
+    nr = IO.trytee(rda, wrb, size, 0)
+    assert_equal 5, nr
+    assert_equal str, rdb.sysread(5)
+    assert_equal str, rda.sysread(5)
+  end
+
   def test_tee_eof
     rda, wra = IO.pipe
     rdb, wrb = IO.pipe
@@ -200,6 +213,13 @@ class Test_IO_Splice < Test::Unit::TestCase
     assert_raises(EOFError) { IO.tee(rda.fileno, wrb.fileno, 4096, 0) }
   end
 
+  def test_trytee_eof
+    rda, wra = IO.pipe
+    rdb, wrb = IO.pipe
+    wra.close
+    assert_nil IO.trytee(rda, wrb, 4096)
+  end
+
   def test_tee_nonblock
     rda, wra = IO.pipe
     rdb, wrb = IO.pipe
@@ -208,6 +228,12 @@ class Test_IO_Splice < Test::Unit::TestCase
     }
   end
 
+  def test_trytee_nonblock
+    rda, wra = IO.pipe
+    rdb, wrb = IO.pipe
+    assert_equal :EAGAIN, IO.trytee(rda, wrb, 4096)
+  end
+
   def test_tee_io
     str = 'abcde'
     size = 5