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 79d64bb78bab8..70eb3f7601430 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 625ca55fce9a4..da070128f0e60 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 0000000000000..2f663bf312bda --- /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 0000000000000..9e25ef92af087 --- /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; + } + } + } +}