lib/io/splice.rb

   1 # -*- encoding: binary -*-
   2 require 'io_splice_ext'
   3
   4 class IO
   5
   6   module Splice
   7
   8     # the version of IO::Splice, currently 0.1.0
   9     VERSION = '1.0.0'
  10
  11     # The maximum capacity of the pipe in bytes.
  12     # Under stock Linux, this is 65536 bytes as of 2.6.11, and 4096 before
  13     # We detect this at runtime as it is easy to recompile the kernel
  14     # and set a new value.
  15     PIPE_CAPA = begin
  16       rd, wr = IO.pipe
  17       buf = ' ' * PIPE_BUF
  18       n = 0
  19       begin
  20         n += wr.write_nonblock(buf)
  21       rescue Errno::EAGAIN
  22         break
  23       end while true
  24       wr.close
  25       rd.close
  26       n
  27     end
  28
  29     # copies the contents of the IO object given by +src+ to +dst+
  30     # If len is specified, then only len bytes are copied.  Otherwise
  31     # the copy will be until EOF is reached on the +src+.
  32     # +src+ and +dst+ must be IO objects or respond to +to_io+
  33     def self.copy_stream(src, dst, len = nil, src_offset = nil)
  34       src.kind_of?(String) and src = File.open(src, 'rb')
  35       dst.kind_of?(String) and dst = File.open(dst, 'wb')
  36       src, dst = src.to_io, dst.to_io
  37       rv = len
  38       src.sysseek(src_offset) if src_offset
  39       if src.stat.pipe? || dst.stat.pipe?
  40         if len
  41           while len > 0
  42             nr = len > PIPE_CAPA ? PIPE_CAPA : len
  43             nr = IO.splice(src, nil, dst, nil, nr, F_MOVE)
  44             if nr == 0
  45               raise EOFError, "unexpected EOF with #{len} bytes left"
  46             end
  47             len -= nr
  48           end
  49         else
  50           rv = 0
  51           begin
  52             nr = IO.splice(src, nil, dst, nil, PIPE_CAPA, F_MOVE)
  53             rv += nr
  54           rescue EOFError
  55             break
  56           end while true
  57         end
  58       else
  59         begin
  60           r, w = IO.pipe
  61           if len
  62             while len > 0
  63               nr = len > PIPE_CAPA ? PIPE_CAPA : len
  64               nr_src = copy_stream(src, w, nr)
  65               nr_src == nr or
  66                 raise RuntimeError, "short splice from: #{nr_src} != #{nr}"
  67               nr_dst = copy_stream(r, dst, nr)
  68               nr_dst == nr or
  69                 raise RuntimeError, "short splice to: #{nr_dst} != #{nr}"
  70               len -= nr
  71             end
  72           else
  73             rv = 0
  74             begin
  75               nr = IO.splice(src, nil, w, nil, PIPE_CAPA, F_MOVE)
  76               rv += nr
  77               begin
  78                 nr -= IO.splice(r, nil, dst, nil, nr, F_MOVE)
  79               end while nr > 0
  80             rescue EOFError
  81               break
  82             end while true
  83           end
  84         ensure
  85           r.close
  86           w.close
  87         end
  88       end
  89
  90       rv
  91     end
  92
  93   end
  94 end