HADOOP-19996: Fix to drain the unread chunk terminator bytes from inp…

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java

@@ -79,6 +79,7 @@
 import org.apache.hadoop.crypto.key.KeyProviderCryptoExtension;
 import org.apache.hadoop.crypto.key.KeyProviderCryptoExtension.CryptoExtension;
 
+import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.util.Preconditions;
@@ -113,6 +114,7 @@ public class KMSClientProvider extends KeyProvider implements CryptoExtension,
   private static final String UTF8 = "UTF-8";
 
   private static final String CONTENT_TYPE = "Content-Type";
+  private static final String TRANSFER_ENCODING = "Transfer-Encoding";
   private static final String APPLICATION_JSON_MIME = "application/json";
 
   private static final String HTTP_GET = "GET";
@@ -130,6 +132,23 @@ public class KMSClientProvider extends KeyProvider implements CryptoExtension,
       + "authentication.retry-count";
   public static final int DEFAULT_AUTH_RETRY = 1;
 
+  /* Whether to fully consume a response body before closing it, so that the
+   * JDK can return the connection to its keep-alive cache. For a chunked
+   * response the terminating chunk is framing that the JSON parser never
+   * reads, and the JDK's non-blocking fallback cannot always recover it, so
+   * the socket is otherwise torn down and the next call pays for a fresh
+   * connection and TLS handshake.
+   * <p/>
+   * On by default: the drain blocks until the remaining bytes arrive (bounded
+   * by the read timeout), so a server that stops writing after the payload
+   * would make a call wait where it previously returned at once. Set to false
+   * to disable draining.
+   */
+  public static final String DRAIN_RESPONSE_UNREAD_BYTES =
+      CommonConfigurationKeysPublic.KMS_CLIENT_DRAIN_RESPONSE_UNREAD_BYTES_KEY;
+  public static final boolean DEFAULT_DRAIN_RESPONSE_UNREAD_BYTES =
+      CommonConfigurationKeysPublic.KMS_CLIENT_DRAIN_RESPONSE_UNREAD_BYTES_DEFAULT;
+
   private final ValueQueue<EncryptedKeyVersion> encKeyVersionQueue;
 
   private KeyProviderDelegationTokenExtension.DelegationTokenExtension
@@ -331,6 +350,7 @@ private KMSClientProvider[] createProviders(Configuration conf,
   private ConnectionConfigurator configurator;
   private DelegationTokenAuthenticatedURL.Token authToken;
   private final int authRetry;
+  private final boolean drainUnreadResponseBody;
 
   @Override
   public String toString() {
@@ -398,6 +418,8 @@ public KMSClientProvider(URI uri, Configuration conf) throws IOException {
             CommonConfigurationKeysPublic.KMS_CLIENT_TIMEOUT_SECONDS,
             CommonConfigurationKeysPublic.KMS_CLIENT_TIMEOUT_DEFAULT);
     authRetry = conf.getInt(AUTH_RETRY, DEFAULT_AUTH_RETRY);
+    drainUnreadResponseBody =
+            conf.getBoolean(DRAIN_RESPONSE_UNREAD_BYTES, DEFAULT_DRAIN_RESPONSE_UNREAD_BYTES);
     configurator = new TimeoutConnConfigurator(timeout, sslFactory);
     encKeyVersionQueue =
         new ValueQueue<KeyProviderCryptoExtension.EncryptedKeyVersion>(
@@ -536,6 +558,11 @@ public HttpURLConnection run() throws Exception {
     return conn;
   }
 
+  @VisibleForTesting
+  static boolean isChunkedHttpResponse(HttpURLConnection conn) {
+    return "chunked".equalsIgnoreCase(conn.getHeaderField(TRANSFER_ENCODING));
+  }
+
   private <T> T call(HttpURLConnection conn, Object jsonOutput,
       int expectedResponse, Class<T> klass) throws IOException {
     return call(conn, jsonOutput, expectedResponse, klass, authRetry);
@@ -582,6 +609,11 @@ private <T> T call(HttpURLConnection conn, Object jsonOutput,
         String contentType = conn.getRequestProperty(CONTENT_TYPE);
         String requestMethod = conn.getRequestMethod();
         URL url = conn.getURL();
+        // Discard the body of the rejected response so that this connection
+        // can be reused for the retry below.
+        if (drainUnreadResponseBody && isChunkedHttpResponse(conn)) {
+          IOUtils.drainAndClose(conn);
+        }
         conn = createConnection(url, requestMethod);
         if (contentType != null && !contentType.isEmpty()) {
           conn.setRequestProperty(CONTENT_TYPE, contentType);
@@ -595,13 +627,24 @@ private <T> T call(HttpURLConnection conn, Object jsonOutput,
         && conn.getContentType().trim().toLowerCase()
             .startsWith(APPLICATION_JSON_MIME)
         && klass != null) {
-      ObjectMapper mapper = new ObjectMapper();
+      ObjectMapper mapper =
+          new ObjectMapper().disable(JsonParser.Feature.AUTO_CLOSE_SOURCE);
       InputStream is = null;
       try {
         is = conn.getInputStream();
         ret = mapper.readValue(is, klass);
       } finally {
-        IOUtils.closeStream(is);
+        // Jackson stops reading at the end of the JSON value, so drain whatever
+        // is left before closing to keep the connection reusable.
+        if (drainUnreadResponseBody && isChunkedHttpResponse(conn)) {
+          IOUtils.drainAndClose(is);
+        } else {
+          IOUtils.closeStream(is);
+        }
+      }
+    } else {
+      if (drainUnreadResponseBody && isChunkedHttpResponse(conn)) {
+        IOUtils.drainAndClose(conn);
       }
     }
     return ret;

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/fs/CommonConfigurationKeysPublic.java

@@ -883,6 +883,17 @@ public class CommonConfigurationKeysPublic {
       "hadoop.security.kms.client.timeout";
   public static final int KMS_CLIENT_TIMEOUT_DEFAULT = 60;
 
+  /**
+   * @see
+   * <a href="{@docRoot}/../hadoop-project-dist/hadoop-common/core-default.xml">
+   * core-default.xml</a>
+   */
+  public static final String KMS_CLIENT_DRAIN_RESPONSE_UNREAD_BYTES_KEY =
+      "hadoop.security.kms.client.drain.response.unread.bytes";
+  /** Default value for KMS_CLIENT_DRAIN_RESPONSE_UNREAD_BYTES_KEY. */
+  public static final boolean KMS_CLIENT_DRAIN_RESPONSE_UNREAD_BYTES_DEFAULT =
+      true;
+
   /**
    * @see
    * <a href="{@docRoot}/../hadoop-project-dist/hadoop-common/core-default.xml">

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/IOUtils.java

@@ -20,6 +20,7 @@
 
 import java.io.*;
 import java.lang.reflect.Constructor;
+import java.net.HttpURLConnection;
 import java.net.Socket;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
@@ -305,6 +306,66 @@ public static void closeSocket(Socket sock) {
       }
     }
   }
+
+  /**
+   * Upper bound on how much of a response body {@link #drainAndClose} will
+   * read in order to keep a connection reusable. A body larger than this is
+   * not worth holding the connection for.
+   */
+  private static final long MAX_DRAIN_BYTES = 8L * 1024 * 1024;
+
+  /**
+   * Fully consumes a stream before closing it, so that the JDK is able to
+   * return the underlying connection to its keep-alive cache.
+   * Reads block until the remaining bytes arrive, bounded by the read timeout
+   * configured on the connection.
+   *
+   * @param is the stream to drain and close; may be null.
+   */
+  public static void drainAndClose(InputStream is) {
+    if (is == null) {
+      return;
+    }
+    try {
+      final byte[] buf = new byte[4096];
+      long drained = 0;
+      while (drained < MAX_DRAIN_BYTES) {
+        final int n = is.read(buf);
+        if (n < 0) {
+          break;
+        }
+        drained += n;
+      }
+    } catch (IOException ex) {
+      LOG.debug("Failed to drain stream, connection will not be reused.", ex);
+    } finally {
+      closeStream(is);
+    }
+  }
+
+  /**
+   * Drains and closes the response body of a connection whose payload the
+   * caller does not need, so that the connection stays reusable. The error
+   * stream is preferred because on a 4xx/5xx response that is where the body
+   * is exposed and {@link HttpURLConnection#getInputStream()} throws.
+   *
+   * @param conn the connection whose response body should be discarded.
+   */
+  public static void drainAndClose(HttpURLConnection conn) {
+    if (conn == null) {
+      return;
+    }
+    InputStream is = conn.getErrorStream();
+    if (is == null) {
+      try {
+        is = conn.getInputStream();
+      } catch (IOException ex) {
+        LOG.debug("Failed to open response body to drain it.", ex);
+        return;
+      }
+    }
+    drainAndClose(is);
+  }
 
   /**
    * The /dev/null of OutputStreams.

hadoop-common-project/hadoop-common/src/main/resources/core-default.xml

@@ -3965,6 +3965,21 @@ are idle beyond the keepalive window.
     to KMS servers.
   </description>
 </property>
+<property>
+  <name>hadoop.security.kms.client.drain.response.unread.bytes</name>
+  <value>true</value>
+  <description>
+    When true, the KMS client fully reads chunked HTTP response bodies before
+    closing them so the JDK can return connections to the keep-alive pool.
+    Jackson stops at the end of the JSON value and does not consume chunk
+    framing; without draining, clients on JDK 17+ often open a new connection
+    and TLS handshake on every KMS call. Draining applies only when the
+    response uses Transfer-Encoding: chunked. Reads block until the remaining
+    bytes arrive or the KMS read timeout (hadoop.security.kms.client.timeout)
+    is hit; set to false if a server stops sending after the payload and calls
+    would otherwise hang until timeout.
+  </description>
+</property>
 
 <property>
   <name>hadoop.security.kms.client.failover.sleep.base.millis</name>

hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestKMSClientProvider.java

@@ -32,14 +32,19 @@
 import org.slf4j.event.Level;
 
 import java.io.IOException;
+import java.net.HttpURLConnection;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
 
 import static org.apache.hadoop.crypto.key.kms.KMSDelegationToken.TOKEN_KIND;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /**
  * Unit test for {@link KMSClientProvider} class.
@@ -132,4 +137,20 @@ private void testURLSelectToken(final Token tok)
     final Token chosen = url.selectDelegationToken(new URL(urlString), creds);
     assertEquals(tok, chosen);
   }
+
+  @Test
+  public void testIsChunkedHttpResponse() {
+    HttpURLConnection conn = mock(HttpURLConnection.class);
+    when(conn.getHeaderField("Transfer-Encoding")).thenReturn("chunked");
+    assertTrue(KMSClientProvider.isChunkedHttpResponse(conn));
+
+    when(conn.getHeaderField("Transfer-Encoding")).thenReturn("CHUNKED");
+    assertTrue(KMSClientProvider.isChunkedHttpResponse(conn));
+
+    when(conn.getHeaderField("Transfer-Encoding")).thenReturn(null);
+    assertFalse(KMSClientProvider.isChunkedHttpResponse(conn));
+
+    when(conn.getHeaderField("Transfer-Encoding")).thenReturn("gzip");
+    assertFalse(KMSClientProvider.isChunkedHttpResponse(conn));
+  }
 }