YARN-8365. Set DNS query type according to client request.

            Contributed by Shane Kumpf

(cherry picked from commit 5889a3c0f931ebe1a9e6b8023c0557922b7d9e49)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-registry/src/main/java/org/apache/hadoop/registry/server/dns/RegistryDNS.java

@@ -1106,7 +1106,7 @@ public class RegistryDNS extends AbstractService implements DNSOperations,
     LOG.debug("calling addAnswer");
     byte rcode = addAnswer(response, name, type, dclass, 0, flags);
     if (rcode != Rcode.NOERROR) {
-      rcode = remoteLookup(response, name, 0);
+      rcode = remoteLookup(response, name, type, 0);
       response.getHeader().setRcode(rcode);
     }
     addAdditional(response, flags);
@@ -1124,9 +1124,10 @@ public class RegistryDNS extends AbstractService implements DNSOperations,
   /**
    * Lookup record from upstream DNS servers.
    */
-  private byte remoteLookup(Message response, Name name, int iterations) {
+  private byte remoteLookup(Message response, Name name, int type,
+      int iterations) {
     // Forward lookup to primary DNS servers
-    Record[] answers = getRecords(name, Type.ANY);
+    Record[] answers = getRecords(name, type);
     try {
       for (Record r : answers) {
         if (r.getType() == Type.SOA) {
@@ -1137,7 +1138,7 @@ public class RegistryDNS extends AbstractService implements DNSOperations,
         if (r.getType() == Type.CNAME) {
           Name cname = ((CNAMERecord) r).getAlias();
           if (iterations < 6) {
-            remoteLookup(response, cname, iterations + 1);
+            remoteLookup(response, cname, Type.CNAME, iterations + 1);
           }
         }
       }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-registry/src/test/java/org/apache/hadoop/registry/server/dns/TestRegistryDNS.java

@@ -424,13 +424,8 @@ public class TestRegistryDNS extends Assert {
     assertEquals("Questions do not match", query.getQuestion(),
         response.getQuestion());
     Record[] recs = response.getSectionArray(Section.ANSWER);
-    boolean found = false;
-    for (Record r : recs) {
-      if (r.getType()==Type.A) {
-        found = true;
-      }
-    }
-    assertTrue("No A records in answer", found);
+    assertEquals(1, recs.length);
+    assertEquals(recs[0].getType(), type);
     return recs;
   }