lib/mogilefs/socket/kgio.rb

   1 # -*- encoding: binary -*-
   2 require "kgio"
   3
   4 class MogileFS::Socket < Kgio::Socket
   5   def self.start(host, port)
   6     sock = super(Socket.sockaddr_in(port, host))
   7     Socket.const_defined?(:TCP_NODELAY) and
   8       sock.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
   9     sock
  10   end
  11
  12   def self.tcp(host, port, timeout = 5)
  13     sock = start(host, port)
  14     unless sock.kgio_wait_writable(timeout)
  15       sock.close
  16       raise MogileFS::Timeout, 'socket connect timeout'
  17     end
  18     sock
  19   end
  20
  21   def timed_read(len, dst = "", timeout = 5)
  22     case rc = kgio_tryread(len, dst)
  23     when :wait_readable
  24       kgio_wait_readable(timeout) or raise MogileFS::Timeout, "read timeout"
  25     else
  26       return rc
  27     end while true
  28   end
  29
  30   def timed_peek(len, dst, timeout = 5)
  31     case rc = kgio_trypeek(len, dst)
  32     when :wait_readable
  33       kgio_wait_readable(timeout) or raise MogileFS::Timeout, "peek timeout"
  34     else
  35       return rc
  36     end while true
  37   end
  38
  39   def timed_write(buffer, timeout = 5)
  40     case rc = kgio_trywrite(buffer)
  41     when :wait_writable
  42       kgio_wait_writable(timeout) or raise MogileFS::Timeout, "write timeout"
  43     when String
  44       buffer = rc
  45     else
  46       return rc
  47     end while true
  48   end
  49
  50   SEP_RE = /\A(.*?#{Regexp.escape("\n")})/
  51   def timed_gets(timeout = 5)
  52     unless defined?(@rbuf)
  53       @rbuf = timed_read(1024, "", timeout) or return # EOF
  54     end
  55     begin
  56       @rbuf.sub!(SEP_RE, "") and return $1
  57       tmp ||= ""
  58       if timed_read(1024, tmp, timeout)
  59         @rbuf << tmp
  60       else
  61         # EOF, return the last buffered bit even without SEP_RE matching
  62         # (not ideal for MogileFS, this is an error)
  63         return @rbuf.empty? ? nil : @rbuf.slice!(0, @rbuf.size)
  64       end
  65     end while true
  66   end
  67 end