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.

1 files changed, 34 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);
 
         /*