This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 8f8f3bb128cb479ed72e104d27e12c69ec14486a Author: Dapeng Sun(孙大鹏) <[email protected]> AuthorDate: Fri Sep 11 23:10:10 2026 +0800 [mosaic] Fix a close race in the reader stream pool and add pool and writer sizing tests (#9744) --- .../format/mosaic/MosaicInputFileAdapter.java | 8 +- .../paimon/format/mosaic/MosaicRecordsReader.java | 3 +- .../paimon/format/mosaic/MosaicRecordsWriter.java | 2 +- .../format/mosaic/MosaicInputFileAdapterTest.java | 103 ++++++++++++++++++++- .../format/mosaic/MosaicReaderWriterTest.java | 40 ++++++++ .../format/mosaic/MosaicRecordsWriterTest.java | 100 ++++++++++++++------ 6 files changed, 225 insertions(+), 31 deletions(-) diff --git a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java index 07385fdd36..27e18f75e3 100644 --- a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java +++ b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapter.java @@ -33,7 +33,8 @@ import java.util.ArrayList; import java.util.List; /** - * Adapter that exposes a Paimon {@link SeekableInputStream} as a Mosaic {@link InputFile}. + * Adapter that exposes a Paimon file as a Mosaic {@link InputFile} through a small pool of {@link + * SeekableInputStream}s. * * <p>Each read borrows one of at most {@code maxStreams} input streams, so concurrent reads do not * serialize on a single stream; a read that finds every stream busy waits for one. @@ -119,19 +120,22 @@ public class MosaicInputFileAdapter implements InputFile, Closeable { } } SeekableInputStream opened = null; + boolean closedMeanwhile; try { opened = fileIO.newInputStream(path); } finally { synchronized (this) { openingStreams--; + closedMeanwhile = closed; if (opened != null && !closed) { allStreams.add(opened); } else { + // The slot is free again for another borrower. notifyAll(); } } } - if (closed) { + if (closedMeanwhile) { opened.close(); throw new IOException("Input file " + path + " is closed"); } diff --git a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java index c36285baea..0f32cb6b28 100644 --- a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java +++ b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsReader.java @@ -438,7 +438,7 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { this.future = future; } - /** Waits for the data; the batch is not released here even when the wait fails. */ + /** Waits for the data; a failed wait leaves the batch queued so close() still drains it. */ @Nullable VectorSchemaRoot await() throws IOException { if (future == null) { @@ -487,6 +487,7 @@ public class MosaicRecordsReader implements FileRecordReader<InternalRow> { // The native read still uses the reader handle; it must complete first. interrupted = true; } catch (ExecutionException e) { + // A failed read holds no data to release. return interrupted; } } diff --git a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java index f4a96937ce..df23bf833b 100644 --- a/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java +++ b/paimon-mosaic/src/main/java/org/apache/paimon/format/mosaic/MosaicRecordsWriter.java @@ -125,7 +125,7 @@ public class MosaicRecordsWriter implements BundleFormatWriter { private static void setInitialCapacity(FieldVector vector, int capacity) { if (vector instanceof BaseRepeatedValueVector) { - // The plain overload would size the element vector for 5 elements per row. + // Avoid Arrow's 5x estimate for fixed- or variable-width element vectors. ((BaseRepeatedValueVector) vector).setInitialCapacity(capacity, 1.0); } else { vector.setInitialCapacity(capacity); diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java index 6a0e4ffb64..9b00ea57ee 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicInputFileAdapterTest.java @@ -82,6 +82,101 @@ class MosaicInputFileAdapterTest { adapter.close(); } + @Test + void testFailedExtraStreamOpenReleasesItsSlot() throws Exception { + CountDownLatch readsStarted = new CountDownLatch(1); + CountDownLatch releaseReads = new CountDownLatch(1); + CountDownLatch thirdOpen = new CountDownLatch(1); + AtomicInteger opens = new AtomicInteger(); + LocalFileIO fileIO = + new LocalFileIO() { + @Override + public SeekableInputStream newInputStream(Path path) throws IOException { + // The second open (the first extra stream) fails once. + int open = opens.incrementAndGet(); + if (open == 2) { + throw new IOException("open failed"); + } + if (open == 3) { + thirdOpen.countDown(); + } + return new BlockingStream(readsStarted, releaseReads); + } + }; + MosaicInputFileAdapter adapter = + new MosaicInputFileAdapter(fileIO, new Path("file:/tmp/mosaic-adapter-test"), 2); + AtomicReference<Throwable> holderFailure = new AtomicReference<>(); + Thread holder = new Thread(() -> read(adapter, holderFailure)); + holder.start(); + readsStarted.await(); + + // The failed open must not keep the second slot reserved. + AtomicReference<Throwable> failed = new AtomicReference<>(); + read(adapter, failed); + assertThat(failed.get()).isInstanceOf(IOException.class).hasMessage("open failed"); + CountDownLatch secondRead = new CountDownLatch(1); + AtomicReference<Throwable> retryFailure = new AtomicReference<>(); + Thread retry = + new Thread( + () -> { + read(adapter, retryFailure); + secondRead.countDown(); + }); + retry.start(); + // The retry must open its own stream while the first one is still held. + thirdOpen.await(); + releaseReads.countDown(); + retry.join(); + holder.join(); + assertThat(retryFailure.get()).isNull(); + assertThat(holderFailure.get()).isNull(); + assertThat(opens.get()).isEqualTo(3); + adapter.close(); + } + + @Test + void testCloseClosesEveryStreamAndRejectsLaterReads() throws Exception { + CountDownLatch readsStarted = new CountDownLatch(2); + CountDownLatch releaseReads = new CountDownLatch(1); + List<BlockingStream> streams = new ArrayList<>(); + LocalFileIO fileIO = + new LocalFileIO() { + @Override + public SeekableInputStream newInputStream(Path path) { + BlockingStream stream = new BlockingStream(readsStarted, releaseReads); + synchronized (streams) { + streams.add(stream); + } + return stream; + } + }; + MosaicInputFileAdapter adapter = + new MosaicInputFileAdapter(fileIO, new Path("file:/tmp/mosaic-adapter-test"), 3); + List<Thread> readers = new ArrayList<>(); + AtomicReference<Throwable> failure = new AtomicReference<>(); + for (int i = 0; i < 2; i++) { + Thread thread = new Thread(() -> read(adapter, failure)); + thread.start(); + readers.add(thread); + } + readsStarted.await(); + releaseReads.countDown(); + for (Thread thread : readers) { + thread.join(); + } + assertThat(failure.get()).isNull(); + assertThat(streams).hasSize(2); + + adapter.close(); + assertThat(streams).allMatch(stream -> stream.closeCount == 1); + AtomicReference<Throwable> afterClose = new AtomicReference<>(); + read(adapter, afterClose); + assertThat(afterClose.get()).isInstanceOf(IOException.class); + // Closing again is a no-op. + adapter.close(); + assertThat(streams).allMatch(stream -> stream.closeCount == 1); + } + @Test void testCloseWakesWaitingReader() throws Exception { CountDownLatch readsStarted = new CountDownLatch(1); @@ -110,6 +205,8 @@ class MosaicInputFileAdapterTest { assertThat(second.get()).isInstanceOf(IOException.class); releaseReads.countDown(); holder.join(); + // The read that held the stream completes; the stream is closed by close(). + assertThat(first.get()).isNull(); } private static void read(MosaicInputFileAdapter adapter, AtomicReference<Throwable> failure) { @@ -157,6 +254,10 @@ class MosaicInputFileAdapterTest { } @Override - public void close() {} + public void close() { + closeCount++; + } + + private volatile int closeCount; } } diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java index aceca5c455..14553219ce 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicReaderWriterTest.java @@ -31,6 +31,7 @@ import org.apache.paimon.format.FormatReaderFactory; import org.apache.paimon.format.FormatWriter; import org.apache.paimon.format.FormatWriterFactory; import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.SeekableInputStream; import org.apache.paimon.fs.local.LocalFileIO; import org.apache.paimon.io.DataFileRecordReader; import org.apache.paimon.options.MemorySize; @@ -57,6 +58,7 @@ import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; @@ -484,6 +486,44 @@ class MosaicReaderWriterTest { assertThatCode(reader::close).doesNotThrowAnyException(); } + @Test + void testReaderOpensAtMostDepthPlusOneStreams() throws IOException { + RowType rowType = DataTypes.ROW(DataTypes.INT(), DataTypes.STRING()); + Path path = newPath(); + GenericRow[] rows = new GenericRow[20_000]; + for (int i = 0; i < rows.length; i++) { + rows[i] = GenericRow.of(i, BinaryString.fromString("value_" + i + "_padding")); + } + writeRows(rowType, path, new Options(), MemorySize.ofKibiBytes(32), rows); + + AtomicInteger opened = new AtomicInteger(); + LocalFileIO fileIO = + new LocalFileIO() { + @Override + public SeekableInputStream newInputStream(Path file) throws IOException { + opened.incrementAndGet(); + return super.newInputStream(file); + } + }; + FormatReaderFactory readerFactory = createReaderFactory(rowType, null, 3); + int count = 0; + try (RecordReader<InternalRow> reader = + readerFactory.createReader( + new FormatReaderContext( + fileIO, path, fileIO.getFileSize(path), null, null))) { + RecordReader.RecordIterator<InternalRow> batch; + while ((batch = reader.readBatch()) != null) { + while (batch.next() != null) { + count++; + } + batch.releaseBatch(); + } + } + assertThat(count).isEqualTo(rows.length); + // One stream per row group being opened plus one for the consumer. + assertThat(opened.get()).isBetween(1, 4); + } + private void writeRows( RowType rowType, Path path, Options options, MemorySize blockSize, GenericRow... rows) throws IOException { diff --git a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java index be69fb6207..cc1f78cfdf 100644 --- a/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java +++ b/paimon-mosaic/src/test/java/org/apache/paimon/format/mosaic/MosaicRecordsWriterTest.java @@ -38,7 +38,9 @@ import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Test; import java.io.ByteArrayOutputStream; +import java.util.Arrays; import java.util.Collections; +import java.util.List; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -210,37 +212,74 @@ class MosaicRecordsWriterTest { void testInitialCapacityNeverExceedsArrowDefaultAllocation() throws Exception { // DOUBLE, STRING and ARRAY<DOUBLE> columns cover the fixed-width, variable-width and // repeated vector families, whose default sizing differs. - RowType rowType = mixedRowType(60); - for (int batchSize : new int[] {4, 1024, 3969, 3970, 65536}) { - long baseline = baselineFirstRowAllocation(rowType, batchSize); - MosaicWriter nativeWriter = mock(MosaicWriter.class); - try (RootAllocator allocator = new RootAllocator()) { - MosaicRecordsWriter writer = - createWriter( - rowType, - allocator, - nativeWriter, - batchSize, - MemorySize.VALUE_128_MB); - try { - writer.addElement(firstRow(rowType)); - assertThat(allocator.getAllocatedMemory()) - .as("batch size %d", batchSize) - .isLessThanOrEqualTo(baseline); - // Arrow rounds buffers to powers of two, so only clearly smaller batches - // allocate less. - if (batchSize <= 1024) { + for (RowType rowType : familyRowTypes()) { + for (int batchSize : new int[] {4, 1024, 3969, 3970, 65536}) { + long baseline = baselineFirstRowAllocation(rowType, batchSize); + MosaicWriter nativeWriter = mock(MosaicWriter.class); + try (RootAllocator allocator = new RootAllocator()) { + MosaicRecordsWriter writer = + createWriter( + rowType, + allocator, + nativeWriter, + batchSize, + MemorySize.VALUE_128_MB); + try { + writer.addElement(firstRow(rowType)); assertThat(allocator.getAllocatedMemory()) .as("batch size %d", batchSize) - .isLessThan(baseline); + .isLessThanOrEqualTo(baseline); + // Arrow rounds buffers to powers of two, so only clearly smaller batches + // allocate less. + if (batchSize <= 1024) { + assertThat(allocator.getAllocatedMemory()) + .as("batch size %d", batchSize) + .isLessThan(baseline); + } + } finally { + writer.close(); } - } finally { - writer.close(); } } } } + @Test + void testFillingOneBatchDoesNotReallocate() throws Exception { + // Fixed-width vectors sized for the batch must hold a full batch without growing. + RowType.Builder builder = RowType.builder(); + for (int i = 0; i < 100; i++) { + builder.field("d" + i, DataTypes.DOUBLE()); + } + RowType rowType = builder.build(); + MosaicWriter nativeWriter = mock(MosaicWriter.class); + try (RootAllocator allocator = new RootAllocator()) { + MosaicRecordsWriter writer = + createWriter(rowType, allocator, nativeWriter, 1024, MemorySize.VALUE_128_MB); + writer.addElement(firstRow(rowType)); + long afterFirstRow = allocator.getAllocatedMemory(); + for (int i = 1; i < 1024; i++) { + writer.addElement(firstRow(rowType)); + } + assertThat(allocator.getAllocatedMemory()).isEqualTo(afterFirstRow); + verify(nativeWriter, never()).write(any()); + writer.close(); + } + } + + /** The mixed schema plus one schema per vector family, so a skipped family shows. */ + private static List<RowType> familyRowTypes() { + RowType.Builder doubles = RowType.builder(); + RowType.Builder strings = RowType.builder(); + RowType.Builder arrays = RowType.builder(); + for (int i = 0; i < 60; i++) { + doubles.field("d" + i, DataTypes.DOUBLE()); + strings.field("s" + i, DataTypes.STRING()); + arrays.field("a" + i, DataTypes.ARRAY(DataTypes.DOUBLE())); + } + return Arrays.asList(mixedRowType(60), doubles.build(), strings.build(), arrays.build()); + } + private static RowType mixedRowType(int columnsPerType) { RowType.Builder builder = RowType.builder(); for (int i = 0; i < columnsPerType; i++) { @@ -253,9 +292,18 @@ class MosaicRecordsWriterTest { private static GenericRow firstRow(RowType rowType) { GenericRow row = new GenericRow(rowType.getFieldCount()); - row.setField(0, 1.0d); - row.setField(1, BinaryString.fromString("one")); - row.setField(2, new GenericArray(new Object[] {1.0d})); + for (int i = 0; i < rowType.getFieldCount(); i++) { + switch (rowType.getTypeAt(i).getTypeRoot()) { + case DOUBLE: + row.setField(i, 1.0d); + break; + case VARCHAR: + row.setField(i, BinaryString.fromString("one")); + break; + default: + row.setField(i, new GenericArray(new Object[] {1.0d})); + } + } return row; }
