From 5e73ae9c477d0ba64f497f0517be20e65ba0c253 Mon Sep 17 00:00:00 2001 From: Eric Wong Date: Mon, 9 May 2011 17:57:10 -0700 Subject: [PATCH] add IO::Splice::WAITALL flag support This allows splice-in-full and tee-in-full behavior to simplify user code. --- ext/io_splice/io_splice_ext.c | 82 ++++++++++++++++++++++++++++++++++++++---- test/test_io_splice_in_full.rb | 39 ++++++++++++++++++++ 2 files changed, 114 insertions(+), 7 deletions(-) create mode 100644 test/test_io_splice_in_full.rb 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 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) -- 2.11.4.GIT