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