splice.rb in io_splice-2.1.0

- old
+ new

@@ -3,12 +3,12 @@
 
 class IO
 
   module Splice
 
-    # the version of IO::Splice, currently 2.0.0
-    VERSION = '2.0.0'
+    # the version of IO::Splice, currently 2.1.0
+    VERSION = '2.1.0'
 
     # The maximum default capacity of the pipe in bytes.
     # Under stock Linux, this is 65536 bytes as of 2.6.11, and 4096 before
     # We detect this at runtime as it is easy to recompile the kernel
     # and set a new value.
@@ -31,66 +31,89 @@
 
     # copies the contents of the IO object given by +src+ to +dst+
     # If len is specified, then only len bytes are copied.  Otherwise
     # the copy will be until EOF is reached on the +src+.
     # +src+ and +dst+ must be IO objects or respond to +to_io+
-    def self.copy_stream(src, dst, len = nil, src_offset = nil)
-      src.kind_of?(String) and src = File.open(src, 'rb')
-      dst.kind_of?(String) and dst = File.open(dst, 'wb')
+    def Splice.copy_stream(src, dst, len = nil, src_offset = nil)
+      close = []
+      src.kind_of?(String) and close << (src = File.open(src, 'rb'))
+      dst.kind_of?(String) and close << (dst = File.open(dst, 'wb'))
       src, dst = src.to_io, dst.to_io
       rv = len
       src.sysseek(src_offset) if src_offset
+      select_args = selectable(src, dst)
+
       if src.stat.pipe? || dst.stat.pipe?
         if len
-          while len > 0
-            nr = len > PIPE_CAPA ? PIPE_CAPA : len
-            nr = IO.splice(src, nil, dst, nil, nr, F_MOVE)
-            if nr == 0
-              raise EOFError, "unexpected EOF with #{len} bytes left"
-            end
-            len -= nr
-          end
+          len -= full(src, dst, len, select_args) until len == 0
         else
           rv = 0
           begin
-            nr = IO.splice(src, nil, dst, nil, PIPE_CAPA, F_MOVE)
-            rv += nr
+            rv += partial(src, dst, PIPE_CAPA, select_args)
           rescue EOFError
             break
           end while true
         end
       else
-        begin
-          r, w = IO.pipe
-          if len
-            while len > 0
-              nr = len > PIPE_CAPA ? PIPE_CAPA : len
-              nr_src = copy_stream(src, w, nr)
-              nr_src == nr or
-                raise RuntimeError, "short splice from: #{nr_src} != #{nr}"
-              nr_dst = copy_stream(r, dst, nr)
-              nr_dst == nr or
-                raise RuntimeError, "short splice to: #{nr_dst} != #{nr}"
-              len -= nr
-            end
-          else
-            rv = 0
-            begin
-              nr = IO.splice(src, nil, w, nil, PIPE_CAPA, F_MOVE)
-              rv += nr
-              begin
-                nr -= IO.splice(r, nil, dst, nil, nr, F_MOVE)
-              end while nr > 0
-            rescue EOFError
-              break
-            end while true
+        r, w = tmp = IO.pipe
+        close.concat(tmp)
+        if len
+          while len != 0
+            nr = partial(src, w, len, select_args)
+            len -= full(r, dst, nr, select_args)
           end
-        ensure
-          r.close
-          w.close
+        else
+          rv = 0
+          begin
+            nr = partial(src, w, PIPE_CAPA, select_args)
+            rv += full(r, dst, nr, select_args)
+          rescue EOFError
+            break
+          end while true
         end
       end
 
+      rv
+      ensure
+        close.each { |io| io.close }
+    end
+
+    # splice the full amount specified from +src+ to +dst+
+    # Either +dst+ or +src+ must be a pipe.  +dst+ and +src+
+    # may BOTH be pipes in Linux 2.6.31 or later.
+    # This will block and wait for IO completion of +len+
+    # bytes.  Returns the nubmer of bytes actually spliced (always +len+)
+    # The +_select_args+ parameter is reserved for internal use and
+    # may be removed in future versions.  Do not write code that
+    # depends on +_select_args+.
+    def Splice.full(src, dst, len, _select_args = selectable(src, dst))
+      nr = len
+      nr -= partial(src, dst, nr, _select_args) until nr == 0
+      len
+    end
+
+    # splice up to +len+ bytes from +src+ to +dst+.
+    # Either +dst+ or +src+ must be a pipe.  +dst+ and +src+
+    # may BOTH be pipes in Linux 2.6.31 or later.
+    # Returns the number of bytes actually spliced.
+    # Like IO#readpartial, this never returns Errno::EAGAIN
+    # The +_select_args+ parameter is reserved for internal use and
+    # may be removed in future versions.  Do not write code that
+    # depends on +_select_args+.
+    def Splice.partial(src, dst, len, _select_args = selectable(src, dst))
+      begin
+        IO.splice(src, nil, dst, nil, len, F_MOVE)
+      rescue Errno::EAGAIN
+        IO.select(*_select_args)
+        retry
+      end
+    end
+
+    # returns an array suitable for splat-ing to IO.select for blocking I/O
+    def Splice.selectable(src, dst)
+      rv = []
+      src.stat.pipe? or rv[0] = [ src ]
+      dst.stat.pipe? or rv[1] = [ dst ]
       rv
     end
 
   end
 end