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 d5db8e63e6fa694570cb02d2b0f333545750a0ee Author: fhan <[email protected]> AuthorDate: Fri Aug 28 10:11:00 2026 +0800 perf(flink): preempt inactive write buckets on memory exhaustion (#19728) * perf(flink): preempt inactive write buckets on memory exhaustion * refactor(flink): simplify preemptive memory reclamation * refactor(flink): refine memory exhaustion recovery --------- Co-authored-by: fhan <[email protected]> (cherry picked from commit 1895de8a06480506793577563f29b93a46317206) --- .../org/apache/hudi/sink/StreamWriteFunction.java | 70 +++++---- .../sink/buffer/PreemptiveMemorySegmentPool.java | 135 ++++++++++++++++ .../TestBucketStreamWriteMemoryExhaustion.java | 169 +++++++++++++++------ .../buffer/TestPreemptiveMemorySegmentPool.java | 156 +++++++++++++++++++ 4 files changed, 447 insertions(+), 83 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 39a5463b597e..c7ff3a709c33 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 @@ -33,6 +33,7 @@ import org.apache.hudi.configuration.FlinkOptions; import org.apache.hudi.exception.HoodieException; import org.apache.hudi.metrics.FlinkStreamWriteMetrics; import org.apache.hudi.sink.buffer.MemorySegmentPoolFactory; +import org.apache.hudi.sink.buffer.PreemptiveMemorySegmentPool; import org.apache.hudi.sink.buffer.RowDataBucket; import org.apache.hudi.sink.buffer.TotalSizeTracer; import org.apache.hudi.sink.bulk.RowDataKeyGen; @@ -60,7 +61,6 @@ 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; @@ -139,7 +139,7 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink */ protected transient FlinkStreamWriteMetrics writeMetrics; - protected transient MemorySegmentPool memorySegmentPool; + protected transient PreemptiveMemorySegmentPool preemptiveMemorySegmentPool; protected transient RecordConverter recordConverter; @@ -195,7 +195,8 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink private void initBuffer() { this.buckets = new LinkedHashMap<>(); - this.memorySegmentPool = MemorySegmentPoolFactory.createMemorySegmentPool(config); + MemorySegmentPool delegate = MemorySegmentPoolFactory.createMemorySegmentPool(config); + this.preemptiveMemorySegmentPool = new PreemptiveMemorySegmentPool(delegate, this::preemptMemory); } private void initWriteFunction() { @@ -259,11 +260,16 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink RowDataBucket bucket = this.buckets.computeIfAbsent(bucketID, k -> new RowDataBucket( bucketID, - BufferUtils.createBuffer(rowType, memorySegmentPool), + BufferUtils.createBuffer(rowType, preemptiveMemorySegmentPool), getBucketInfo(record), this.config.get(FlinkOptions.WRITE_BATCH_SIZE))); - return bucket.writeRow(record.getRowData()); + this.preemptiveMemorySegmentPool.setCurrentOwner(bucketID); + try { + return bucket.writeRow(record.getRowData()); + } finally { + this.preemptiveMemorySegmentPool.clearCurrentOwner(); + } } catch (MemoryPagesExhaustedException e) { LOG.info("There are not enough free pages in the memory pool to create a buffer; flushing is required first."); return false; @@ -323,45 +329,43 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink // 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) { + if (!preemptMemory(bucketID)) { 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"); + // Allocation failures during writeRow have already tried to preempt inactive buckets. The + // diverged bucket only needs to flush its committed rows and return its own pages before retry. + flushAndDisposeBucket(failedBucket); + } - 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); - } + /** + * Flushes the largest non-empty bucket other than the excluded bucket to return its pages to the + * shared memory pool. + * + * <p>The excluded bucket is either in the middle of serializing a row or is about to retry buffer + * creation and must never be flushed here. + */ + private boolean preemptMemory(String excludedBucketID) { + RowDataBucket bucketToFlush = findLargestNonEmptyBucketExcluding(excludedBucketID); + if (bucketToFlush == null) { + return false; } + flushAndDisposeBucket(bucketToFlush); + return true; + } - if (failure != null) { - throw failure; - } + private RowDataBucket findLargestNonEmptyBucketExcluding(String excludedBucketID) { + return this.buckets.values().stream() + .filter(bucket -> !excludedBucketID.equals(bucket.getBucketId()) && !bucket.isEmpty()) + .max(Comparator.comparingLong(RowDataBucket::getBufferSize)) + .orElse(null); } private void retryBufferRecord( @@ -541,8 +545,8 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink } try { - if (this.memorySegmentPool instanceof Closeable) { - ((Closeable) this.memorySegmentPool).close(); + if (this.preemptiveMemorySegmentPool != null) { + this.preemptiveMemorySegmentPool.close(); } } catch (Exception e) { closeFailure = addCloseFailure(closeFailure, e); diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/PreemptiveMemorySegmentPool.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/PreemptiveMemorySegmentPool.java new file mode 100644 index 000000000000..4fa21cf74dfe --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/buffer/PreemptiveMemorySegmentPool.java @@ -0,0 +1,135 @@ +/* + * 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.common.util.ValidationUtils; + +import org.apache.flink.core.memory.MemorySegment; +import org.apache.flink.table.runtime.util.MemorySegmentPool; + +import javax.annotation.Nullable; + +import java.io.Closeable; +import java.io.IOException; +import java.util.List; + +/** + * A {@link MemorySegmentPool} wrapper that can reclaim pages from an inactive owner when the + * delegate pool is exhausted. + * + * <p>The owner represents the bucket whose buffer is currently requesting pages. The reclaimer + * must never reclaim that owner because the request may occur in the middle of serializing a row. + * If no other owner can release pages, this pool returns {@code null} and lets the caller handle + * the allocation failure. + */ +public class PreemptiveMemorySegmentPool implements MemorySegmentPool, Closeable { + + /** Callback that reclaims memory from an owner other than the excluded in-flight owner. */ + @FunctionalInterface + public interface MemoryReclaimer { + /** + * Reclaims memory from an inactive owner. + * + * @param excludedOwnerId the owner currently requesting a page + * @return {@code true} if memory was reclaimed and allocation should be retried + */ + boolean reclaim(String excludedOwnerId); + } + + private final MemorySegmentPool delegate; + private final MemoryReclaimer memoryReclaimer; + + /** + * The owner whose buffer is currently serializing a row. This is {@code null} outside + * {@code writeRow}, including while a new buffer is being created, so allocation failures in + * those contexts are handled by the caller's existing fallback path. + */ + @Nullable + private String currentOwnerId; + + /** Prevents an allocation made from the reclamation callback from recursively reclaiming. */ + private boolean preempting; + + public PreemptiveMemorySegmentPool( + MemorySegmentPool delegate, + MemoryReclaimer memoryReclaimer) { + ValidationUtils.checkArgument(delegate != null, "Delegate memory segment pool must not be null"); + ValidationUtils.checkArgument(memoryReclaimer != null, "Memory reclaimer must not be null"); + this.delegate = delegate; + this.memoryReclaimer = memoryReclaimer; + } + + /** Marks the owner whose buffer is currently requesting memory pages. */ + public void setCurrentOwner(String ownerId) { + ValidationUtils.checkArgument(ownerId != null, "Memory segment pool owner must not be null"); + ValidationUtils.checkState( + currentOwnerId == null, + "A memory segment pool owner is already active: " + currentOwnerId); + this.currentOwnerId = ownerId; + } + + /** Clears the current owner after its buffer write finishes. */ + public void clearCurrentOwner() { + this.currentOwnerId = null; + } + + @Override + public int pageSize() { + return delegate.pageSize(); + } + + @Override + public void returnAll(List<MemorySegment> memorySegments) { + delegate.returnAll(memorySegments); + } + + @Override + public int freePages() { + return delegate.freePages(); + } + + @Override + public MemorySegment nextSegment() { + MemorySegment segment = delegate.nextSegment(); + // Reclamation requires an in-flight owner and is disabled while a reclamation callback is + // running to prevent recursively selecting a victim that has not yet been disposed. + if (segment != null || currentOwnerId == null || preempting) { + return segment; + } + + preempting = true; + try { + if (!memoryReclaimer.reclaim(currentOwnerId)) { + return null; + } + // Retry this allocation once. A later page request may reclaim another inactive owner, + // while this bounded retry prevents recursion if the callback does not return any pages. + return delegate.nextSegment(); + } finally { + preempting = false; + } + } + + @Override + public void close() throws IOException { + if (delegate instanceof Closeable) { + ((Closeable) delegate).close(); + } + } +} 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 index 14d47b7bce5d..ae1ffa0de8d6 100644 --- 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 @@ -37,10 +37,11 @@ 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 org.junit.jupiter.api.Test; import java.io.File; +import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -57,54 +58,35 @@ class TestBucketStreamWriteMemoryExhaustion { @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(); + DataType dataType = dataType(); 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); + Configuration conf = writeConfig(dataType, 4); Map<String, String> expectedPayloads = new HashMap<>(); Map<String, String> expectedAttributes = new HashMap<>(); - List<RowData> rows = createRows(rowType, expectedPayloads, expectedAttributes); + List<RowData> rows = createRows(rowType, 40, expectedPayloads, expectedAttributes); TrackingBucketWriteFunctionWrapper pipeline = new TrackingBucketWriteFunctionWrapper(tempFile.getAbsolutePath(), conf); pipeline.openFunction(); int initialFreePages = pipeline.freePages(); try { - boolean reclaimedOtherBucketBeforeDivergedBucket = false; + boolean preemptedInactiveBucket = 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; - } + if (!invocationFlushes.isEmpty() + && invocationFlushes.stream().noneMatch(flushedBucket -> flushedBucket.diverged)) { + preemptedInactiveBucket = true; } } int recoveryFlushCount = pipeline.writeBatchCount(); - assertTrue(recoveryFlushCount > 0, "the tight pool should trigger at least one recovery flush"); + assertTrue(recoveryFlushCount > 0, "the tight pool should trigger at least one preemptive flush"); assertTrue( - reclaimedOtherBucketBeforeDivergedBucket, - "a write failure should reclaim another bucket before disposing the diverged bucket"); + preemptedInactiveBucket, + "memory exhaustion should flush an inactive bucket before the current buffer diverges"); pipeline.checkpointFunction(1); assertEquals( @@ -115,31 +97,94 @@ class TestBucketStreamWriteMemoryExhaustion { 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()); + assertWrittenRows(rowType, rows.size(), expectedPayloads, expectedAttributes); + } finally { + pipeline.close(); + } + } + + @Test + void testFlushesDivergedBucketWhenNoInactiveBucketCanBePreempted() + throws Exception { + DataType dataType = dataType(); + RowType rowType = (RowType) dataType.getLogicalType(); + // A single bucket guarantees that memory reclamation cannot find an inactive victim while + // the current bucket is serializing a row. + Configuration conf = writeConfig(dataType, 1); + + Map<String, String> expectedPayloads = new HashMap<>(); + Map<String, String> expectedAttributes = new HashMap<>(); + List<RowData> rows = createRows(rowType, 12, expectedPayloads, expectedAttributes); + + TrackingBucketWriteFunctionWrapper pipeline = + new TrackingBucketWriteFunctionWrapper(tempFile.getAbsolutePath(), conf); + pipeline.openFunction(); + int initialFreePages = pipeline.freePages(); + try { + boolean flushedDivergedBucket = false; + for (RowData row : rows) { + int flushCountBeforeInvoke = pipeline.writeBatchCount(); + pipeline.invoke(row); + List<FlushedBucket> invocationFlushes = pipeline.flushesFrom(flushCountBeforeInvoke); + if (!invocationFlushes.isEmpty()) { + assertTrue( + invocationFlushes.stream().allMatch(flushedBucket -> flushedBucket.diverged), + "without an inactive victim, recovery should only flush the diverged current bucket"); + flushedDivergedBucket = true; + } } + + int recoveryFlushCount = pipeline.writeBatchCount(); + assertTrue( + flushedDivergedBucket, + "memory exhaustion should flush the diverged bucket when no inactive victim exists"); + 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); + + assertWrittenRows(rowType, rows.size(), expectedPayloads, expectedAttributes); } finally { pipeline.close(); } } + private static DataType dataType() { + return 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(); + } + + private Configuration writeConfig(DataType dataType, int bucketCount) { + 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, bucketCount); + // 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); + // Prevent batch-size based flushes so mid-invocation flushes come from memory exhaustion. + conf.set(FlinkOptions.WRITE_BATCH_SIZE, 1024.0); + return conf; + } + private static List<RowData> createRows( RowType rowType, + int rowCount, Map<String, String> expectedPayloads, Map<String, String> expectedAttributes) { List<RowData> rows = new ArrayList<>(); - for (int i = 0; i < 40; i++) { + for (int i = 0; i < rowCount; i++) { String id = "uuid-" + i; String payload = "payload-" + i + "-" + repeat((char) ('a' + i % 26), 256 * 1024); String attributeKey = "key-" + i; @@ -160,6 +205,30 @@ class TestBucketStreamWriteMemoryExhaustion { return rows; } + private void assertWrittenRows( + RowType rowType, + int expectedRowCount, + Map<String, String> expectedPayloads, + Map<String, String> expectedAttributes) throws IOException { + List<GenericRecord> actualRecords = TestData.readAllData(tempFile, rowType, 1); + assertEquals( + expectedRowCount, + 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()); + } + } + private static void handleWriteEvents( TrackingBucketWriteFunctionWrapper pipeline, int eventCount) { for (int i = 0; i < eventCount; i++) { @@ -169,7 +238,9 @@ class TestBucketStreamWriteMemoryExhaustion { } private static String repeat(char value, int count) { - return String.valueOf(value).repeat(count); + char[] chars = new char[count]; + java.util.Arrays.fill(chars, value); + return new String(chars); } private static class TrackingBucketWriteFunctionWrapper @@ -212,7 +283,7 @@ class TestBucketStreamWriteMemoryExhaustion { } int freePages() { - return memorySegmentPool.freePages(); + return preemptiveMemorySegmentPool.freePages(); } int writeBatchCount() { @@ -226,17 +297,15 @@ class TestBucketStreamWriteMemoryExhaustion { @Override protected List<WriteStatus> writeRecords(String instant, RowDataBucket rowDataBucket) { List<WriteStatus> writeStatuses = super.writeRecords(instant, rowDataBucket); - flushedBuckets.add(new FlushedBucket(rowDataBucket.getBucketId(), rowDataBucket.isDiverged())); + flushedBuckets.add(new FlushedBucket(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; + private FlushedBucket(boolean diverged) { this.diverged = diverged; } } diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/buffer/TestPreemptiveMemorySegmentPool.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/buffer/TestPreemptiveMemorySegmentPool.java new file mode 100644 index 000000000000..6ed2023749f3 --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/buffer/TestPreemptiveMemorySegmentPool.java @@ -0,0 +1,156 @@ +/* + * 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.flink.core.memory.MemorySegment; +import org.apache.flink.runtime.memory.MemoryManager; +import org.apache.flink.table.runtime.util.LazyMemorySegmentPool; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Tests for {@link PreemptiveMemorySegmentPool}. */ +class TestPreemptiveMemorySegmentPool { + + private static final int PAGE_SIZE = 32 * 1024; + + @Test + void testReclaimsPageForCurrentOwner() { + HeapMemorySegmentPool delegate = new HeapMemorySegmentPool(PAGE_SIZE, PAGE_SIZE); + AtomicReference<String> excludedOwner = new AtomicReference<>(); + AtomicReference<PreemptiveMemorySegmentPool> poolReference = new AtomicReference<>(); + MemorySegment heldSegment = delegate.nextSegment(); + + PreemptiveMemorySegmentPool pool = new PreemptiveMemorySegmentPool(delegate, ownerId -> { + excludedOwner.set(ownerId); + poolReference.get().returnAll(Collections.singletonList(heldSegment)); + return true; + }); + poolReference.set(pool); + + pool.setCurrentOwner("bucket-0"); + MemorySegment reclaimedSegment; + try { + reclaimedSegment = pool.nextSegment(); + } finally { + pool.clearCurrentOwner(); + } + + assertEquals("bucket-0", excludedOwner.get()); + assertSame(heldSegment, reclaimedSegment); + pool.returnAll(Collections.singletonList(reclaimedSegment)); + assertEquals(1, pool.freePages()); + } + + @Test + void testDoesNotPreemptWithoutCurrentOwner() { + HeapMemorySegmentPool delegate = new HeapMemorySegmentPool(PAGE_SIZE, PAGE_SIZE); + delegate.nextSegment(); + AtomicInteger preemptionCount = new AtomicInteger(); + PreemptiveMemorySegmentPool pool = new PreemptiveMemorySegmentPool(delegate, ownerId -> { + preemptionCount.incrementAndGet(); + return true; + }); + + assertNull(pool.nextSegment()); + assertEquals(0, preemptionCount.get()); + } + + @Test + void testRetriesAllocationOnceWithoutNestedPreemption() { + HeapMemorySegmentPool delegate = new HeapMemorySegmentPool(PAGE_SIZE, PAGE_SIZE); + delegate.nextSegment(); + AtomicInteger preemptionCount = new AtomicInteger(); + AtomicReference<PreemptiveMemorySegmentPool> poolReference = new AtomicReference<>(); + PreemptiveMemorySegmentPool pool = new PreemptiveMemorySegmentPool(delegate, ownerId -> { + preemptionCount.incrementAndGet(); + assertNull(poolReference.get().nextSegment(), "nested allocation must not trigger preemption"); + return true; + }); + poolReference.set(pool); + + pool.setCurrentOwner("bucket-0"); + try { + assertNull(pool.nextSegment(), "allocation should fail when the callback returns no pages"); + } finally { + pool.clearCurrentOwner(); + } + assertEquals(1, preemptionCount.get()); + } + + @Test + void testResetsPreemptionStateAfterCallbackFailure() { + HeapMemorySegmentPool delegate = new HeapMemorySegmentPool(PAGE_SIZE, PAGE_SIZE); + delegate.nextSegment(); + AtomicInteger preemptionCount = new AtomicInteger(); + PreemptiveMemorySegmentPool pool = new PreemptiveMemorySegmentPool(delegate, ownerId -> { + if (preemptionCount.incrementAndGet() == 1) { + throw new IllegalStateException("reclamation failed"); + } + return false; + }); + + pool.setCurrentOwner("bucket-0"); + try { + assertThrows(IllegalStateException.class, pool::nextSegment); + assertNull(pool.nextSegment(), "a later allocation should be allowed to invoke the callback again"); + } finally { + pool.clearCurrentOwner(); + } + assertEquals(2, preemptionCount.get()); + } + + @Test + void testCloseReleasesManagedMemoryDelegate() throws Exception { + int numPages = 3; + MemoryManager memoryManager = MemoryManager.create((long) PAGE_SIZE * numPages, PAGE_SIZE); + LazyMemorySegmentPool delegate = + new LazyMemorySegmentPool(new Object(), memoryManager, numPages); + PreemptiveMemorySegmentPool pool = new PreemptiveMemorySegmentPool(delegate, ownerId -> false); + List<MemorySegment> allocatedSegments = new ArrayList<>(); + + try { + try { + for (int i = 0; i < numPages; i++) { + MemorySegment segment = pool.nextSegment(); + assertNotNull(segment); + allocatedSegments.add(segment); + } + assertEquals(0, pool.freePages()); + pool.returnAll(allocatedSegments); + } finally { + pool.close(); + } + assertTrue(memoryManager.verifyEmpty(), "closing the wrapper should release managed pages"); + } finally { + memoryManager.shutdown(); + } + } +}
