lib/mogilefs/util.rb

   1 require 'mogilefs'
   2 require 'socket'
   3
   4 module MogileFS::Util
   5
   6   CHUNK_SIZE = 65536
   7
   8   # for copying large files while avoiding GC thrashing as much as possible
   9   # writes the contents of io_rd into io_wr, running through filter if
  10   # it is a Proc object.  The filter proc must respond to a string
  11   # argument (and return a string) and to nil (possibly returning a
  12   # string or nil).  This can be used to filter I/O through an
  13   # Zlib::Inflate or Digest::MD5 object
  14   def sysrwloop(io_rd, io_wr, filter = nil)
  15     copied = 0
  16     # avoid making sysread repeatedly allocate a new String
  17     # This is not well-documented, but both read/sysread can take
  18     # an optional second argument to use as the buffer to avoid
  19     # GC overhead of creating new strings in a loop
  20     buf = ' ' * CHUNK_SIZE # preallocate to avoid GC thrashing
  21     io_wr.sync = true
  22     loop do
  23       begin
  24         b = begin
  25           io_rd.sysread(CHUNK_SIZE, buf)
  26         rescue Errno::EAGAIN, Errno::EINTR
  27           select([io_rd], nil, nil, nil)
  28           retry
  29         end
  30         b = filter.call(b) if filter
  31         copied += syswrite_full(io_wr, b)
  32       rescue EOFError
  33         break
  34       end
  35     end
  36
  37     # filter must take nil as a possible argument to indicate EOF
  38     if filter
  39       b = filter.call(nil)
  40       copied += syswrite_full(io_wr, b) if b && b.length > 0
  41     end
  42     copied
  43   end # sysrwloop
  44
  45   # given an array of URIs, verify that at least one of them is accessible
  46   # with the expected HTTP code within the timeout period (in seconds).
  47   def verify_uris(uris = [], expect = '200', timeout = 2.00)
  48     uri_socks = {}
  49     ok_uris = []
  50     sockets = []
  51
  52     # first, we asynchronously connect to all of them
  53     uris.each do |uri|
  54       sock = Socket.mogilefs_new_nonblock(uri.host, uri.port) rescue next
  55       uri_socks[sock] = uri
  56     end
  57
  58     # wait for at least one of them to finish connecting and send
  59     # HTTP requests to the connected ones
  60     begin
  61       t0 = Time.now
  62       r = select(nil, uri_socks.keys, nil, timeout > 0 ? timeout : 0)
  63       timeout -= (Time.now - t0)
  64       break unless r && r[1]
  65       r[1].each do |sock|
  66         begin
  67           sock.syswrite "HEAD #{uri_socks[sock].request_uri} HTTP/1.0\r\n\r\n"
  68           sockets << sock
  69         rescue
  70           sock.close rescue nil
  71         end
  72       end
  73     end until sockets[0] || timeout < 0
  74
  75     # Await a response from the sockets we had written to, we only need one
  76     # valid response, but we'll take more if they return simultaneously
  77     if sockets[0]
  78       begin
  79         t0 = Time.now
  80         r = select(sockets, nil, nil, timeout > 0 ? timeout : 0)
  81         timeout -= (Time.now - t0)
  82         break unless r && r[0]
  83         r[0].each do |sock|
  84           buf = sock.recv_nonblock(128, Socket::MSG_PEEK) rescue next
  85           if buf && /\AHTTP\/[\d\.]+ #{expect} / =~ buf
  86             ok_uris << uri_socks.delete(sock)
  87             sock.close rescue nil
  88           end
  89         end
  90       end
  91     end until ok_uris[0] || timeout < 0
  92
  93     ok_uris
  94     ensure
  95       uri_socks.keys.each { |sock| sock.close rescue nil }
  96   end
  97
  98   private
  99
 100     # writes the contents of buf to io_wr in full w/o blocking
 101     def syswrite_full(io_wr, buf)
 102       written = 0
 103       loop do
 104         w = begin
 105           io_wr.syswrite(buf)
 106         rescue Errno::EAGAIN, Errno::EINTR
 107           select(nil, [io_wr], nil, nil)
 108           retry
 109         end
 110         written += w
 111         break if w == buf.size
 112         buf = buf[w..-1]
 113       end
 114
 115       written
 116     end
 117
 118 end
 119
 120 require 'timeout'
 121 ##
 122 # Timeout error class.  Subclassing it from Timeout::Error is the only
 123 # reason we require the 'timeout' module, otherwise that module is
 124 # broken and worthless to us.
 125 class MogileFS::Timeout < Timeout::Error; end
 126
 127 class Socket
 128   attr_accessor :mogilefs_addr, :mogilefs_connected
 129
 130   # Socket lacks peeraddr method of the IPSocket/TCPSocket classes
 131   def mogilefs_peername
 132     Socket.unpack_sockaddr_in(getpeername).reverse.map {|x| x.to_s }.join(':')
 133   end
 134
 135   def mogilefs_init(host = nil, port = nil)
 136     return true if defined?(@mogilefs_connected)
 137
 138     @mogilefs_addr = Socket.sockaddr_in(port, host).freeze if port && host
 139
 140     begin
 141       connect_nonblock(@mogilefs_addr)
 142       @mogilefs_connected = true
 143     rescue Errno::EINPROGRESS
 144       nil
 145     rescue Errno::EISCONN
 146       @mogilefs_connected = true
 147     end
 148   end
 149
 150   class << self
 151
 152     # Creates a new (TCP) Socket and initiates (but does not wait for) the
 153     # connection
 154     def mogilefs_new_nonblock(host, port)
 155       sock = Socket.new(Socket::AF_INET, Socket::SOCK_STREAM, 0)
 156       sock.sync = true
 157       sock.mogilefs_init(host, port)
 158       sock
 159     end
 160
 161     # Like TCPSocket.new(host, port), but with an explicit timeout
 162     # (and we don't care for local address/port we're binding to).
 163     # This raises MogileFS::Timeout if timeout expires
 164     def mogilefs_new(host, port, timeout = 5.0)
 165       sock = mogilefs_new_nonblock(host, port) or return sock
 166
 167       while timeout > 0
 168         t0 = Time.now
 169         r = IO.select(nil, [sock], nil, timeout)
 170         return sock if r && r[1] && sock.mogilefs_init
 171         timeout -= (Time.now - t0)
 172       end
 173
 174       sock.close rescue nil
 175       raise MogileFS::Timeout, 'socket write timeout'
 176     end
 177
 178     # Makes a request on a new TCP Socket and returns with a readble socket
 179     # within the given timeout.
 180     # This raises MogileFS::Timeout if timeout expires
 181     def mogilefs_new_request(host, port, request, timeout = 5.0)
 182       t0 = Time.now
 183       sock = mogilefs_new(host, port, timeout)
 184       sock.syswrite(request)
 185       timeout -= (Time.now - t0)
 186       raise MogileFS::Timeout, 'socket read timeout' if timeout < 0
 187       r = IO.select([sock], nil, nil, timeout)
 188       return sock if r && r[0]
 189       raise MogileFS::Timeout, 'socket read timeout'
 190     end
 191
 192   end
 193
 194 end
 195