avoid partial writes on IO.vmsplice unless F_NONBLOCK

vmsplice-ing a partial array of strings into a pipe is not very
useful under Ruby, so wait for I/O availability if the pipe is
full.  This code/logic was also contributed to the io-extra
gem (by me) for the IO.writev implementation.

Signed-off-by: Eric Wong <normalperson@yhbt.net>

1 files changed, 60 insertions, 6 deletions

diff --git a/ext/io_splice/io_splice_ext.c b/ext/io_splice/io_splice_ext.c
index 69c0288..9531633 100644
--- a/ext/io_splice/io_splice_ext.c
+++ b/ext/io_splice/io_splice_ext.c
@@ -245,6 +245,36 @@ do { \
         } \
 } while (0)
 
+static void advance_vmsplice_args(struct vmsplice_args *a, long n)
+{
+        struct iovec *new_iov = a->iov;
+        int i;
+
+        /* skip over iovecs we've already written completely */
+        for (i = 0; i < a->nr_segs; i++, new_iov++) {
+                if (n == 0)
+                        break;
+                /*
+                 * partially written iov,
+                 * modify and retry with current iovec in
+                 * front
+                 */
+                if (new_iov->iov_len > (size_t)n) {
+                        VALUE base = (VALUE)new_iov->iov_base;
+
+                        new_iov->iov_len -= n;
+                        new_iov->iov_base = (void *)(base + n);
+                        break;
+                }
+
+                n -= new_iov->iov_len;
+        }
+
+        /* setup to retry without the already-written iovecs */
+        a->nr_segs -= i;
+        a->iov = new_iov;
+}
+
 /*
  * call-seq:
  *   IO.vmsplice(fd, string_array, flags) => integer
@@ -260,7 +290,7 @@ do { \
  */
 static VALUE my_vmsplice(VALUE self, VALUE fd, VALUE data, VALUE flags)
 {
-        long n;
+        long rv = 0;
         ssize_t left;
         struct vmsplice_args a;
 
@@ -268,11 +298,35 @@ static VALUE my_vmsplice(VALUE self, VALUE fd, VALUE data, VALUE flags)
         a.fd = NUM2INT(fd);
         a.flags = NUM2UINT(flags);
 
-        n = (long)nb_io_run(nogvl_vmsplice, &a, a.flags);
-        if (n < 0)
-                rb_sys_fail("vmsplice");
-
-        return LONG2NUM(n);
+        for (;;) {
+                long n = (long)nb_io_run(nogvl_vmsplice, &a, a.flags);
+
+                if (n < 0) {
+                        if (errno == EAGAIN) {
+                                if (a.flags & SPLICE_F_NONBLOCK)
+                                        rb_sys_fail("vmsplice");
+                                else if (rb_io_wait_writable(a.fd))
+                                        continue;
+                                /* fall through on error */
+                        }
+                        /*
+                         * unlikely to hit this case, return the
+                         * already written bytes, we'll let the next
+                         * write (or close) fail instead
+                         */
+                        if (rv > 0)
+                                break;
+                        rb_sys_fail("vmsplice");
+                }
+
+                rv += n;
+                left -= n;
+                if (left == 0)
+                        break;
+                advance_vmsplice_args(&a, n);
+        }
+
+        return LONG2NUM(rv);
 }
 
 void Init_io_splice_ext(void)