forked from kylejginavan/youtube_it
-
Notifications
You must be signed in to change notification settings - Fork 2
/
chain_io.rb
71 lines (61 loc) · 2.3 KB
/
chain_io.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
require 'delegate'
#:stopdoc:
# Stream wrapper that reads IOs in succession. Can be fed to Net::HTTP as post body stream. We use it internally to stream file content
# instead of reading whole video files into memory. Strings passed to the constructor will be wrapped in StringIOs. By default it will auto-close
# file handles when they have been read completely to prevent our uploader from leaking file handles
#
# chain = ChainIO.new(File.open(__FILE__), File.open('/etc/passwd'), "abcd")
class YouTubeG::ChainIO
attr_accessor :autoclose
def initialize(*any_ios)
@autoclose = true
@chain = any_ios.flatten.map{|e| e.respond_to?(:read) ? e : StringIO.new(e.to_s) }
end
def read(buffer_size = 1024)
# Read off the first element in the stack
current_io = @chain.shift
return false if !current_io
buf = current_io.read(buffer_size)
if !buf && @chain.empty? # End of streams
release_handle(current_io) if @autoclose
false
elsif !buf # This IO is depleted, but next one is available
release_handle(current_io) if @autoclose
read(buffer_size)
elsif buf.length < buffer_size # This IO is depleted, but we were asked for more
release_handle(current_io) if @autoclose
buf + (read(buffer_size - buf.length) || '') # and recurse
else # just return the buffer
@chain.unshift(current_io) # put the current back
buf
end
end
# Predict the length of all embedded IOs. Will automatically send file size.
def expected_length
@chain.inject(0) do | len, io |
if io.respond_to?(:length)
len + (io.length - io.pos)
elsif io.is_a?(File)
len + File.size(io.path) - io.pos
else
raise "Cannot predict length of #{io.inspect}"
end
end
end
private
def release_handle(io)
io.close if io.respond_to?(:close)
end
end
# Net::HTTP only can send chunks of 1024 bytes. This is very inefficient, so we have a spare IO that will send more when asked for 1024.
# We use delegation because the read call is recursive.
class YouTubeG::GreedyChainIO < DelegateClass(YouTubeG::ChainIO)
BIG_CHUNK = 512 * 1024 # 500 kb
def initialize(*with_ios)
__setobj__(YouTubeG::ChainIO.new(with_ios))
end
def read(any_buffer_size)
__getobj__.read(BIG_CHUNK)
end
end
#:startdoc: