diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/index/RecordLevelIndexBackend.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/index/RecordLevelIndexBackend.java index 8128c625158d0..783a3bd83befa 100644 --- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/index/RecordLevelIndexBackend.java +++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/index/RecordLevelIndexBackend.java @@ -245,7 +245,8 @@ private void updateEvictableCkp(long checkpointId) { minRetainedCheckpointId = checkpointId; } - private long inferMemorySizeForCache() { + @VisibleForTesting + long inferMemorySizeForCache() { int concurrentPartitionsNum = conf.get(FlinkOptions.INDEX_RLI_CACHE_CONCURRENT_PARTITIONS_NUM); if (partitionBucketCaches.isEmpty()) { return maxCacheSizeInBytes / concurrentPartitionsNum; diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithBufferSort.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithBufferSort.java index 70dc60f96a6ec..fd183ea9b8bc4 100644 --- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithBufferSort.java +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithBufferSort.java @@ -22,29 +22,47 @@ import org.apache.hudi.common.util.queue.HoodieConsumer; import org.apache.hudi.configuration.FlinkOptions; import org.apache.hudi.sink.buffer.BufferType; +import org.apache.hudi.sink.utils.TestWriteBase; +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.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.logical.RowType; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import java.io.File; import java.util.Arrays; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; +import java.util.stream.Collectors; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; /** - * Unit tests for buffer sort configuration resolution in {@link AppendWriteFunctions}. + * Unit tests for append write functions with buffer sorting. */ public class TestAppendWriteFunctionWithBufferSort { private Configuration conf; + @TempDir + File tempFile; + @BeforeEach void setUp() { conf = new Configuration(); @@ -81,6 +99,114 @@ public void testCustomSortKeysOverrideDefault() { assertEquals(Arrays.asList("age", "name"), resolvedSortKeys); } + @Test + public void testEmptySortKeysAreRejected() { + conf.set(FlinkOptions.WRITE_BUFFER_SORT_KEYS, " , , "); + + assertThrows(IllegalArgumentException.class, () -> AppendWriteFunctions.resolveSortKeys(conf)); + } + + @ParameterizedTest + @EnumSource(BufferType.class) + public void testFactorySelectsConfiguredBuffer(BufferType bufferType) { + conf.set(FlinkOptions.WRITE_BUFFER_TYPE, bufferType.name()); + RowType rowType = TestConfigurations.ROW_TYPE; + + AppendWriteFunction function = AppendWriteFunctions.create(conf, rowType); + + switch (bufferType) { + case CONTINUOUS_SORT: + assertInstanceOf(AppendWriteFunctionWithContinuousSort.class, function); + break; + case DISRUPTOR: + assertInstanceOf(AppendWriteFunctionWithDisruptorBufferSort.class, function); + break; + case BOUNDED_IN_MEMORY: + assertInstanceOf(AppendWriteFunctionWithBIMBufferSort.class, function); + break; + case NONE: + assertEquals(AppendWriteFunction.class, function.getClass()); + break; + default: + throw new AssertionError("Unexpected buffer type " + bufferType); + } + } + + @ParameterizedTest + @EnumSource(value = BufferType.class, names = {"CONTINUOUS_SORT", "DISRUPTOR", "BOUNDED_IN_MEMORY"}) + public void testBufferedAppendFunctionsWriteAndFlushOnCheckpoint(BufferType bufferType) throws Exception { + Configuration writeConf = TestConfigurations.getDefaultConf(tempFile.getAbsolutePath()); + writeConf.set(FlinkOptions.OPERATION, "insert"); + writeConf.set(FlinkOptions.METADATA_ENABLED, false); + writeConf.set(FlinkOptions.WRITE_BUFFER_TYPE, bufferType.name()); + writeConf.set(FlinkOptions.WRITE_BUFFER_SORT_KEYS, "name,age"); + writeConf.set(FlinkOptions.WRITE_BUFFER_SIZE, 3L); + if (bufferType == BufferType.CONTINUOUS_SORT) { + writeConf.set(FlinkOptions.WRITE_BUFFER_SORT_CONTINUOUS_DRAIN_SIZE, 1); + } else if (bufferType == BufferType.DISRUPTOR) { + writeConf.set(FlinkOptions.WRITE_BUFFER_DISRUPTOR_RING_SIZE, 16); + } + + TestWriteBase.TestHarness harness = TestWriteBase.TestHarness.instance() + .preparePipeline(tempFile, writeConf); + try { + harness + .consume(TestData.DATA_SET_INSERT) + .checkpoint(1) + .assertNextEvent(4, "par1,par2,par3,par4") + .checkpointComplete(1); + } finally { + harness.end(); + } + } + + @ParameterizedTest + @EnumSource(value = BufferType.class, names = {"CONTINUOUS_SORT", "DISRUPTOR", "BOUNDED_IN_MEMORY"}) + public void testBufferedAppendFunctionsPersistRowsInSortOrder(BufferType bufferType) throws Exception { + Configuration writeConf = TestConfigurations.getDefaultConf(tempFile.getAbsolutePath()); + writeConf.set(FlinkOptions.OPERATION, "insert"); + writeConf.set(FlinkOptions.METADATA_ENABLED, false); + writeConf.set(FlinkOptions.WRITE_BUFFER_TYPE, bufferType.name()); + writeConf.set(FlinkOptions.WRITE_BUFFER_SORT_KEYS, "name,age"); + writeConf.set(FlinkOptions.WRITE_BUFFER_SIZE, 3L); + if (bufferType == BufferType.CONTINUOUS_SORT) { + writeConf.set(FlinkOptions.WRITE_BUFFER_SORT_CONTINUOUS_DRAIN_SIZE, 1); + } else if (bufferType == BufferType.DISRUPTOR) { + writeConf.set(FlinkOptions.WRITE_BUFFER_DISRUPTOR_RING_SIZE, 16); + } + + List inputData = Arrays.asList( + TestData.insertRow(StringData.fromString("uuid1"), StringData.fromString("Bob"), 30, + TimestampData.fromEpochMillis(1123), StringData.fromString("p1")), + TestData.insertRow(StringData.fromString("uuid2"), StringData.fromString("Alice"), 25, + TimestampData.fromEpochMillis(1124), StringData.fromString("p1")), + TestData.insertRow(StringData.fromString("uuid3"), StringData.fromString("Bob"), 21, + TimestampData.fromEpochMillis(31124), StringData.fromString("p1"))); + List expected = Arrays.asList( + "uuid2,Alice,25,1970-01-01 00:00:01.124,p1", + "uuid3,Bob,21,1970-01-01 00:00:31.124,p1", + "uuid1,Bob,30,1970-01-01 00:00:01.123,p1"); + + TestWriteBase.TestHarness harness = TestWriteBase.TestHarness.instance() + .preparePipeline(tempFile, writeConf); + try { + harness + .consume(inputData) + .checkpoint(1) + .assertNextEvent(1, "p1") + .checkpointComplete(1); + + List persistedRows = + TestData.readAllData(new File(writeConf.get(FlinkOptions.PATH)), TestConfigurations.ROW_TYPE, 1); + List actual = persistedRows.stream() + .map(TestData::filterOutVariablesWithoutHudiMetadata) + .collect(Collectors.toList()); + assertEquals(expected, actual); + } finally { + harness.end(); + } + } + // ------------------------------------------------------------------------- // Tests for DisruptorMessageQueue.waitUntilDrained() // ------------------------------------------------------------------------- diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bulk/TestBulkInsertWriteHelper.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bulk/TestBulkInsertWriteHelper.java index 773e0eefd3e54..c79f331237697 100644 --- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bulk/TestBulkInsertWriteHelper.java +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/bulk/TestBulkInsertWriteHelper.java @@ -48,6 +48,7 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; /** * Test cases for {@link BulkInsertWriterHelper}. @@ -105,6 +106,22 @@ void testWrite() throws Exception { TestData.checkWrittenData(tempFile, expected2, 4, TestBulkInsertWriteHelper::filterCommitTime); } + @Test + void testInvalidRecordIsWrappedAsIOException() { + HoodieFlinkTable table = FlinkTables.createTable(conf); + BulkInsertWriterHelper writerHelper = new BulkInsertWriterHelper( + conf, + table, + table.getConfig(), + WriteClientTestUtils.createNewInstantTime(), + 1, + 1, + 0, + TestConfigurations.ROW_TYPE); + + assertThrows(IOException.class, () -> writerHelper.write(new GenericRowData(0))); + } + private void assertWriteStatus(List writeStatusList) { String partitions = writeStatusList.stream() .map(writeStatus -> StringUtils.nullToEmpty(writeStatus.getStat().getPartitionPath())) diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/TestDynamicBucketAssignFunction.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/TestDynamicBucketAssignFunction.java new file mode 100644 index 0000000000000..277b5e61e4f9a --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/TestDynamicBucketAssignFunction.java @@ -0,0 +1,197 @@ +/* + * 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.partitioner; + +import org.apache.hudi.client.model.HoodieFlinkInternalRow; +import org.apache.hudi.common.table.view.FileSystemViewStorageConfig; +import org.apache.hudi.common.table.view.FileSystemViewStorageType; +import org.apache.hudi.configuration.FlinkOptions; +import org.apache.hudi.sink.event.Correspondent; +import org.apache.hudi.sink.partitioner.index.DummyPartitionedIndexBackend; +import org.apache.hudi.sink.partitioner.index.PartitionedIndexBackend; +import org.apache.hudi.sink.utils.MockStreamingRuntimeContext; +import org.apache.hudi.table.action.commit.BucketInfo; +import org.apache.hudi.table.action.commit.BucketType; +import org.apache.hudi.util.StreamerUtil; +import org.apache.hudi.util.ViewStorageProperties; +import org.apache.hudi.utils.TestConfigurations; + +import org.apache.flink.api.common.functions.RuntimeContext; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.metrics.groups.UnregisteredMetricsGroup; +import org.apache.flink.runtime.state.FunctionInitializationContext; +import org.apache.flink.runtime.state.FunctionSnapshotContext; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.util.Collector; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Test cases for {@link DynamicBucketAssignFunction}. + */ +class TestDynamicBucketAssignFunction { + + @TempDir + File tempFile; + + @Test + void testRoutesExistingRecordToUpdateBucket() throws Exception { + DynamicBucketAssignFunction function = new DynamicBucketAssignFunction(new Configuration()); + PartitionedIndexBackend indexBackend = mock(PartitionedIndexBackend.class); + BucketAssigner bucketAssigner = mock(BucketAssigner.class); + when(indexBackend.get("partition", "key")).thenReturn("existing-file"); + when(bucketAssigner.addUpdate("partition", "existing-file")) + .thenReturn(new BucketInfo(BucketType.UPDATE, "existing-file", "partition")); + setField(function, "indexBackend", indexBackend); + setField(function, "bucketAssigner", bucketAssigner); + HoodieFlinkInternalRow record = record("key", "partition", "U"); + List output = new ArrayList<>(); + + function.processElement(record, null, collector(output)); + + assertEquals(1, output.size()); + assertEquals("existing-file", record.getFileId()); + assertEquals("U", record.getInstantTime()); + assertEquals("U", record.getOperationType()); + verify(indexBackend, never()).update("partition", "key", "existing-file"); + } + + @Test + void testAssignsAndCachesNewRecord() throws Exception { + DynamicBucketAssignFunction function = new DynamicBucketAssignFunction(new Configuration()); + PartitionedIndexBackend indexBackend = mock(PartitionedIndexBackend.class); + BucketAssigner bucketAssigner = mock(BucketAssigner.class); + when(indexBackend.get("partition", "key")).thenReturn(null); + when(bucketAssigner.addInsert("partition")) + .thenReturn(new BucketInfo(BucketType.INSERT, "new-file", "partition")); + setField(function, "indexBackend", indexBackend); + setField(function, "bucketAssigner", bucketAssigner); + HoodieFlinkInternalRow record = record("key", "partition", "U"); + List output = new ArrayList<>(); + + function.processElement(record, null, collector(output)); + + assertEquals(1, output.size()); + assertEquals("new-file", record.getFileId()); + assertEquals("I", record.getInstantTime()); + assertEquals("I", record.getOperationType()); + verify(indexBackend).update("partition", "key", "new-file"); + } + + @Test + void testCheckpointLifecycleDelegatesToBackends() throws Exception { + DynamicBucketAssignFunction function = new DynamicBucketAssignFunction(new Configuration()); + PartitionedIndexBackend indexBackend = mock(PartitionedIndexBackend.class); + BucketAssigner bucketAssigner = mock(BucketAssigner.class); + Correspondent correspondent = mock(Correspondent.class); + FunctionSnapshotContext snapshotContext = mock(FunctionSnapshotContext.class); + when(snapshotContext.getCheckpointId()).thenReturn(7L); + setField(function, "indexBackend", indexBackend); + setField(function, "bucketAssigner", bucketAssigner); + function.setCorrespondent(correspondent); + + function.snapshotState(snapshotContext); + function.notifyCheckpointComplete(7L); + function.close(); + + verify(bucketAssigner).reset(); + verify(indexBackend).onCheckpoint(7L); + verify(bucketAssigner).reload(7L); + verify(indexBackend).onCheckpointComplete(correspondent, 7L); + verify(indexBackend).close(); + verify(bucketAssigner).close(); + } + + @Test + void testInsertOverwriteUsesDummyIndexBackend() throws Exception { + Configuration conf = new Configuration(); + conf.set(FlinkOptions.OPERATION, "insert_overwrite"); + DynamicBucketAssignFunction function = new DynamicBucketAssignFunction(conf); + RuntimeContext runtimeContext = mock(RuntimeContext.class); + when(runtimeContext.getMetricGroup()).thenReturn(UnregisteredMetricsGroup.createOperatorMetricGroup()); + function.setRuntimeContext(runtimeContext); + + function.initializeState(mock(FunctionInitializationContext.class)); + + assertInstanceOf(DummyPartitionedIndexBackend.class, getField(function, "indexBackend")); + } + + @Test + void testOpenInitializesBucketAssignerAndTaskOwnership() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempFile.getAbsolutePath()); + conf.set(FlinkOptions.OPERATION, "insert_overwrite"); + StreamerUtil.initTableIfNotExists(conf); + ViewStorageProperties.createProperties( + conf.get(FlinkOptions.PATH), + FileSystemViewStorageConfig.newBuilder() + .withStorageType(FileSystemViewStorageType.SPILLABLE_DISK) + .build(), + conf); + DynamicBucketAssignFunction function = new DynamicBucketAssignFunction(conf); + function.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0)); + + function.open(new Configuration()); + function.initializeState(mock(FunctionInitializationContext.class)); + + assertInstanceOf(BucketAssigner.class, getField(function, "bucketAssigner")); + assertInstanceOf(DummyPartitionedIndexBackend.class, getField(function, "indexBackend")); + function.close(); + } + + private static HoodieFlinkInternalRow record(String recordKey, String partitionPath, String operationType) { + return new HoodieFlinkInternalRow(recordKey, partitionPath, operationType, new GenericRowData(0)); + } + + private static Collector collector(List output) { + return new Collector() { + @Override + public void collect(HoodieFlinkInternalRow record) { + output.add(record); + } + + @Override + public void close() { + } + }; + } + + private static Object getField(Object target, String fieldName) throws Exception { + Field field = DynamicBucketAssignFunction.class.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(target); + } + + private static void setField(Object target, String fieldName, Object value) throws Exception { + Field field = DynamicBucketAssignFunction.class.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(target, value); + } +} diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/index/TestRecordLevelIndexBackend.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/index/TestRecordLevelIndexBackend.java index cea4db07294a1..c6f1a4191eedc 100644 --- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/index/TestRecordLevelIndexBackend.java +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/partitioner/index/TestRecordLevelIndexBackend.java @@ -40,6 +40,8 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.when; @@ -119,6 +121,54 @@ public void testGetDoesNotRefreshLastUpdatedCheckpoint() throws Exception { } } + @Test + public void testGetAndUpdateLazilyCreatedPartitionCache() throws Exception { + try (RecordLevelIndexBackend backend = createBackend()) { + backend.registerMetrics(new UnregisteredMetricsGroup()); + + assertNull(backend.get("new-partition", "new-key")); + + backend.onCheckpoint(3L); + backend.update("new-partition", "new-key", "new-file-group"); + + assertEquals("new-file-group", backend.get("new-partition", "new-key")); + assertEquals(1, backend.getPartitionBucketCaches().size()); + } + } + + @Test + public void testNewPartitionUsesAverageExistingCacheSize() throws Exception { + try (RecordLevelIndexBackend backend = createBackend()) { + backend.registerMetrics(new UnregisteredMetricsGroup()); + backend.getPartitionBucketCaches().put( + "existing", cacheWithHeapSize(backend, ONE_MB / 2, 1L)); + + assertEquals(ONE_MB / 2, backend.inferMemorySizeForCache()); + assertNull(backend.get("new-partition", "new-key")); + + assertTrue(backend.getPartitionBucketCaches().containsKey("existing")); + assertTrue(backend.getPartitionBucketCaches().containsKey("new-partition")); + } + } + + @Test + public void testCheckpointCompletionUsesOldestInflightCheckpoint() throws Exception { + try (RecordLevelIndexBackend backend = createBackend()) { + backend.getPartitionBucketCaches().put("old", cacheWithHeapSize(backend, 2 * ONE_MB, 1L)); + Map inflightInstants = new HashMap<>(); + inflightInstants.put(2L, "002"); + inflightInstants.put(4L, "004"); + + backend.onCheckpointComplete(new TestCorrespondent(inflightInstants), 5L); + backend.cleanIfNecessary(0L, null); + + assertFalse(backend.getPartitionBucketCaches().containsKey("old")); + assertThrows(IllegalArgumentException.class, + () -> backend.onCheckpointComplete( + new TestCorrespondent(Collections.singletonMap(1L, "001")), 6L)); + } + } + @Test public void testPartitionCacheDictionaryEncodesFileGroupId() throws Exception { try (RecordLevelIndexBackend backend = createBackend()) { @@ -128,10 +178,13 @@ public void testPartitionCacheDictionaryEncodesFileGroupId() throws Exception { cache.putRecordKey("key1", "file-group-id-000000000000000000000001"); cache.putRecordKey("key2", "file-group-id-000000000000000000000001"); cache.putRecordKey("key3", "file-group-id-000000000000000000000002"); + cache.bootstrapRecordKey("key4", "file-group-id-000000000000000000000002"); assertEquals("file-group-id-000000000000000000000001", cache.getFileGroupId("key1")); assertEquals("file-group-id-000000000000000000000001", cache.getFileGroupId("key2")); assertEquals("file-group-id-000000000000000000000002", cache.getFileGroupId("key3")); + assertEquals("file-group-id-000000000000000000000002", cache.getFileGroupId("key4")); + assertEquals(4, cache.size()); assertEquals(Integer.valueOf(0), recordKeyToFileGroupIdCode.get("key1")); assertEquals(Integer.valueOf(0), recordKeyToFileGroupIdCode.get("key2")); assertEquals(Integer.valueOf(1), recordKeyToFileGroupIdCode.get("key3")); @@ -166,6 +219,18 @@ public void testRegisterMetricsRegistersPartitionBootstrapHistograms() throws Ex } } + @Test + public void testLazyEvictReturnsWhenOnlyProtectedPartitionExceedsLimit() throws Exception { + try (RecordLevelIndexBackend backend = createBackend()) { + backend.getPartitionBucketCaches().put("protected", cacheWithHeapSize(backend, 2 * ONE_MB, 1L)); + backend.onCheckpointComplete(new TestCorrespondent(Collections.emptyMap()), 2L); + + backend.cleanIfNecessary(0L, "protected"); + + assertTrue(backend.getPartitionBucketCaches().containsKey("protected")); + } + } + private RecordLevelIndexBackend createBackend() { return new RecordLevelIndexBackend(conf, (partitionPath, recordKey, fileId) -> true); } diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/ClusteringFunctionWrapper.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/ClusteringFunctionWrapper.java index b01208a9ff4f7..e874078e371af 100644 --- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/ClusteringFunctionWrapper.java +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/ClusteringFunctionWrapper.java @@ -123,6 +123,10 @@ public void cluster(long checkpointID) throws Exception { } public void close() throws Exception { - ioManager.close(); + TestFunctionWrapper.closeAll( + clusteringPlanOperator == null ? null : clusteringPlanOperator::close, + clusteringOperator == null ? null : clusteringOperator::close, + commitSink == null ? null : commitSink::close, + ioManager::close); } } diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/InsertFunctionWrapper.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/InsertFunctionWrapper.java index 9a8262155973a..aabef24d5d04b 100644 --- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/InsertFunctionWrapper.java +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/InsertFunctionWrapper.java @@ -24,6 +24,7 @@ import org.apache.hudi.sink.StreamWriteOperatorCoordinator; import org.apache.hudi.sink.append.AppendWriteFunction; import org.apache.hudi.sink.append.AppendWriteFunctions; +import org.apache.hudi.sink.buffer.MemorySegmentPoolFactory; import org.apache.hudi.sink.bulk.BulkInsertWriterHelper; import org.apache.hudi.sink.common.AbstractWriteFunction; import org.apache.hudi.sink.event.WriteMetadataEvent; @@ -66,6 +67,7 @@ public class InsertFunctionWrapper implements TestFunctionWrapper { private final MockOperatorEventGateway gateway; private final MockSubtaskGateway subtaskGateway; private final MockOperatorCoordinatorContext coordinatorContext; + private final IOManager ioManager; @Getter private StreamWriteOperatorCoordinator coordinator; private final MockStateInitializationContext stateInitializationContext; @@ -84,11 +86,11 @@ public InsertFunctionWrapper(String tablePath, Configuration conf) throws Except } public InsertFunctionWrapper(String tablePath, Configuration conf, ExecutionConfig executionConfig) throws Exception { - IOManager ioManager = new IOManagerAsync(); + this.ioManager = new IOManagerAsync(); MockEnvironment environment = new MockEnvironmentBuilder() .setTaskName("mockTask") .setManagedMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE) - .setIOManager(ioManager) + .setIOManager(this.ioManager) .setExecutionConfig(executionConfig) .build(); this.runtimeContext = new MockStreamingRuntimeContext(false, 1, 0, environment, executionConfig); @@ -209,10 +211,11 @@ public AbstractWriteFunction getWriteFunction() { @Override public void close() throws Exception { - this.coordinator.close(); - if (clusteringFunctionWrapper != null) { - clusteringFunctionWrapper.close(); - } + TestFunctionWrapper.closeAll( + writeFunction == null ? null : writeFunction::close, + coordinator::close, + clusteringFunctionWrapper == null ? null : clusteringFunctionWrapper::close, + ioManager::close); } public BulkInsertWriterHelper getWriterHelper() { @@ -228,6 +231,7 @@ private void setupWriteFunction() throws Exception { writeFunction.setRuntimeContext(runtimeContext); writeFunction.setOperatorEventGateway(gateway); writeFunction.initializeState(this.stateInitializationContext); + writeFunction.setMemorySegmentPoolFactory(new MemorySegmentPoolFactory(null, null, -1)); writeFunction.open(conf); writeFunction.setCorrespondent(new MockCorrespondent(this.coordinator)); // set up subtask gateway diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFunctionWrapper.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFunctionWrapper.java index 238306b801f0b..1e9c428f18921 100644 --- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFunctionWrapper.java +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFunctionWrapper.java @@ -35,6 +35,30 @@ * Define the common interfaces for test function wrappers. */ public interface TestFunctionWrapper { + /** + * Closes all the given resources in order and preserves any subsequent failures as suppressed exceptions. + */ + static void closeAll(AutoCloseable... closeables) throws Exception { + Exception firstException = null; + for (AutoCloseable closeable : closeables) { + if (closeable == null) { + continue; + } + try { + closeable.close(); + } catch (Exception e) { + if (firstException == null) { + firstException = e; + } else { + firstException.addSuppressed(e); + } + } + } + if (firstException != null) { + throw firstException; + } + } + /** * Open all the functions within this wrapper. */