[FLINK-17440][network] Resolve potential buffer leak in output unspil…

…ling for unaligned checkpoints

If there are any exceptions thrown while interacting with ChannelStateReader#readOutputData, we should close the
respective BufferConsumer to avoid leak.

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PipelinedSubpartition.java

@@ -103,17 +103,23 @@ public class PipelinedSubpartition extends ResultSubpartition {
 
 	@Override
 	public void readRecoveredState(ChannelStateReader stateReader) throws IOException, InterruptedException {
+		boolean recycleBuffer = true;
 		for (ReadResult readResult = ReadResult.HAS_MORE_DATA; readResult == ReadResult.HAS_MORE_DATA;) {
 			BufferBuilder bufferBuilder = parent.getBufferPool().requestBufferBuilderBlocking(subpartitionInfo.getSubPartitionIdx());
 			BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
-			readResult = stateReader.readOutputData(subpartitionInfo, bufferBuilder);
-
-			// check whether there are some states data filled in this time
-			if (bufferConsumer.isDataAvailable()) {
-				add(bufferConsumer, false, false);
-				bufferBuilder.finish();
-			} else {
-				bufferConsumer.close();
+			try {
+				readResult = stateReader.readOutputData(subpartitionInfo, bufferBuilder);
+
+				// check whether there are some states data filled in this time
+				if (bufferConsumer.isDataAvailable()) {
+					add(bufferConsumer, false, false);
+					recycleBuffer = false;
+					bufferBuilder.finish();
+				}
+			} finally {
+				if (recycleBuffer) {
+					bufferConsumer.close();
+				}
 			}
 		}
 	}

flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionTest.java

@@ -529,6 +529,32 @@ public void testInitializeMoreStateThanBuffer() throws Exception {
 		}
 	}
 
+	/**
+	 * Tests that the buffer is recycled correctly if exception is thrown during
+	 * {@link ChannelStateReader#readOutputData(ResultSubpartitionInfo, BufferBuilder)}.
+	 */
+	@Test
+	public void testReadRecoveredStateWithException() throws Exception {
+		final int totalBuffers = 2;
+		final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 1, 1);
+		final ResultPartition partition = new ResultPartitionBuilder()
+			.setNetworkBufferPool(globalPool)
+			.build();
+		final ChannelStateReader stateReader = new ChannelStateReaderWithException();
+
+		try {
+			partition.setup();
+			partition.readRecoveredState(stateReader);
+		} catch (IOException e) {
+			assertThat("should throw custom exception message", e.getMessage().contains("test"));
+		} finally {
+			globalPool.destroyAllBufferPools();
+			// verify whether there are any buffers leak
+			assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
+			globalPool.destroy();
+		}
+	}
+
 	/**
 	 * The {@link ChannelStateReader} instance for restoring the specific number of states.
 	 */
@@ -562,4 +588,25 @@ public ReadResult readOutputData(ResultSubpartitionInfo info, BufferBuilder buff
 		public void close() {
 		}
 	}
+
+	/**
+	 * The {@link ChannelStateReader} instance for throwing exception when
+	 * {@link #readOutputData(ResultSubpartitionInfo, BufferBuilder)}.
+	 */
+	private static final class ChannelStateReaderWithException implements ChannelStateReader {
+
+		@Override
+		public ReadResult readInputData(InputChannelInfo info, Buffer buffer) {
+			return ReadResult.NO_MORE_DATA;
+		}
+
+		@Override
+		public ReadResult readOutputData(ResultSubpartitionInfo info, BufferBuilder bufferBuilder) throws IOException {
+			throw new IOException("test");
+		}
+
+		@Override
+		public void close() {
+		}
+	}
 }