This is an automated email from the ASF dual-hosted git repository.
danny0405 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new 1895de8a0648 perf(flink): preempt inactive write buckets on memory
exhaustion (#19728)
1895de8a0648 is described below
commit 1895de8a06480506793577563f29b93a46317206
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]>
---
.../org/apache/hudi/sink/StreamWriteFunction.java | 73 +++++----
.../sink/buffer/PreemptiveMemorySegmentPool.java | 135 ++++++++++++++++
.../TestBucketStreamWriteMemoryExhaustion.java | 178 +++++++++++++++------
.../buffer/TestPreemptiveMemorySegmentPool.java | 156 ++++++++++++++++++
4 files changed, 457 insertions(+), 85 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 cdce134c3c82..ddf6a3fd1be1 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
@@ -34,6 +34,7 @@ import org.apache.hudi.configuration.FlinkOptions;
import org.apache.hudi.configuration.OptionsResolver;
import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.metrics.FlinkStreamWriteMetrics;
+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;
@@ -65,7 +66,6 @@ import org.apache.flink.table.runtime.util.MemorySegmentPool;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.util.Collector;
-import java.io.Closeable;
import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
@@ -147,7 +147,7 @@ public class StreamWriteFunction extends
AbstractStreamWriteFunction<HoodieFlink
*/
protected transient FlinkStreamWriteMetrics writeMetrics;
- protected transient MemorySegmentPool memorySegmentPool;
+ protected transient PreemptiveMemorySegmentPool preemptiveMemorySegmentPool;
protected transient RecordConverter recordConverter;
@@ -209,7 +209,9 @@ public class StreamWriteFunction extends
AbstractStreamWriteFunction<HoodieFlink
private void initBuffer() {
this.buckets = new LinkedHashMap<>();
- this.memorySegmentPool =
this.memorySegmentPoolFactory.createMemorySegmentPool(config,
OptionsResolver.getWriteBufferSizeInBytes(config));
+ MemorySegmentPool delegate =
this.memorySegmentPoolFactory.createMemorySegmentPool(
+ config, OptionsResolver.getWriteBufferSizeInBytes(config));
+ this.preemptiveMemorySegmentPool = new
PreemptiveMemorySegmentPool(delegate, this::preemptMemory);
}
private void initRecordKeySort() {
@@ -322,7 +324,12 @@ public class StreamWriteFunction extends
AbstractStreamWriteFunction<HoodieFlink
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;
@@ -379,45 +386,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(
@@ -563,12 +568,12 @@ public class StreamWriteFunction extends
AbstractStreamWriteFunction<HoodieFlink
private BinaryInMemorySortBuffer createDataBuffer() {
if (recordKeyComputer == null) {
- return BufferUtils.createBuffer(rowType, memorySegmentPool);
+ return BufferUtils.createBuffer(rowType, preemptiveMemorySegmentPool);
}
try {
return BufferUtils.createBuffer(
rowType,
- memorySegmentPool,
+ preemptiveMemorySegmentPool,
recordKeyComputer,
recordKeyComparator);
} catch (MemoryPagesExhaustedException e) {
@@ -621,8 +626,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..af44348215cd 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
@@ -20,6 +20,7 @@ package org.apache.hudi.sink.bucket;
import org.apache.hudi.client.WriteStatus;
import org.apache.hudi.common.model.HoodieTableType;
+import org.apache.hudi.common.table.HoodieTableConfig;
import org.apache.hudi.configuration.FlinkOptions;
import org.apache.hudi.sink.StreamWriteFunction;
import org.apache.hudi.sink.buffer.RowDataBucket;
@@ -37,10 +38,12 @@ 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.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.io.File;
+import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
@@ -55,56 +58,38 @@ 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();
+ @ParameterizedTest(name = "lsmTreeLayout={0}")
+ @ValueSource(booleans = {false, true})
+ void testRecoveryPreservesVariableLengthValuesAndReleasesPages(boolean
lsmTreeLayout) throws Exception {
+ 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, lsmTreeLayout);
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 +100,100 @@ 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();
+ }
+ }
+
+ @ParameterizedTest(name = "lsmTreeLayout={0}")
+ @ValueSource(booleans = {false, true})
+ void testFlushesDivergedBucketWhenNoInactiveBucketCanBePreempted(boolean
lsmTreeLayout)
+ 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, lsmTreeLayout);
+
+ 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,
boolean lsmTreeLayout) {
+ 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);
+ if (lsmTreeLayout) {
+ conf.setString(
+ HoodieTableConfig.TABLE_STORAGE_LAYOUT.key(),
+ HoodieTableConfig.TableStorageLayout.LSM_TREE.configValue());
+ }
+ // 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 +214,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++) {
@@ -212,7 +290,7 @@ class TestBucketStreamWriteMemoryExhaustion {
}
int freePages() {
- return memorySegmentPool.freePages();
+ return preemptiveMemorySegmentPool.freePages();
}
int writeBatchCount() {
@@ -226,17 +304,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();
+ }
+ }
+}
