add IO::Splice::WAITALL flag support

This allows splice-in-full and tee-in-full behavior
to simplify user code.

2 files changed, 114 insertions, 7 deletions

diff --git a/ext/io_splice/io_splice_ext.c b/ext/io_splice/io_splice_ext.c
index a1f511c..80c8cde 100644
--- a/ext/io_splice/io_splice_ext.c
+++ b/ext/io_splice/io_splice_ext.c
@@ -13,6 +13,7 @@
 #include <sys/utsname.h>
 
 static VALUE sym_EAGAIN;
+#define WAITALL 0x4000000
 
 #ifndef F_LINUX_SPECIFIC_BASE
 #  define F_LINUX_SPECIFIC_BASE 1024
@@ -135,10 +136,12 @@ static VALUE nogvl_splice(void *ptr)
 
 static ssize_t do_splice(int argc, VALUE *argv, unsigned dflags)
 {
-        off_t i, o;
+        off_t i = 0, o = 0;
         VALUE fd_in, off_in, fd_out, off_out, len, flags;
         struct splice_args a;
         ssize_t bytes;
+        ssize_t total = 0;
+        unsigned waitall;
 
         rb_scan_args(argc, argv, "51",
                      &fd_in, &off_in, &fd_out, &off_out, &len, &flags);
@@ -147,14 +150,40 @@ static ssize_t do_splice(int argc, VALUE *argv, unsigned dflags)
         a.off_out = NIL_P(off_out) ? NULL : (o = NUM2OFFT(off_out), &o);
         a.len = NUM2SIZET(len);
         a.flags = NIL_P(flags) ? dflags : NUM2UINT(flags) | dflags;
+        waitall = a.flags & WAITALL;
+        if (waitall)
+                a.flags ^= WAITALL;
 
-        do {
+        for (;;) {
                 a.fd_in = check_fileno(fd_in);
                 a.fd_out = check_fileno(fd_out);
                 bytes = (ssize_t)io_run(nogvl_splice, &a);
-        } while (bytes == -1 && errno == EINTR);
+                if (bytes == -1) {
+                        if (errno == EINTR)
+                                continue;
+                        if (waitall && errno == EAGAIN) {
+                                rb_io_wait_readable(check_fileno(fd_in));
+                                errno = EAGAIN;
+                                rb_io_wait_writable(check_fileno(fd_out));
+                                continue;
+                        }
+                        if (total > 0)
+                                return total;
+                        return bytes;
+                } else if (bytes == 0) {
+                        break;
+                } else if (waitall) {
+                        total += bytes;
+                        if ((a.len -= bytes) == 0)
+                                return total;
+                        i += bytes;
+                        o += bytes;
+                } else {
+                        return bytes;
+                }
+        }
 
-        return bytes;
+        return total;
 }
 
 /*
@@ -257,18 +286,44 @@ static ssize_t do_tee(int argc, VALUE *argv, unsigned dflags)
         VALUE fd_in, fd_out, len, flags;
         struct tee_args a;
         ssize_t bytes;
+        ssize_t total = 0;
+        unsigned waitall;
 
         rb_scan_args(argc, argv, "31", &fd_in, &fd_out, &len, &flags);
         a.len = (size_t)NUM2SIZET(len);
         a.flags = NIL_P(flags) ? dflags : NUM2UINT(flags) | dflags;
+        waitall = a.flags & WAITALL;
+        if (waitall)
+                a.flags ^= WAITALL;
 
-        do {
+        for (;;) {
                 a.fd_in = check_fileno(fd_in);
                 a.fd_out = check_fileno(fd_out);
                 bytes = (ssize_t)io_run(nogvl_tee, &a);
-        } while (bytes == -1 && errno == EINTR);
+                if (bytes == -1) {
+                        if (errno == EINTR)
+                                continue;
+                        if (waitall && errno == EAGAIN) {
+                                rb_io_wait_readable(check_fileno(fd_in));
+                                errno = EAGAIN;
+                                rb_io_wait_writable(check_fileno(fd_out));
+                                continue;
+                        }
+                        if (total > 0)
+                                return total;
+                        return bytes;
+                } else if (bytes == 0) {
+                        break;
+                } else if (waitall) {
+                        total += bytes;
+                        if ((a.len -= bytes) == 0)
+                                return total;
+                } else {
+                        return bytes;
+                }
+        }
 
-        return bytes;
+        return total;
 }
 
 /*
@@ -573,6 +628,7 @@ void Init_io_splice_ext(void)
          * re-added for FUSE filesystems only in Linux 2.6.35.
          */
         rb_define_const(mSplice, "F_MOVE", UINT2NUM(SPLICE_F_MOVE));
+        assert(WAITALL != SPLICE_F_MOVE && "WAITALL == F_MOVE");
 
         /*
          * Do not block on pipe I/O.  This flag only affects the pipe(s)
@@ -586,6 +642,7 @@ void Init_io_splice_ext(void)
          * out of them.
          */
         rb_define_const(mSplice, "F_NONBLOCK", UINT2NUM(SPLICE_F_NONBLOCK));
+        assert(WAITALL != SPLICE_F_NONBLOCK && "WAITALL == F_NONBLOCK");
 
         /*
          * Indicate that there may be more data coming into the outbound
@@ -593,12 +650,23 @@ void Init_io_splice_ext(void)
          * frames from sockets.  Currently only used with splice.
          */
         rb_define_const(mSplice, "F_MORE", UINT2NUM(SPLICE_F_MORE));
+        assert(WAITALL != SPLICE_F_MORE && "WAITALL == F_MORE");
 
         /*
          * Only usable by vmsplice.  This flag probably not useful in the
          * context of Ruby applications which cannot control alignment.
          */
         rb_define_const(mSplice, "F_GIFT", UINT2NUM(SPLICE_F_GIFT));
+        assert(WAITALL != SPLICE_F_GIFT && "WAITALL == F_GIFT");
+
+        /*
+         * Retry until the requested transfer is complete, this will
+         * cause IO.splice/IO.tee to never return less than the requested
+         * transfer size unless an error occored.
+         *
+         * IO.vmsplice always defaults to this behavior.
+         */
+        rb_define_const(mSplice, "WAITALL", UINT2NUM(WAITALL));
 
         /*
          * The maximum size of an atomic write to a pipe
diff --git a/test/test_io_splice_in_full.rb b/test/test_io_splice_in_full.rb
new file mode 100644
index 0000000..4426351
--- /dev/null
+++ b/test/test_io_splice_in_full.rb
@@ -0,0 +1,39 @@
+# -*- encoding: binary -*-
+require 'test/unit'
+require 'tempfile'
+$-w = true
+require 'io/splice'
+Thread.abort_on_exception = true
+
+class Test_IO_Splice_In_Full < Test::Unit::TestCase
+  def test_splice_in_full
+    rd, wr = IO.pipe
+    tmp = Tempfile.new 'splice-read'
+    Thread.new do
+      111.times do
+        sleep 0.002
+        wr.write "HIHIHI"
+      end
+    end
+    nr = IO.splice rd, nil, tmp, nil, 666, IO::Splice::WAITALL
+    assert_equal 666, nr
+    tmp.rewind
+    assert_equal "HIHIHI" * 111, tmp.read
+  end
+
+  def test_tee_in_full
+    rd, wr = IO.pipe
+    a, b = IO.pipe
+    thr = Thread.new { a.read(666) }
+    Thread.new do
+      111.times do
+        sleep 0.001
+        wr.syswrite "HIHIHI"
+      end
+    end
+    nr = IO.tee rd, b, 666, IO::Splice::WAITALL
+    assert_equal 666, nr
+    thr.join
+    assert_equal "HIHIHI" * 111, thr.value
+  end
+end if defined?(RUBY_ENGINE)