lib/io/splice.rb

   1 # -*- encoding: binary -*-
   2 require 'io_splice_ext'
   3
   4 module IO::Splice
   5
   6   # The maximum default capacity of the pipe in bytes.
   7   # Under stock Linux, this is 65536 bytes as of 2.6.11, and 4096 before
   8   # We detect this at runtime as it is easy to recompile the kernel
   9   # and set a new value.
  10   # Starting with Linux 2.6.35, pipe capacity will be tunable
  11   # and this will only represent the default capacity of a
  12   # newly-created pipe.
  13   PIPE_CAPA = begin
  14     rd, wr = IO.pipe
  15     buf = ' ' * PIPE_BUF
  16     n = 0
  17     begin
  18       n += wr.write_nonblock(buf)
  19     rescue Errno::EAGAIN
  20       break
  21     end while true
  22     wr.close
  23     rd.close
  24     n
  25   end
  26
  27   # copies the contents of the IO object given by +src+ to +dst+
  28   # If +len+ is specified, then only +len+ bytes are copied and
  29   # +EOFError+ is raised if fewer than +len+ bytes could be copied.
  30   # Otherwise the copy will be until EOF is reached on the +src+.
  31   # +src+ and +dst+ must be IO objects or respond to +to_io+
  32   #
  33   # This is nearly a drop-in replacement for IO.copy_stream (in Ruby 1.9)
  34   # but does not take into account userspace I/O buffers nor IO-like
  35   # objects with no underlying file descriptor (e.g. StringIO).
  36   def self.copy_stream(src, dst, len = nil, src_offset = nil)
  37     close = []
  38     src.kind_of?(String) and close << (src = File.open(src))
  39     dst.kind_of?(String) and close << (dst = File.open(dst, "w"))
  40     src, dst = src.to_io, dst.to_io
  41     rv = len
  42     select_args = selectable(src, dst)
  43
  44     if src.stat.pipe? || dst.stat.pipe?
  45       if len
  46         len -= full(src, dst, len, src_offset, select_args) until len == 0
  47       else
  48         rv = 0
  49         while n = partial(src, dst, PIPE_CAPA, src_offset, select_args)
  50           rv += n
  51         end
  52       end
  53     else
  54       r, w = tmp = IO.pipe
  55       close.concat(tmp)
  56       if len
  57         while len != 0 && n = partial(src, w, len, src_offset, select_args)
  58           len -= full(r, dst, n, nil, select_args)
  59         end
  60       else
  61         rv = 0
  62         while n = partial(src, w, PIPE_CAPA, src_offset, select_args)
  63           rv += full(r, dst, n, nil, select_args)
  64         end
  65       end
  66     end
  67
  68     rv
  69     ensure
  70       close.each { |io| io.close }
  71   end
  72
  73   # splice the full amount specified from +src+ to +dst+
  74   # Either +dst+ or +src+ must be a pipe.  +dst+ and +src+
  75   # may BOTH be pipes in Linux 2.6.31 or later.
  76   # This will block and wait for IO completion of +len+
  77   # Raises +EOFError+ if end of file is reached.
  78   # bytes.  Returns the number of bytes actually spliced (always +len+)
  79   # The +_select_args+ parameter is reserved for internal use and
  80   # may be removed in future versions.  Do not write code that
  81   # depends on +_select_args+.
  82   def self.full(src, dst, len, src_offset, _select_args = selectable(src, dst))
  83     nr = len
  84     while nr > 0
  85       n = partial(src, dst, nr, src_offset, _select_args) or
  86                                      raise EOFError, "end of file reached"
  87       nr -= n
  88     end
  89     len
  90   end
  91
  92   # splice up to +len+ bytes from +src+ to +dst+.
  93   # Either +dst+ or +src+ must be a pipe.  +dst+ and +src+
  94   # may BOTH be pipes in Linux 2.6.31 or later.
  95   # Returns the number of bytes actually spliced.
  96   # Like IO#readpartial, this never returns Errno::EAGAIN
  97   # The +_select_args+ parameter is reserved for internal use and
  98   # may be removed in future versions.  Do not write code that
  99   # depends on +_select_args+.
 100   def self.partial(src, dst, len, src_offset,
 101                    _select_args = selectable(src, dst))
 102     begin
 103       rv = IO.trysplice(src, src_offset, dst, nil, len, F_MOVE)
 104     end while rv == :EAGAIN and IO.select(*_select_args)
 105     rv
 106   end
 107
 108   # returns an array suitable for splat-ing to IO.select for blocking I/O
 109   def self.selectable(src, dst) # :nodoc:
 110     rv = []
 111     src.stat.pipe? or rv[0] = [ src ]
 112     dst.stat.pipe? or rv[1] = [ dst ]
 113     rv
 114   end
 115 end