--- a/test/jdk/com/sun/jndi/dns/lib/DNSServer.java	Tue Jul 24 08:58:32 2018 -0700
+++ b/test/jdk/com/sun/jndi/dns/lib/DNSServer.java	Wed Jul 25 11:03:07 2018 +0800
@@ -26,6 +26,8 @@
 import java.io.IOException;
 import java.net.DatagramPacket;
 import java.net.DatagramSocket;
+import java.net.InetAddress;
+import java.net.SocketException;
 import java.nio.ByteBuffer;
 import java.nio.file.Paths;
 import java.util.ArrayList;
@@ -49,7 +51,7 @@
  * Typically, DNS protocol exchange is generated by DNSTracer who captures
  * communication messages between DNS application program and real DNS server
  */
-public class DNSServer implements Runnable {
+public class DNSServer extends Thread implements Server {
 
     public class Pair<F, S> {
         private F first;
@@ -87,19 +89,21 @@
     private boolean loop;
     private final List<Pair<byte[], byte[]>> cache = new ArrayList<>();
     private ByteBuffer reqBuffer = ByteBuffer.allocate(DNS_PACKET_SIZE);
+    private volatile boolean isRunning;
 
-    public DNSServer(DatagramSocket socket, String filename) {
-        this(socket, filename, false);
+    public DNSServer(String filename) throws SocketException {
+        this(filename, false);
     }
 
-    public DNSServer(DatagramSocket socket, String filename, boolean loop) {
-        this.socket = socket;
+    public DNSServer(String filename, boolean loop) throws SocketException {
+        this.socket = new DatagramSocket(0, InetAddress.getLoopbackAddress());
         this.filename = filename;
         this.loop = loop;
     }
 
     public void run() {
         try {
+            isRunning = true;
             System.out.println(
                     "DNSServer: Loading DNS cache data from : " + filename);
             loadCaptureFile(filename);
@@ -112,30 +116,25 @@
             int playbackIndex = 0;
 
             while (playbackIndex < cache.size()) {
-                DatagramPacket reqPacket = new DatagramPacket(reqBuffer.array(),
-                        reqBuffer.array().length);
-                socket.receive(reqPacket);
-
-                System.out.println(
-                        "DNSServer: received query message from " + reqPacket
-                                .getSocketAddress());
+                DatagramPacket reqPacket = receiveQuery();
 
                 if (!verifyRequestMsg(reqPacket, playbackIndex)) {
-                    throw new RuntimeException(
-                            "DNSServer: Error: Failed to verify DNS request. "
-                                    + "Not identical request message : \n"
-                                    + encoder.encodeBuffer(
-                                    Arrays.copyOf(reqPacket.getData(),
-                                            reqPacket.getLength())));
+                    if (playbackIndex > 0 && verifyRequestMsg(reqPacket,
+                            playbackIndex - 1)) {
+                        System.out.println(
+                                "DNSServer: received retry query, resend");
+                        playbackIndex--;
+                    } else {
+                        throw new RuntimeException(
+                                "DNSServer: Error: Failed to verify DNS request. "
+                                        + "Not identical request message : \n"
+                                        + encoder.encodeBuffer(
+                                        Arrays.copyOf(reqPacket.getData(),
+                                                reqPacket.getLength())));
+                    }
                 }
 
-                byte[] payload = generateResponsePayload(reqPacket,
-                        playbackIndex);
-                socket.send(new DatagramPacket(payload, payload.length,
-                        reqPacket.getSocketAddress()));
-                System.out.println(
-                        "DNSServer: send response message to " + reqPacket
-                                .getSocketAddress());
+                sendResponse(reqPacket, playbackIndex);
 
                 playbackIndex++;
                 if (loop && playbackIndex >= cache.size()) {
@@ -145,11 +144,54 @@
 
             System.out.println(
                     "DNSServer: Done for all cached messages playback");
+
+            System.out.println(
+                    "DNSServer: Still listening for possible retry query");
+            while (true) {
+                DatagramPacket reqPacket = receiveQuery();
+
+                // here we only handle the retry query for last one
+                if (!verifyRequestMsg(reqPacket, playbackIndex - 1)) {
+                    throw new RuntimeException(
+                            "DNSServer: Error: Failed to verify DNS request. "
+                                    + "Not identical request message : \n"
+                                    + encoder.encodeBuffer(
+                                    Arrays.copyOf(reqPacket.getData(),
+                                            reqPacket.getLength())));
+                }
+
+                sendResponse(reqPacket, playbackIndex - 1);
+            }
         } catch (Exception e) {
-            System.err.println("DNSServer: Error: " + e);
+            if (isRunning) {
+                System.err.println("DNSServer: Error: " + e);
+                e.printStackTrace();
+            } else {
+                System.out.println("DNSServer: Exit");
+            }
         }
     }
 
+    private DatagramPacket receiveQuery() throws IOException {
+        DatagramPacket reqPacket = new DatagramPacket(reqBuffer.array(),
+                reqBuffer.array().length);
+        socket.receive(reqPacket);
+
+        System.out.println("DNSServer: received query message from " + reqPacket
+                .getSocketAddress());
+
+        return reqPacket;
+    }
+
+    private void sendResponse(DatagramPacket reqPacket, int playbackIndex)
+            throws IOException {
+        byte[] payload = generateResponsePayload(reqPacket, playbackIndex);
+        socket.send(new DatagramPacket(payload, payload.length,
+                reqPacket.getSocketAddress()));
+        System.out.println("DNSServer: send response message to " + reqPacket
+                .getSocketAddress());
+    }
+
     /*
      * Load a capture file containing an DNS protocol exchange in the
      * hexadecimal dump format emitted by sun.misc.HexDumpEncoder:
@@ -299,4 +341,23 @@
         }
         return -1;
     }
+
+    @Override public void stopServer() {
+        isRunning = false;
+        if (socket != null) {
+            try {
+                socket.close();
+            } catch (Exception e) {
+                // ignore
+            }
+        }
+    }
+
+    @Override public int getPort() {
+        if (socket != null) {
+            return socket.getLocalPort();
+        } else {
+            return -1;
+        }
+    }
 }