Examples with TestingLeaderCallbackHandler org.apache.flink.kubernetes.kubeclient.resources.TestingLeaderCallbackHandler used on opensource projects

Example 1 with TestingLeaderCallbackHandler

use of org.apache.flink.kubernetes.kubeclient.resources.TestingLeaderCallbackHandler in project flink by apache.

the class KubernetesStateHandleStoreITCase method testMultipleKubernetesStateHandleStores.

@Test
public void testMultipleKubernetesStateHandleStores() throws Exception {
    final Configuration configuration = kubernetesResource.getConfiguration();
    final String leaderConfigMapName = LEADER_CONFIGMAP_NAME + System.currentTimeMillis();
    final int leaderNum = 3;
    final KubernetesLeaderElector[] leaderElectors = new KubernetesLeaderElector[leaderNum];
    final FlinkKubeClient[] kubeClients = new FlinkKubeClient[leaderNum];
    final TestingLeaderCallbackHandler[] leaderCallbackHandlers = new TestingLeaderCallbackHandler[leaderNum];
    @SuppressWarnings("unchecked") final KubernetesStateHandleStore<TestingLongStateHandleHelper.LongStateHandle>[] stateHandleStores = new KubernetesStateHandleStore[leaderNum];
    try {
        for (int i = 0; i < leaderNum; i++) {
            final String lockIdentity = UUID.randomUUID().toString();
            kubeClients[i] = kubeClientFactory.fromConfiguration(configuration, "testing");
            leaderCallbackHandlers[i] = new TestingLeaderCallbackHandler(lockIdentity);
            leaderElectors[i] = kubeClients[i].createLeaderElector(new KubernetesLeaderElectionConfiguration(leaderConfigMapName, lockIdentity, configuration), leaderCallbackHandlers[i]);
            stateHandleStores[i] = new KubernetesStateHandleStore<>(kubeClients[i], leaderConfigMapName, new TestingLongStateHandleHelper(), (ignore) -> true, lockIdentity);
            leaderElectors[i].run();
        }
        // Wait for the leader
        final String lockIdentity = TestingLeaderCallbackHandler.waitUntilNewLeaderAppears(TIMEOUT);
        Long expectedState = null;
        for (int i = 0; i < leaderNum; i++) {
            // leader
            if (leaderCallbackHandlers[i].getLockIdentity().equals(lockIdentity)) {
                expectedState = (long) i;
            }
            stateHandleStores[i].addAndLock(KEY, new TestingLongStateHandleHelper.LongStateHandle(i));
        }
        // Only the leader could add successfully
        assertThat(expectedState, is(notNullValue()));
        assertThat(stateHandleStores[0].getAllAndLock().size(), is(1));
        assertThat(stateHandleStores[0].getAllAndLock().get(0).f0.retrieveState().getValue(), is(expectedState));
        assertThat(stateHandleStores[0].getAllAndLock().get(0).f1, is(KEY));
    } finally {
        TestingLongStateHandleHelper.clearGlobalState();
        // Cleanup the resources
        for (int i = 0; i < leaderNum; i++) {
            if (leaderElectors[i] != null) {
                leaderElectors[i].stop();
            }
            if (kubeClients[i] != null) {
                kubeClients[i].close();
            }
        }
        kubernetesResource.getFlinkKubeClient().deleteConfigMap(leaderConfigMapName).get();
    }
}