Examples with KvStateRegistryListener org.apache.flink.runtime.query.KvStateRegistryListener used on opensource projects

Example 1 with KvStateRegistryListener

use of org.apache.flink.runtime.query.KvStateRegistryListener in project flink by apache.

the class StateBackendTestBase method testQueryableStateRegistration.

/**
	 * Tests registration with the KvStateRegistry.
	 */
@Test
public void testQueryableStateRegistration() throws Exception {
    DummyEnvironment env = new DummyEnvironment("test", 1, 0);
    KvStateRegistry registry = env.getKvStateRegistry();
    CheckpointStreamFactory streamFactory = createStreamFactory();
    AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE, env);
    KeyGroupRange expectedKeyGroupRange = backend.getKeyGroupRange();
    KvStateRegistryListener listener = mock(KvStateRegistryListener.class);
    registry.registerListener(listener);
    ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("test", IntSerializer.INSTANCE);
    desc.setQueryable("banana");
    backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
    // Verify registered
    verify(listener, times(1)).notifyKvStateRegistered(eq(env.getJobID()), eq(env.getJobVertexId()), eq(expectedKeyGroupRange), eq("banana"), any(KvStateID.class));
    KeyGroupsStateHandle snapshot = FutureUtil.runIfNotDoneAndGet(backend.snapshot(682375462379L, 4, streamFactory, CheckpointOptions.forFullCheckpoint()));
    backend.dispose();
    verify(listener, times(1)).notifyKvStateUnregistered(eq(env.getJobID()), eq(env.getJobVertexId()), eq(expectedKeyGroupRange), eq("banana"));
    backend.dispose();
    // Initialize again
    backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot, env);
    snapshot.discardState();
    backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
    // Verify registered again
    verify(listener, times(2)).notifyKvStateRegistered(eq(env.getJobID()), eq(env.getJobVertexId()), eq(expectedKeyGroupRange), eq("banana"), any(KvStateID.class));
    backend.dispose();
}