diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java index a3293620ab9e4b..ee0c017b972c30 100644 --- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java +++ b/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. + *

+ * 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 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( @@ -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 call(HttpURLConnection conn, Object jsonOutput, int expectedResponse, Class klass) throws IOException { return call(conn, jsonOutput, expectedResponse, klass, authRetry); @@ -582,6 +609,11 @@ private 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 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; diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/fs/CommonConfigurationKeysPublic.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/fs/CommonConfigurationKeysPublic.java index 86cad28d3ee33d..750d5d4ddcabf7 100644 --- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/fs/CommonConfigurationKeysPublic.java +++ b/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 + * + * core-default.xml + */ + 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 * diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/IOUtils.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/IOUtils.java index 7d792f8dc7d4e7..93290e4d888f19 100644 --- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/IOUtils.java +++ b/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. diff --git a/hadoop-common-project/hadoop-common/src/main/resources/core-default.xml b/hadoop-common-project/hadoop-common/src/main/resources/core-default.xml index 6eed9a13fa222c..a2625ae3b950e7 100644 --- a/hadoop-common-project/hadoop-common/src/main/resources/core-default.xml +++ b/hadoop-common-project/hadoop-common/src/main/resources/core-default.xml @@ -3965,6 +3965,21 @@ are idle beyond the keepalive window. to KMS servers. + + hadoop.security.kms.client.drain.response.unread.bytes + true + + 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. + + hadoop.security.kms.client.failover.sleep.base.millis diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestKMSClientProvider.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestKMSClientProvider.java index 21ffa2f47fe8b1..633ea25ab8f630 100644 --- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestKMSClientProvider.java +++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestKMSClientProvider.java @@ -32,6 +32,7 @@ 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; @@ -39,7 +40,11 @@ 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)); + } } \ No newline at end of file diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/TestIOUtils.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/TestIOUtils.java index 0fa638a51ebf23..561fb1276b6cc4 100644 --- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/TestIOUtils.java +++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/TestIOUtils.java @@ -21,11 +21,13 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; -import static org.mockito.Mockito.atMost; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.atMost; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -39,9 +41,11 @@ import java.io.InputStream; import java.io.OutputStream; import java.io.RandomAccessFile; +import java.net.HttpURLConnection; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.nio.charset.CharacterCodingException; +import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.HashSet; import java.util.List; @@ -354,4 +358,129 @@ public void testWrapException() throws Exception { }); assertEquals("/tmp/abc.txt", returnedEx.getPath().toString()); } + + @Test + public void testDrainAndCloseNullStream() { + IOUtils.drainAndClose((InputStream) null); + } + + @Test + public void testDrainAndCloseConsumesUntilEofAndCloses() throws Exception { + byte[] payload = "json{}trailing-bytes".getBytes(StandardCharsets.UTF_8); + TrackingInputStream in = new TrackingInputStream(payload); + in.read(new byte[6]); + IOUtils.drainAndClose(in); + assertTrue(in.isClosed()); + assertEquals(payload.length, in.getTotalRead()); + } + + @Test + public void testDrainAndCloseSwallowsReadIOException() throws Exception { + InputStream in = mock(InputStream.class); + when(in.read(any(byte[].class))).thenThrow(new IOException("read failed")); + IOUtils.drainAndClose(in); + verify(in).close(); + } + + @Test + public void testDrainAndCloseHttpUrlConnectionPrefersErrorStream() + throws Exception { + HttpURLConnection conn = mock(HttpURLConnection.class); + byte[] body = "error-body".getBytes(StandardCharsets.UTF_8); + when(conn.getErrorStream()).thenReturn(new ByteArrayInputStream(body)); + IOUtils.drainAndClose(conn); + verify(conn, never()).getInputStream(); + } + + @Test + public void testDrainAndCloseHttpUrlConnectionWhenInputStreamThrows() + throws Exception { + HttpURLConnection conn = mock(HttpURLConnection.class); + when(conn.getErrorStream()).thenReturn(null); + when(conn.getInputStream()).thenThrow(new IOException("no body")); + IOUtils.drainAndClose(conn); + } + + @Test + public void testDrainAndCloseRespectsMaxBytes() throws Exception { + InfiniteInputStream in = new InfiniteInputStream(); + IOUtils.drainAndClose(in); + assertTrue(in.isClosed()); + assertTrue(in.getTotalRead() >= 8L * 1024 * 1024); + assertTrue(in.getTotalRead() <= 8L * 1024 * 1024 + 4096); + } + + private static final class TrackingInputStream extends InputStream { + private final byte[] data; + private int pos; + private int totalRead; + private boolean closed; + + TrackingInputStream(byte[] data) { + this.data = data; + } + + @Override + public int read() throws IOException { + byte[] one = new byte[1]; + int n = read(one, 0, 1); + return n < 0 ? -1 : one[0] & 0xff; + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + if (pos >= data.length) { + return -1; + } + int n = Math.min(len, data.length - pos); + System.arraycopy(data, pos, b, off, n); + pos += n; + totalRead += n; + return n; + } + + boolean isClosed() { + return closed; + } + + int getTotalRead() { + return totalRead; + } + + @Override + public void close() throws IOException { + closed = true; + } + } + + private static final class InfiniteInputStream extends InputStream { + private long totalRead; + private boolean closed; + + @Override + public int read() { + totalRead++; + return 0; + } + + @Override + public int read(byte[] b, int off, int len) { + int n = Math.min(len, 4096); + totalRead += n; + return n; + } + + long getTotalRead() { + return totalRead; + } + + boolean isClosed() { + return closed; + } + + @Override + public void close() throws IOException { + closed = true; + } + } }