[hotfix] Fix RocksDB resource handling in RocksKeyGroupsRocksSingleSt…

…ateIteratorTest
ruslandanilin · Feb 5, 2021 · f63b9ad · f63b9ad
1 parent f56021c
commit f63b9ad
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/.../org/apache/flink/contrib/streaming/state/RocksKeyGroupsRocksSingleStateIteratorTest.java b/.../org/apache/flink/contrib/streaming/state/RocksKeyGroupsRocksSingleStateIteratorTest.java
@@ -118,20 +118,21 @@ public void testMergeIterator(int maxParallelism) throws Exception {
                 totalKeysExpected += numKeys;
             }
 
+            CloseableRegistry closeableRegistry = new CloseableRegistry();
             int id = 0;
             for (Tuple2<ColumnFamilyHandle, Integer> columnFamilyHandle :
                     columnFamilyHandlesWithKeyCount) {
-                rocksIteratorsWithKVStateId.add(
-                        new Tuple2<>(
-                                RocksDBOperationUtils.getRocksIterator(
-                                        rocksDB, columnFamilyHandle.f0, readOptions),
-                                id));
+                RocksIteratorWrapper rocksIterator =
+                        RocksDBOperationUtils.getRocksIterator(
+                                rocksDB, columnFamilyHandle.f0, readOptions);
+                closeableRegistry.registerCloseable(rocksIterator);
+                rocksIteratorsWithKVStateId.add(new Tuple2<>(rocksIterator, id));
                 ++id;
             }
 
             try (RocksStatesPerKeyGroupMergeIterator mergeIterator =
                     new RocksStatesPerKeyGroupMergeIterator(
-                            new CloseableRegistry(),
+                            closeableRegistry,
                             rocksIteratorsWithKVStateId,
                             maxParallelism <= Byte.MAX_VALUE ? 1 : 2)) {