Code review comment: I think it's better to move `if enableIncrementalCheckpointing` check here from `verifyRocksDBStateUploaderClosed` so that the uploader isn't created unnecessarily.  We could also use `RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue()` instead of `4` here.

Code before review:
public void setupRocksKeyedStateBackend() throws Exception {
blocker = new OneShotLatch();
waiter = new OneShotLatch();
testStreamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
testStreamFactory.setBlockerLatch(blocker);
testStreamFactory.setWaiterLatch(waiter);
testStreamFactory.setAfterNumberInvocations(10);
rocksDBStateUploader = spy(new RocksDBStateUploader(4));
prepareRocksDB();
keyedStateBackend =
RocksDBTestUtils.builderForTestDB(
TEMP_FOLDER
.newFolder(),
IntSerializer.INSTANCE,
spy(db),
defaultCFHandle,
optionsContainer.getColumnOptions())
.setEnableIncrementalCheckpointing(enableIncrementalCheckpointing)
.setRocksDBStateUploader(rocksDBStateUploader)
.build();
testState1 =
keyedStateBackend.getPartitionedState(
VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE,
new ValueStateDescriptor<>("TestState-1", Integer.class, 0));
testState2 =
keyedStateBackend.getPartitionedState(
VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE,
new ValueStateDescriptor<>("TestState-2", String.class, ""));
allCreatedCloseables = new ArrayList<>();
doAnswer(
new Answer<Object>() {
@Override
public Object answer(InvocationOnMock invocationOnMock)
throws Throwable {
RocksIterator rocksIterator =
spy((RocksIterator) invocationOnMock.callRealMethod());
allCreatedCloseables.add(rocksIterator);
return rocksIterator;
}
})
.when(keyedStateBackend.db)
.newIterator(any(ColumnFamilyHandle.class), any(ReadOptions.class));
doAnswer(
new Answer<Object>() {
@Override
public Object answer(InvocationOnMock invocationOnMock)
throws Throwable {
Snapshot snapshot =
spy((Snapshot) invocationOnMock.callRealMethod());
allCreatedCloseables.add(snapshot);
return snapshot;
}
})
.when(keyedStateBackend.db)
.getSnapshot();
doAnswer(
new Answer<Object>() {
@Override
public Object answer(InvocationOnMock invocationOnMock)
throws Throwable {
ColumnFamilyHandle snapshot =
spy((ColumnFamilyHandle) invocationOnMock.callRealMethod());
allCreatedCloseables.add(snapshot);
return snapshot;
}
})
.when(keyedStateBackend.db)
.createColumnFamily(any(ColumnFamilyDescriptor.class));
for (int i = 0; i < 100; ++i) {
keyedStateBackend.setCurrentKey(i);
testState1.update(4200 + i);
testState2.update("S-" + (4200 + i));
}
}

Line under review: rocksDBStateUploader = spy(new RocksDBStateUploader(4));

Code after addressing the review:
public void setupRocksKeyedStateBackend() throws Exception {
blocker = new OneShotLatch();
waiter = new OneShotLatch();
testStreamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
testStreamFactory.setBlockerLatch(blocker);
testStreamFactory.setWaiterLatch(waiter);
testStreamFactory.setAfterNumberInvocations(10);
prepareRocksDB();
RocksDBKeyedStateBackendBuilder keyedStateBackendBuilder =
RocksDBTestUtils.builderForTestDB(
TEMP_FOLDER
.newFolder(),
IntSerializer.INSTANCE,
spy(db),
defaultCFHandle,
optionsContainer.getColumnOptions())
.setEnableIncrementalCheckpointing(enableIncrementalCheckpointing);
if (enableIncrementalCheckpointing) {
rocksDBStateUploader =
spy(
new RocksDBStateUploader(
RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue()));
keyedStateBackendBuilder.setRocksDBStateUploader(rocksDBStateUploader);
}
keyedStateBackend = keyedStateBackendBuilder.build();
testState1 =
keyedStateBackend.getPartitionedState(
VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE,
new ValueStateDescriptor<>("TestState-1", Integer.class, 0));
testState2 =
keyedStateBackend.getPartitionedState(
VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE,
new ValueStateDescriptor<>("TestState-2", String.class, ""));
allCreatedCloseables = new ArrayList<>();
doAnswer(
new Answer<Object>() {
@Override
public Object answer(InvocationOnMock invocationOnMock)
throws Throwable {
RocksIterator rocksIterator =
spy((RocksIterator) invocationOnMock.callRealMethod());
allCreatedCloseables.add(rocksIterator);
return rocksIterator;
}
})
.when(keyedStateBackend.db)
.newIterator(any(ColumnFamilyHandle.class), any(ReadOptions.class));
doAnswer(
new Answer<Object>() {
@Override
public Object answer(InvocationOnMock invocationOnMock)
throws Throwable {
Snapshot snapshot =
spy((Snapshot) invocationOnMock.callRealMethod());
allCreatedCloseables.add(snapshot);
return snapshot;
}
})
.when(keyedStateBackend.db)
.getSnapshot();
doAnswer(
new Answer<Object>() {
@Override
public Object answer(InvocationOnMock invocationOnMock)
throws Throwable {
ColumnFamilyHandle snapshot =
spy((ColumnFamilyHandle) invocationOnMock.callRealMethod());
allCreatedCloseables.add(snapshot);
return snapshot;
}
})
.when(keyedStateBackend.db)
.createColumnFamily(any(ColumnFamilyDescriptor.class));
for (int i = 0; i < 100; ++i) {
keyedStateBackend.setCurrentKey(i);
testState1.update(4200 + i);
testState2.update("S-" + (4200 + i));
}
}