module Redis::Connection::SocketMixin

def _read_from_socket(nbytes, buffer = nil)
  loop do
    case chunk = read_nonblock(nbytes, buffer, exception: false)
    when :wait_readable
      unless wait_readable(@timeout)
        raise Redis::TimeoutError
      end
    when :wait_writable
      unless wait_writable(@timeout)
        raise Redis::TimeoutError
      end
    when nil
      raise Errno::ECONNRESET
    when String
      return chunk
    end
  end
end

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

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

def read(nbytes)
  result = @buffer.slice!(0, nbytes)
  buffer = String.new(capacity: nbytes, encoding: Encoding::ASCII_8BIT)
  result << _read_from_socket(nbytes - result.bytesize, buffer) while result.bytesize < nbytes
  result
end

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

def write(buffer)
  return super(buffer) unless @write_timeout
  bytes_to_write = buffer.bytesize
  total_bytes_written = 0
  loop do
    case bytes_written = write_nonblock(buffer, exception: false)
    when :wait_readable
      unless wait_readable(@write_timeout)
        raise Redis::TimeoutError
      end
    when :wait_writable
      unless wait_writable(@write_timeout)
        raise Redis::TimeoutError
      end
    when nil
      raise Errno::ECONNRESET
    when Integer
      total_bytes_written += bytes_written
      if total_bytes_written >= bytes_to_write
        return total_bytes_written
      end
      buffer = buffer.byteslice(bytes_written..-1)
    end
  end
end

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