From 0494ffcb427c759847544d2045f57a86215e6e61 Mon Sep 17 00:00:00 2001 From: Seungjoo Choi Date: Mon, 14 Sep 2026 09:30:28 +0900 Subject: [PATCH] [FLINK-40644][state/forst] Close the remote stream when a CachedDataInputStream is closed FileCacheEntry#open hands the original (remote) stream over to the CachedDataInputStream and keeps no other reference to it, but close() only closed the cached local stream. Every remote stream that had read from S3 kept its leased connection of the S3A connection pool, so each close of a stream on a still-cached file (table cache eviction, compaction file deletion, stream pool overflow) leaked one connection until the pool was exhausted and all reads blocked in AbstractConnPool.getPoolEntryBlocking. Close the original stream as well (keeping the first exception and attaching the second as suppressed), and close a pooled stream of ByteBufferReadableFSDataInputStream when its positioned read fails, since it is neither returned to the pool nor closed otherwise. Co-Authored-By: Claude Fable 5.1 --- .../ByteBufferReadableFSDataInputStream.java | 18 +- .../forst/fs/cache/CachedDataInputStream.java | 24 ++- ...teBufferReadableFSDataInputStreamTest.java | 99 +++++++++++ .../fs/cache/CachedDataInputStreamTest.java | 154 ++++++++++++++++++ 4 files changed, 289 insertions(+), 6 deletions(-) create mode 100644 flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStreamTest.java create mode 100644 flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStreamTest.java diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStream.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStream.java index 79d64bb78bab86..70eb3f7601430f 100644 --- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStream.java +++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStream.java @@ -20,6 +20,7 @@ import org.apache.flink.core.fs.ByteBufferReadable; import org.apache.flink.core.fs.FSDataInputStream; +import org.apache.flink.util.IOUtils; import java.io.IOException; import java.nio.ByteBuffer; @@ -128,11 +129,18 @@ public int readFully(long position, ByteBuffer bb) throws Exception { } int result; - if (fsDataInputStream instanceof ByteBufferReadable) { - result = ((ByteBufferReadable) fsDataInputStream).read(position, bb); - } else { - fsDataInputStream.seek(position); - result = readFullyFromFSDataInputStream(fsDataInputStream, bb); + try { + if (fsDataInputStream instanceof ByteBufferReadable) { + result = ((ByteBufferReadable) fsDataInputStream).read(position, bb); + } else { + fsDataInputStream.seek(position); + result = readFullyFromFSDataInputStream(fsDataInputStream, bb); + } + } catch (Exception ex) { + // The stream is neither returned to the pool nor closed otherwise: close it here so + // that a failed read does not leak the underlying (remote) stream. + IOUtils.closeQuietly(fsDataInputStream); + throw ex; } boolean offered; diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStream.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStream.java index 625ca55fce9a4d..da070128f0e60d 100644 --- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStream.java +++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStream.java @@ -306,7 +306,29 @@ public void close() throws IOException { return; } closed = true; - closeCachedStream(); + // Both the cached stream and the original (remote) stream belong to this wrapper: + // FileCacheEntry#open hands the original stream over and keeps no reference to it, so + // nobody else can close it. Leaving it open leaks whatever the remote file system holds + // for it (for S3A: a leased connection of the HTTP connection pool). Close both, and if + // both fail keep the first exception and attach the second as suppressed. + IOException failure = null; + try { + closeCachedStream(); + } catch (IOException e) { + failure = e; + } + try { + originalStream.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + if (failure != null) { + throw failure; + } } public boolean isClosed() { diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStreamTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStreamTest.java new file mode 100644 index 00000000000000..2f663bf312bda8 --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/ByteBufferReadableFSDataInputStreamTest.java @@ -0,0 +1,99 @@ +/* + * 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.flink.state.forst.fs; + +import org.apache.flink.core.fs.FSDataInputStream; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for the stream pool of {@link ByteBufferReadableFSDataInputStream}. */ +class ByteBufferReadableFSDataInputStreamTest { + + /** + * A stream taken from (or created for) the pool that fails while reading is neither returned to + * the pool nor closed by the caller, so {@code readFully} has to close it itself — otherwise + * the underlying (remote) stream leaks. + */ + @Test + void testPooledStreamIsClosedWhenPositionedReadFails() throws Exception { + List created = new ArrayList<>(); + ByteBufferReadableFSDataInputStream stream = + new ByteBufferReadableFSDataInputStream( + () -> { + // The first stream is the sequential "original" stream of the wrapper; + // every further one is created for the positioned-read pool. + FailingStream s = new FailingStream(!created.isEmpty()); + created.add(s); + return s; + }, + 4, + 1024); + + assertThatThrownBy(() -> stream.readFully(10, ByteBuffer.allocate(8))) + .isInstanceOf(IOException.class) + .hasMessage("read failed"); + + assertThat(created).hasSize(2); + assertThat(created.get(1).closed).as("failed pooled stream must be closed").isTrue(); + assertThat(created.get(0).closed).as("the original stream is untouched").isFalse(); + stream.close(); + } + + private static final class FailingStream extends FSDataInputStream { + private final boolean failOnRead; + private long pos; + volatile boolean closed; + + FailingStream(boolean failOnRead) { + this.failOnRead = failOnRead; + } + + @Override + public void seek(long desired) { + pos = desired; + } + + @Override + public long getPos() { + return pos; + } + + @Override + public int read() throws IOException { + if (failOnRead) { + throw new IOException("read failed"); + } + pos++; + return 0; + } + + @Override + public void close() { + closed = true; + } + } +} diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStreamTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStreamTest.java new file mode 100644 index 00000000000000..9e25ef92af087c --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/fs/cache/CachedDataInputStreamTest.java @@ -0,0 +1,154 @@ +/* + * 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.flink.state.forst.fs.cache; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.core.fs.FSDataInputStream; +import org.apache.flink.core.fs.FileSystem; +import org.apache.flink.core.fs.Path; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for the stream life cycle of {@link CachedDataInputStream}. */ +class CachedDataInputStreamTest { + + @TempDir java.nio.file.Path tempDir; + + private FileBasedCache newCache() { + return new FileBasedCache( + new Configuration(), + new SizeBasedCacheLimitPolicy(1024 * 1024, 64 * 1024 * 1024), + FileSystem.getLocalFileSystem(), + new Path(tempDir.toString(), "cache"), + null); + } + + private FileCacheEntry newEntry(FileBasedCache cache) { + return new FileCacheEntry( + cache, + new Path(tempDir.toString(), "remote/1.sst"), + new Path(tempDir.toString(), "cache/1.sst"), + 16); + } + + /** + * The original (remote) stream is handed over to the cached stream by {@link + * FileCacheEntry#open} and nothing else references it, so closing the cached stream must close + * it — otherwise it leaks whatever the remote file system holds for it (for S3A: a leased + * connection of the HTTP connection pool). + */ + @Test + void testCloseClosesOriginalStream() throws Exception { + try (FileBasedCache cache = newCache()) { + FileCacheEntry entry = newEntry(cache); + TrackingStream original = new TrackingStream(); + + // The entry is not loaded, so the stream is opened on the original stream only. + CachedDataInputStream stream = entry.open(original); + assertThat(original.closed).isFalse(); + + stream.close(); + assertThat(original.closed).as("original stream must be closed").isTrue(); + + // Idempotent. + stream.close(); + assertThat(original.closeCalls).isEqualTo(1); + } + } + + @Test + void testCloseClosesCachedAndOriginalStream() throws Exception { + try (FileBasedCache cache = newCache()) { + FileCacheEntry entry = newEntry(cache); + TrackingStream cached = new TrackingStream(); + TrackingStream original = new TrackingStream(); + + CachedDataInputStream stream = + new CachedDataInputStream(cache, entry, cached, original); + stream.close(); + + assertThat(cached.closed).isTrue(); + assertThat(original.closed).isTrue(); + } + } + + @Test + void testOriginalStreamIsClosedEvenIfCachedStreamFailsToClose() throws Exception { + try (FileBasedCache cache = newCache()) { + FileCacheEntry entry = newEntry(cache); + TrackingStream cached = new TrackingStream(); + cached.failOnClose = new IOException("cached close failed"); + TrackingStream original = new TrackingStream(); + original.failOnClose = new IOException("original close failed"); + + CachedDataInputStream stream = + new CachedDataInputStream(cache, entry, cached, original); + + assertThatThrownBy(stream::close) + .isInstanceOf(IOException.class) + .hasMessage("cached close failed") + .satisfies( + e -> + assertThat(e.getSuppressed()) + .extracting(Throwable::getMessage) + .containsExactly("original close failed")); + assertThat(cached.closeCalls).isEqualTo(1); + assertThat(original.closeCalls).as("original stream must still be closed").isEqualTo(1); + } + } + + /** An in-memory {@link FSDataInputStream} that records whether it has been closed. */ + static final class TrackingStream extends FSDataInputStream { + private long pos; + volatile boolean closed; + int closeCalls; + IOException failOnClose; + + @Override + public void seek(long desired) { + pos = desired; + } + + @Override + public long getPos() { + return pos; + } + + @Override + public int read() { + pos++; + return 0; + } + + @Override + public void close() throws IOException { + closeCalls++; + closed = true; + if (failOnClose != null) { + throw failOnClose; + } + } + } +}