From 25f5635c690deeb96cc39c0b878cd8fa14575118 Mon Sep 17 00:00:00 2001 From: xza <45376739+xza-m@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:47:29 +0800 Subject: [PATCH] HIVE-29876: Close response streams in YarnQueueHelper Close the error response stream and fallback input stream on success and failure. Add regression coverage for stream ownership, response formatting, and exception propagation. --- .../hive/ql/exec/tez/YarnQueueHelper.java | 10 +- .../hive/ql/exec/tez/TestYarnQueueHelper.java | 165 ++++++++++++++++++ 2 files changed, 170 insertions(+), 5 deletions(-) create mode 100644 ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestYarnQueueHelper.java diff --git a/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/YarnQueueHelper.java b/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/YarnQueueHelper.java index 67ad8a72d1a5..7b76901d51da 100644 --- a/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/YarnQueueHelper.java +++ b/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/YarnQueueHelper.java @@ -202,12 +202,12 @@ public String handleUnexpectedStatusCode( // TODO: handle 401 and return a new connection? nothing for now InputStream errorStream = connection.getErrorStream(); String error = "Received " + statusCode + (errorStr == null ? "" : (" (" + errorStr + ")")); - if (errorStream != null) { - error += ": " + IOUtils.toString(errorStream); - } else { + if (errorStream == null) { errorStream = connection.getInputStream(); - if (errorStream != null) { - error += ": " + IOUtils.toString(errorStream); + } + try (InputStream stream = errorStream) { + if (stream != null) { + error += ": " + IOUtils.toString(stream); } } return error; diff --git a/ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestYarnQueueHelper.java b/ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestYarnQueueHelper.java new file mode 100644 index 000000000000..2c8846b79ee3 --- /dev/null +++ b/ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestYarnQueueHelper.java @@ -0,0 +1,165 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.hadoop.hive.ql.exec.tez; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.net.HttpURLConnection; +import java.nio.charset.StandardCharsets; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.conf.YarnConfiguration; +import org.junit.Before; +import org.junit.Test; + +public class TestYarnQueueHelper { + private YarnQueueHelper helper; + private HttpURLConnection connection; + + @Before + public void setUp() { + HiveConf conf = new HiveConf(); + conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, false); + conf.set(YarnConfiguration.YARN_HTTP_POLICY_KEY, "HTTP_ONLY"); + conf.set(YarnConfiguration.RM_WEBAPP_ADDRESS, "localhost:8088"); + helper = new YarnQueueHelper(conf); + connection = mock(HttpURLConnection.class); + } + + @Test + public void testClosesErrorStream() throws IOException { + InputStream stream = responseStream("Access denied"); + when(connection.getErrorStream()).thenReturn(stream); + + assertEquals("Received 403 (check queue administrators): Access denied", + helper.handleUnexpectedStatusCode(connection, 403, "check queue administrators")); + + verify(stream).close(); + verify(connection, never()).getInputStream(); + } + + @Test + public void testClosesFallbackInputStream() throws IOException { + InputStream stream = responseStream("Unavailable"); + when(connection.getInputStream()).thenReturn(stream); + + assertEquals("Received 503: Unavailable", helper.handleUnexpectedStatusCode(connection, 503, null)); + + verify(stream).close(); + } + + @Test + public void testClosesEmptyErrorStreamWithoutFallingBack() throws IOException { + InputStream stream = responseStream(""); + when(connection.getErrorStream()).thenReturn(stream); + + assertEquals("Received 500: ", helper.handleUnexpectedStatusCode(connection, 500, null)); + + verify(stream).close(); + verify(connection, never()).getInputStream(); + } + + @Test + public void testMissingResponseStreams() throws IOException { + assertEquals("Received 200 (No input on successful API call)", + helper.handleUnexpectedStatusCode(connection, 200, "No input on successful API call")); + } + + @Test + public void testClosesErrorStreamOnReadFailure() throws IOException { + assertClosesStreamOnReadFailure(true); + verify(connection, never()).getInputStream(); + } + + @Test + public void testClosesFallbackInputStreamOnReadFailure() throws IOException { + assertClosesStreamOnReadFailure(false); + } + + @Test + public void testPropagatesFallbackInputStreamFailure() throws IOException { + IOException failure = new IOException("Cannot get input stream"); + when(connection.getInputStream()).thenThrow(failure); + + assertSame(failure, assertThrows(IOException.class, + () -> helper.handleUnexpectedStatusCode(connection, 500, null))); + } + + @Test + public void testPropagatesCloseFailure() throws IOException { + InputStream stream = responseStream("Unavailable"); + IOException failure = new IOException("Cannot close response"); + doThrow(failure).when(stream).close(); + when(connection.getErrorStream()).thenReturn(stream); + + assertSame(failure, assertThrows(IOException.class, + () -> helper.handleUnexpectedStatusCode(connection, 503, null))); + verify(stream).close(); + } + + @Test + public void testPreservesReadFailureWhenCloseAlsoFails() throws IOException { + InputStream stream = mock(InputStream.class); + IOException readFailure = new IOException("Cannot read response"); + IOException closeFailure = new IOException("Cannot close response"); + when(stream.read(any(byte[].class), anyInt(), anyInt())).thenThrow(readFailure); + doThrow(closeFailure).when(stream).close(); + when(connection.getErrorStream()).thenReturn(stream); + + IOException actual = assertThrows(IOException.class, + () -> helper.handleUnexpectedStatusCode(connection, 500, null)); + + assertSame(readFailure, actual); + assertEquals(1, actual.getSuppressed().length); + assertSame(closeFailure, actual.getSuppressed()[0]); + verify(stream).close(); + } + + private void assertClosesStreamOnReadFailure(boolean useErrorStream) throws IOException { + InputStream stream = mock(InputStream.class); + IOException failure = new IOException("Cannot read response"); + when(stream.read(any(byte[].class), anyInt(), anyInt())).thenThrow(failure); + if (useErrorStream) { + when(connection.getErrorStream()).thenReturn(stream); + } else { + when(connection.getInputStream()).thenReturn(stream); + } + + assertSame(failure, assertThrows(IOException.class, + () -> helper.handleUnexpectedStatusCode(connection, 500, null))); + verify(stream).close(); + } + + private InputStream responseStream(String body) { + return spy(new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8))); + } +}