DnsNameResolver: Fail query if id space is exhausted (#13784)

normanmaurer · web-flow · commit 0e7c27c6536d · 2024-01-17T20:19:22.000+01:00
Motivation:

When we try to execute a query we will try to select / generate an id
that is used. When we are not able to find one we throw an
IllegalStateException. In this case we also need to ensure we fail the
original promise as otherwise the user might never be notified of the
problem.

Modifications:

- Move throw code out of the DnsQueryContextManager to make it easier to
reason about
- Fail query with IllegalStateException

Result:

Query will be correctly failed in the case of id space exhausting
diff --git a/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContext.java b/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContext.java
@@ -81,7 +81,7 @@ abstract class DnsQueryContext {
 
     private volatile Future<?> timeoutFuture;
 
-    private int id = -1;
+    private int id = Integer.MIN_VALUE;
 
     DnsQueryContext(Channel channel,
                     Future<? extends Channel> channelReadyFuture,
@@ -172,8 +172,15 @@ final DnsQuestion question() {
      * @return                      the {@link ChannelFuture} that is notified once once the write completes.
      */
     final ChannelFuture writeQuery(boolean flush) {
-        assert id == -1 : this.getClass().getSimpleName() + ".writeQuery(...) can only be executed once.";
-        id = queryContextManager.add(nameServerAddr, this);
+        assert id == Integer.MIN_VALUE : this.getClass().getSimpleName() +
+                ".writeQuery(...) can only be executed once.";
+
+        if ((id = queryContextManager.add(nameServerAddr, this)) == -1) {
+            // We did exhaust the id space, fail the query
+            IllegalStateException e = new IllegalStateException("query ID space exhausted: " + question());
+            finishFailure("failed to send a query via " + protocol(), e, false);
+            return channel.newFailedFuture(e);
+        }
 
         // Ensure we remove the id from the QueryContextManager once the query completes.
         promise.addListener(new FutureListener<AddressedEnvelope<DnsResponse, InetSocketAddress>>() {
diff --git a/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContextManager.java b/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContextManager.java
@@ -38,11 +38,27 @@ final class DnsQueryContextManager {
     private final Map<InetSocketAddress, DnsQueryContextMap> map =
             new HashMap<InetSocketAddress, DnsQueryContextMap>();
 
+    /**
+     * Add {@link DnsQueryContext} to the context manager and return the ID that should be used for the query.
+     * This method will return {@code -1} if an ID could not be generated and the context was not stored.
+     *
+     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver to query.
+     * @param qCtx              The {@link {@link DnsQueryContext} to store.
+     * @return                  the ID that should be used or {@code -1} if none could be generated.
+     */
     int add(InetSocketAddress nameServerAddr, DnsQueryContext qCtx) {
         final DnsQueryContextMap contexts = getOrCreateContextMap(nameServerAddr);
         return contexts.add(qCtx);
     }
 
+    /**
+     * Return the {@link DnsQueryContext} for the given {@link InetSocketAddress} and id or {@code null} if
+     * none could be found.
+     *
+     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver.
+     * @param id                The id that identifies the {@link DnsQueryContext} and was used for the query.
+     * @return                  The context or {@code null} if none could be found.
+     */
     DnsQueryContext get(InetSocketAddress nameServerAddr, int id) {
         final DnsQueryContextMap contexts = getContextMap(nameServerAddr);
         if (contexts == null) {
@@ -51,6 +67,14 @@ DnsQueryContext get(InetSocketAddress nameServerAddr, int id) {
         return contexts.get(id);
     }
 
+    /**
+     * Remove the {@link DnsQueryContext} for the given {@link InetSocketAddress} and id or {@code null} if
+     * none could be found.
+     *
+     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver.
+     * @param id                The id that identifies the {@link DnsQueryContext} and was used for the query.
+     * @return                  The context or {@code null} if none could be removed.
+     */
     DnsQueryContext remove(InetSocketAddress nameServerAddr, int id) {
         final DnsQueryContextMap contexts = getContextMap(nameServerAddr);
         if (contexts == null) {
@@ -150,8 +174,7 @@ synchronized int add(DnsQueryContext ctx) {
 
                 id = id + 1 & 0xFFFF;
                 if (++tries >= MAX_TRIES) {
-                    throw new IllegalStateException(
-                            "query ID space exhausted after " + MAX_TRIES + ": " + ctx.question());
+                    return -1;
                 }
             }
         }
