This is an automated email from the ASF dual-hosted git repository. danny0405 pushed a commit to branch release-1.1.2 in repository https://gitbox.apache.org/repos/asf/hudi.git
commit 4d14a91b0ccd6df2487bd617c970e988a93e9376 Author: fhan <[email protected]> AuthorDate: Tue Aug 25 10:45:15 2026 +0800 fix(flink): prevent reusing diverged buffers on memory exhaustion (#19692) * fix(flink): prevent reusing diverged buffers on memory exhaustion --------- Co-authored-by: fhan <[email protected]> (cherry picked from commit 81fd4992b8dc0687b8ccfa6190d0570682e55ee1) --- .../org/apache/hudi/sink/StreamWriteFunction.java | 203 +++++++++++++---- .../org/apache/hudi/sink/buffer/RowDataBucket.java | 13 ++ .../TestBucketStreamWriteMemoryExhaustion.java | 243 +++++++++++++++++++++ .../apache/hudi/sink/buffer/TestRowDataBucket.java | 104 +++++++++ 4 files changed, 516 insertions(+), 47 deletions(-) diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/StreamWriteFunction.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/StreamWriteFunction.java index 56734988567a..39a5463b597e 100644 --- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/StreamWriteFunction.java +++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/StreamWriteFunction.java @@ -60,6 +60,7 @@ import org.apache.flink.util.Collector; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.Closeable; import java.io.IOException; import java.io.Serializable; import java.util.ArrayList; @@ -70,7 +71,6 @@ import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import java.util.NoSuchElementException; import static org.apache.hudi.common.util.HoodieRecordUtils.getOrderingFieldNames; @@ -179,13 +179,6 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink bufferRecord(record); } - @Override - public void close() { - if (this.writeClient != null) { - this.writeClient.close(); - } - } - /** * End input action for batch source. */ @@ -272,7 +265,7 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink return bucket.writeRow(record.getRowData()); } catch (MemoryPagesExhaustedException e) { - LOG.info("There is no enough free pages in memory pool to create buffer, need flushing first.", e); + LOG.info("There are not enough free pages in the memory pool to create a buffer; flushing is required first."); return false; } } @@ -298,31 +291,18 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink // 1. try buffer the record into the memory pool boolean success = doBufferRecord(bucketID, record); if (!success) { - // 2. flushes the bucket if the memory pool is full - RowDataBucket bucketToFlush = this.buckets.values().stream() - .max(Comparator.comparingLong(RowDataBucket::getBufferSize)) - .orElseThrow(NoSuchElementException::new); - if (flushBucket(bucketToFlush)) { - // 2.1 flushes the data bucket with maximum size - this.tracer.countDown(bucketToFlush.getBufferSize()); - disposeBucket(bucketToFlush); - } else { - LOG.warn("The buffer size hits the threshold {}, but still flush the max size data bucket failed!", this.tracer.maxBufferSize); - } - // 2.2 try to write row again - success = doBufferRecord(bucketID, record); - if (!success) { - throw new RuntimeException("Buffer is too small to hold a single record."); - } + // 2. reclaim pages. A buffer whose write returned false must never be reused because + // BinaryInMemorySortBuffer may already have changed its variable-length storage state. + reclaimMemoryAfterFailedWrite(bucketID); + + // 2.1 retry once with a newly-created buffer + retryBufferRecord(bucketID, record); } RowDataBucket bucket = this.buckets.get(bucketID); this.tracer.trace(bucket.getLastRecordSize()); // 3. flushes the bucket if it is full if (bucket.isFull()) { - if (flushBucket(bucket)) { - this.tracer.countDown(bucket.getBufferSize()); - disposeBucket(bucket); - } + flushAndDisposeBucket(bucket); } // update buffer metrics after tracing buffer size writeMetrics.setWriteBufferedSize(this.tracer.bufferSize); @@ -332,8 +312,103 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink * RowData data bucket can not be used after disposing. */ private void disposeBucket(RowDataBucket rowDataBucket) { - rowDataBucket.dispose(); - this.buckets.remove(rowDataBucket.getBucketId()); + try { + rowDataBucket.dispose(); + } finally { + this.buckets.remove(rowDataBucket.getBucketId()); + } + } + + private void reclaimMemoryAfterFailedWrite(String bucketID) { + // A creation failure leaves no bucket in the map, while a write failure leaves the + // diverged bucket in the map so that its committed records can be flushed and disposed. + RowDataBucket failedBucket = this.buckets.get(bucketID); + RowDataBucket bucketToFlush = this.buckets.values().stream() + .filter(bucket -> !bucketID.equals(bucket.getBucketId()) && !bucket.isEmpty()) + .max(Comparator.comparingLong(RowDataBucket::getBufferSize)) + .orElse(null); + + if (failedBucket == null) { + if (bucketToFlush == null) { + throw new HoodieException( + "Not enough memory pages to create a RowData buffer and no non-empty bucket can be flushed"); + } + flushAndDisposeBucket(bucketToFlush); + return; + } + + ValidationUtils.checkState( + failedBucket.isDiverged(), "The failed RowData bucket has not diverged"); + + RuntimeException failure = null; + if (bucketToFlush != null) { + try { + flushAndDisposeBucket(bucketToFlush); + } catch (RuntimeException e) { + failure = e; + } + } + + try { + flushAndDisposeBucket(failedBucket); + } catch (RuntimeException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + + if (failure != null) { + throw failure; + } + } + + private void retryBufferRecord( + String bucketID, HoodieFlinkInternalRow record) throws IOException { + final boolean success; + try { + success = doBufferRecord(bucketID, record); + } catch (IOException | RuntimeException e) { + disposeFailedRetryBucket(bucketID, e); + throw e; + } + + if (!success) { + HoodieException exception = new HoodieException( + this.buckets.get(bucketID) == null + ? "Not enough memory pages to create a RowData buffer after flushing" + : "The write buffer is too small to hold a single record"); + disposeFailedRetryBucket(bucketID, exception); + throw exception; + } + } + + private void disposeFailedRetryBucket(String bucketID, Throwable failure) { + RowDataBucket bucket = this.buckets.get(bucketID); + if (bucket == null) { + return; + } + try { + disposeBucket(bucket); + } catch (RuntimeException cleanupFailure) { + failure.addSuppressed(cleanupFailure); + } + } + + private void flushAndDisposeBucket(RowDataBucket bucket) { + long bufferSize = bucket.getBufferSize(); + try { + if (!bucket.isEmpty()) { + flushBucket(bucket); + } + } finally { + try { + this.tracer.countDown(bufferSize); + } finally { + disposeBucket(bucket); + } + } } private static BucketInfo getBucketInfo(HoodieFlinkInternalRow internalRow) { @@ -356,14 +431,10 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink && this.buckets.values().stream().anyMatch(bucket -> !bucket.isEmpty()); } - private boolean flushBucket(RowDataBucket bucket) { + private void flushBucket(RowDataBucket bucket) { String instant = instantToWrite(true); - if (instant == null) { - // in case there are empty checkpoints that has no input data - LOG.info("No inflight instant when flushing data, skip."); - return false; - } + ValidationUtils.checkState(instant != null, "No inflight instant when flushing a non-empty bucket"); ValidationUtils.checkState(!bucket.isEmpty(), "Data bucket to flush has no buffering records"); final List<WriteStatus> writeStatus = writeRecords(instant, bucket); @@ -378,7 +449,6 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink this.eventGateway.sendEventToCoordinator(event); writeStatuses.addAll(writeStatus); - return true; } public void flushRemaining(boolean endInput) { @@ -391,15 +461,14 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink final List<WriteStatus> writeStatus; if (!buckets.isEmpty()) { writeStatus = new ArrayList<>(); - this.buckets.values() - // The records are partitioned by the bucket ID and each batch sent to - // the writer belongs to one bucket. - .forEach(bucket -> { - if (!bucket.isEmpty()) { - writeStatus.addAll(writeRecords(currentInstant, bucket)); - bucket.dispose(); - } - }); + // The records are partitioned by the bucket ID and each batch sent to + // the writer belongs to one bucket. + for (RowDataBucket bucket : new ArrayList<>(this.buckets.values())) { + if (!bucket.isEmpty()) { + writeStatus.addAll(writeRecords(currentInstant, bucket)); + } + disposeBucket(bucket); + } } else { LOG.info("No data to write in subtask [{}] for instant [{}]", taskID, currentInstant); writeStatus = Collections.emptyList(); @@ -458,6 +527,46 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink writeMetrics.registerMetrics(); } + @Override + public void close() { + Exception closeFailure = null; + if (this.buckets != null) { + for (RowDataBucket bucket : new ArrayList<>(this.buckets.values())) { + try { + disposeBucket(bucket); + } catch (RuntimeException e) { + closeFailure = addCloseFailure(closeFailure, e); + } + } + } + + try { + if (this.memorySegmentPool instanceof Closeable) { + ((Closeable) this.memorySegmentPool).close(); + } + } catch (Exception e) { + closeFailure = addCloseFailure(closeFailure, e); + } + try { + if (this.writeClient != null) { + this.writeClient.close(); + } + } catch (Exception e) { + closeFailure = addCloseFailure(closeFailure, e); + } + if (closeFailure != null) { + throw new HoodieException("Failed to close stream writer", closeFailure); + } + } + + private static Exception addCloseFailure(Exception failure, Exception nextFailure) { + if (failure == null) { + return nextFailure; + } + failure.addSuppressed(nextFailure); + return failure; + } + // ------------------------------------------------------------------------- // Getter/Setter // ------------------------------------------------------------------------- diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/RowDataBucket.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/RowDataBucket.java index 385359c6745f..3371555d1d88 100644 --- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/RowDataBucket.java +++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/RowDataBucket.java @@ -18,6 +18,7 @@ package org.apache.hudi.sink.buffer; +import org.apache.hudi.common.util.ValidationUtils; import org.apache.hudi.table.action.commit.BucketInfo; import org.apache.flink.table.data.RowData; @@ -36,6 +37,7 @@ public class RowDataBucket { private final BucketInfo bucketInfo; private final BufferSizeDetector detector; private final String bucketId; + private boolean diverged; public RowDataBucket( String bucketId, @@ -57,9 +59,16 @@ public class RowDataBucket { } public boolean writeRow(RowData rowData) throws IOException { + ValidationUtils.checkState( + !diverged, + "RowData bucket " + bucketId + " diverged after a failed write and cannot be reused"); boolean success = dataBuffer.write(rowData); if (success) { detector.detect(rowData); + } else { + // BinaryInMemorySortBuffer may have partially appended variable-length data before + // returning false. Its internal pointers can no longer be trusted for another write. + diverged = true; } return success; } @@ -72,6 +81,10 @@ public class RowDataBucket { return detector.totalSize; } + public boolean isDiverged() { + return diverged; + } + public boolean isEmpty() { return dataBuffer.isEmpty(); } diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bucket/TestBucketStreamWriteMemoryExhaustion.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bucket/TestBucketStreamWriteMemoryExhaustion.java new file mode 100644 index 000000000000..14d47b7bce5d --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bucket/TestBucketStreamWriteMemoryExhaustion.java @@ -0,0 +1,243 @@ +/* + * 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.hudi.sink.bucket; + +import org.apache.hudi.client.WriteStatus; +import org.apache.hudi.common.model.HoodieTableType; +import org.apache.hudi.configuration.FlinkOptions; +import org.apache.hudi.sink.StreamWriteFunction; +import org.apache.hudi.sink.buffer.RowDataBucket; +import org.apache.hudi.sink.utils.BucketStreamWriteFunctionWrapper; +import org.apache.hudi.utils.TestConfigurations; +import org.apache.hudi.utils.TestData; + +import org.apache.avro.generic.GenericRecord; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.runtime.operators.coordination.OperatorEvent; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.RowType; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** End-to-end tests for bucket writes under a tight write memory pool. */ +class TestBucketStreamWriteMemoryExhaustion { + + @TempDir + File tempFile; + + @Test + void testRecoveryPreservesVariableLengthValuesAndReleasesPages() throws Exception { + DataType dataType = DataTypes.ROW( + DataTypes.FIELD("uuid", DataTypes.VARCHAR(20)), + DataTypes.FIELD("payload", DataTypes.VARCHAR(Integer.MAX_VALUE)), + DataTypes.FIELD("attributes", DataTypes.MAP(DataTypes.VARCHAR(64), DataTypes.VARCHAR(64))), + DataTypes.FIELD("ts", DataTypes.TIMESTAMP(3)), + DataTypes.FIELD("partition", DataTypes.VARCHAR(10))) + .notNull(); + RowType rowType = (RowType) dataType.getLogicalType(); + + Configuration conf = TestConfigurations.getDefaultConf(tempFile.getAbsolutePath(), dataType); + conf.set(FlinkOptions.TABLE_TYPE, HoodieTableType.COPY_ON_WRITE.name()); + conf.set(FlinkOptions.OPERATION, "upsert"); + conf.set(FlinkOptions.INDEX_TYPE, "BUCKET"); + conf.set(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, 4); + // Leaves 1 MB for the shared RowData memory pool. + conf.set(FlinkOptions.WRITE_TASK_MAX_SIZE, 201.0); + conf.set(FlinkOptions.WRITE_MERGE_MAX_MEMORY, 100); + + Map<String, String> expectedPayloads = new HashMap<>(); + Map<String, String> expectedAttributes = new HashMap<>(); + List<RowData> rows = createRows(rowType, expectedPayloads, expectedAttributes); + + TrackingBucketWriteFunctionWrapper pipeline = + new TrackingBucketWriteFunctionWrapper(tempFile.getAbsolutePath(), conf); + pipeline.openFunction(); + int initialFreePages = pipeline.freePages(); + try { + boolean reclaimedOtherBucketBeforeDivergedBucket = false; + for (RowData row : rows) { + int flushCountBeforeInvoke = pipeline.writeBatchCount(); + pipeline.invoke(row); + List<FlushedBucket> invocationFlushes = pipeline.flushesFrom(flushCountBeforeInvoke); + for (int i = 1; i < invocationFlushes.size(); i++) { + FlushedBucket previous = invocationFlushes.get(i - 1); + FlushedBucket current = invocationFlushes.get(i); + if (current.diverged && !previous.diverged) { + assertTrue( + !current.bucketId.equals(previous.bucketId), + "memory reclamation should flush another bucket before the diverged bucket"); + reclaimedOtherBucketBeforeDivergedBucket = true; + } + } + } + int recoveryFlushCount = pipeline.writeBatchCount(); + assertTrue(recoveryFlushCount > 0, "the tight pool should trigger at least one recovery flush"); + assertTrue( + reclaimedOtherBucketBeforeDivergedBucket, + "a write failure should reclaim another bucket before disposing the diverged bucket"); + pipeline.checkpointFunction(1); + + assertEquals( + initialFreePages, + pipeline.freePages(), + "all pages should be returned after the checkpoint flush disposes every bucket"); + + handleWriteEvents(pipeline, recoveryFlushCount + 1); + pipeline.checkpointComplete(1); + + List<GenericRecord> actualRecords = TestData.readAllData(tempFile, rowType, 1); + assertEquals(rows.size(), actualRecords.size(), "memory exhaustion recovery must not lose or duplicate records"); + for (GenericRecord actualRecord : actualRecords) { + String id = actualRecord.get("uuid").toString(); + assertTrue( + actualRecord.get("payload").toString().equals(expectedPayloads.get(id)), + "variable-length payload should remain intact for " + id); + Map<?, ?> actualAttributes = (Map<?, ?>) actualRecord.get("attributes"); + assertEquals(1, actualAttributes.size()); + Map.Entry<?, ?> attribute = actualAttributes.entrySet().iterator().next(); + assertEquals( + expectedAttributes.get(id), + attribute.getKey().toString() + "=" + attribute.getValue().toString()); + } + } finally { + pipeline.close(); + } + } + + private static List<RowData> createRows( + RowType rowType, + Map<String, String> expectedPayloads, + Map<String, String> expectedAttributes) { + List<RowData> rows = new ArrayList<>(); + for (int i = 0; i < 40; i++) { + String id = "uuid-" + i; + String payload = "payload-" + i + "-" + repeat((char) ('a' + i % 26), 256 * 1024); + String attributeKey = "key-" + i; + String attributeValue = "value-" + i; + Map<Object, Object> attributes = new HashMap<>(); + attributes.put(StringData.fromString(attributeKey), StringData.fromString(attributeValue)); + + rows.add(TestData.insertRow( + rowType, + StringData.fromString(id), + StringData.fromString(payload), + new GenericMapData(attributes), + TimestampData.fromEpochMillis(i), + StringData.fromString("par0"))); + expectedPayloads.put(id, payload); + expectedAttributes.put(id, attributeKey + "=" + attributeValue); + } + return rows; + } + + private static void handleWriteEvents( + TrackingBucketWriteFunctionWrapper pipeline, int eventCount) { + for (int i = 0; i < eventCount; i++) { + OperatorEvent event = pipeline.getNextEvent(); + pipeline.getCoordinator().handleEventFromOperator(0, event); + } + } + + private static String repeat(char value, int count) { + return String.valueOf(value).repeat(count); + } + + private static class TrackingBucketWriteFunctionWrapper + extends BucketStreamWriteFunctionWrapper<RowData> { + + TrackingBucketWriteFunctionWrapper(String basePath, Configuration conf) throws Exception { + super(basePath, conf); + } + + @Override + protected StreamWriteFunction createWriteFunction() { + return new TrackingBucketStreamWriteFunction(conf, rowType); + } + + int freePages() { + TrackingBucketStreamWriteFunction function = + (TrackingBucketStreamWriteFunction) writeFunction; + return function.freePages(); + } + + int writeBatchCount() { + TrackingBucketStreamWriteFunction function = + (TrackingBucketStreamWriteFunction) writeFunction; + return function.writeBatchCount(); + } + + List<FlushedBucket> flushesFrom(int startIndex) { + TrackingBucketStreamWriteFunction function = + (TrackingBucketStreamWriteFunction) writeFunction; + return function.flushesFrom(startIndex); + } + } + + private static class TrackingBucketStreamWriteFunction extends BucketStreamWriteFunction { + + private final List<FlushedBucket> flushedBuckets = new ArrayList<>(); + + TrackingBucketStreamWriteFunction(Configuration conf, RowType rowType) { + super(conf, rowType); + } + + int freePages() { + return memorySegmentPool.freePages(); + } + + int writeBatchCount() { + return flushedBuckets.size(); + } + + List<FlushedBucket> flushesFrom(int startIndex) { + return new ArrayList<>(flushedBuckets.subList(startIndex, flushedBuckets.size())); + } + + @Override + protected List<WriteStatus> writeRecords(String instant, RowDataBucket rowDataBucket) { + List<WriteStatus> writeStatuses = super.writeRecords(instant, rowDataBucket); + flushedBuckets.add(new FlushedBucket(rowDataBucket.getBucketId(), rowDataBucket.isDiverged())); + return writeStatuses; + } + } + + private static class FlushedBucket { + private final String bucketId; + private final boolean diverged; + + private FlushedBucket(String bucketId, boolean diverged) { + this.bucketId = bucketId; + this.diverged = diverged; + } + } +} diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/buffer/TestRowDataBucket.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/buffer/TestRowDataBucket.java new file mode 100644 index 000000000000..90d23b367ff0 --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/buffer/TestRowDataBucket.java @@ -0,0 +1,104 @@ +/* + * 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.hudi.sink.buffer; + +import org.apache.hudi.sink.utils.BufferUtils; +import org.apache.hudi.table.action.commit.BucketInfo; +import org.apache.hudi.table.action.commit.BucketType; +import org.apache.hudi.utils.TestData; + +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.binary.BinaryRowData; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.util.MutableObjectIterator; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Tests for {@link RowDataBucket}. */ +class TestRowDataBucket { + + @Test + void testDivergedBucketCannotBeReusedAndReleasesAllPages() throws Exception { + RowType rowType = (RowType) DataTypes.ROW( + DataTypes.FIELD("uuid", DataTypes.VARCHAR(20)), + DataTypes.FIELD("payload", DataTypes.VARCHAR(Integer.MAX_VALUE))) + .notNull() + .getLogicalType(); + HeapMemorySegmentPool pool = new HeapMemorySegmentPool(32 * 1024, 256 * 1024); + int initialFreePages = pool.freePages(); + RowDataBucket bucket = new RowDataBucket( + "bucket-0", + BufferUtils.createBuffer(rowType, pool), + new BucketInfo(BucketType.INSERT, "file-0", "partition-0"), + 256.0); + + List<String> expectedIds = new ArrayList<>(); + String payload = repeat('x', 64 * 1024); + boolean writeFailed = false; + try { + for (int i = 0; i < 100; i++) { + String id = "uuid-" + i; + RowData row = TestData.insertRow( + rowType, StringData.fromString(id), StringData.fromString(payload)); + if (!bucket.writeRow(row)) { + writeFailed = true; + break; + } + expectedIds.add(id); + } + + assertTrue(writeFailed, "the tiny pool should cause BinaryInMemorySortBuffer.write to return false"); + assertTrue(bucket.isDiverged()); + assertFalse(bucket.isEmpty(), "successful rows written before exhaustion should remain readable"); + + MutableObjectIterator<BinaryRowData> iterator = bucket.getDataIterator(); + BinaryRowData reuse = new BinaryRowData(rowType.getFieldCount()); + List<String> actualIds = new ArrayList<>(); + BinaryRowData row; + while ((row = iterator.next(reuse)) != null) { + actualIds.add(row.getString(0).toString()); + assertTrue(payload.equals(row.getString(1).toString()), "variable-length payload should remain intact"); + } + assertEquals(expectedIds, actualIds, "only successfully indexed rows should be readable"); + + RowData extraRow = TestData.insertRow( + rowType, StringData.fromString("uuid-extra"), StringData.fromString(payload)); + assertThrows(IllegalStateException.class, () -> bucket.writeRow(extraRow)); + } finally { + bucket.dispose(); + } + + assertEquals(initialFreePages, pool.freePages(), "disposing the diverged bucket should return every page"); + } + + private static String repeat(char value, int count) { + char[] chars = new char[count]; + java.util.Arrays.fill(chars, value); + return new String(chars); + } +}
