module Redis::Connection::SocketMixin

def _read_from_socket(nbytes)
  begin
    read_nonblock(nbytes)
  rescue *NBIO_READ_EXCEPTIONS
    if IO.select([self], nil, nil, @timeout)
      retry
    else
      raise Redis::TimeoutError
    end
  rescue *NBIO_WRITE_EXCEPTIONS
    if IO.select(nil, [self], nil, @timeout)
      retry
    else
      raise Redis::TimeoutError
    end
  end
rescue EOFError
  raise Errno::ECONNRESET
end

def _write_to_socket(data)
  begin
    write_nonblock(data)
  rescue *NBIO_READ_EXCEPTIONS
    if IO.select([self], nil, nil, @write_timeout)
      retry
    else
      raise Redis::TimeoutError
    end
  rescue *NBIO_WRITE_EXCEPTIONS
    if IO.select(nil, [self], nil, @write_timeout)
      retry
    else
      raise Redis::TimeoutError
    end
  end
rescue EOFError
  raise Errno::ECONNRESET
end

def gets
  crlf = nil
  while (crlf = @buffer.index(CRLF)) == nil
    @buffer << _read_from_socket(1024)
  end
  @buffer.slice!(0, crlf + CRLF.bytesize)
end

def initialize(*args)
  super(*args)
  @timeout = @write_timeout = nil
  @buffer = ""
end

def read(nbytes)
  result = @buffer.slice!(0, nbytes)
  while result.bytesize < nbytes
    result << _read_from_socket(nbytes - result.bytesize)
  end
  result
end

def timeout=(timeout)
  if timeout && timeout > 0
    @timeout = timeout
  else
    @timeout = nil
  end
end

def write(data)
  return super(data) unless @write_timeout
  length = data.bytesize
  total_count = 0
  loop do
    count = _write_to_socket(data)
    total_count += count
    return total_count if total_count >= length
    data = data.byteslice(count..-1)
  end
end

def write_timeout=(timeout)
  if timeout && timeout > 0
    @write_timeout = timeout
  else
    @write_timeout = nil
  end
end