class EventMachine::DNS::Request

def id
  begin
    @id = rand(65535)
    @socket.register_request(@id, self)
  rescue RequestIdAlreadyUsed
    retry
  end unless defined?(@id)
  @id
end

def initialize(socket, hostname)
  @socket = socket
  @hostname = hostname
  @tries = 0
  @last_send = Time.at(0)
  @retry_interval = 3
  @max_tries = 5
  if addrs = Resolver.hosts[hostname]
    succeed addrs
  else
    EM.next_tick { tick }
  end
end

def packet
  msg = Resolv::DNS::Message.new
  msg.id = id
  msg.rd = 1
  msg.add_question @hostname, Resolv::DNS::Resource::IN::A
  msg
end

def receive_answer(msg)
  addrs = []
  msg.each_answer do |name,ttl,data|
    if data.kind_of?(Resolv::DNS::Resource::IN::A) ||
        data.kind_of?(Resolv::DNS::Resource::IN::AAAA)
      addrs << data.address.to_s
    end
  end
  if addrs.empty?
    fail "rcode=#{msg.rcode}"
  else
    succeed addrs
  end
end

def send
  @tries += 1
  @last_send = Time.now
  @socket.send_packet(packet.encode)
end

def tick
  # Break early if nothing to do
  return if @last_send + @retry_interval > Time.now
  if @tries < @max_tries
    send
  else
    @socket.deregister_request(@id, self)
    fail 'retries exceeded'
  end
end