Examples with RetrievableStateHandle org.apache.flink.runtime.state.RetrievableStateHandle used on opensource projects

Example 1 with RetrievableStateHandle

use of org.apache.flink.runtime.state.RetrievableStateHandle in project flink by apache.

the class ZooKeeperStateHandleStoreITCase method testAddWithCreateMode.

/**
	 * Tests that {@link CreateMode} is respected.
	 */
@Test
public void testAddWithCreateMode() throws Exception {
    LongStateStorage longStateStorage = new LongStateStorage();
    ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<Long>(ZooKeeper.getClient(), longStateStorage, Executors.directExecutor());
    // Config
    Long state = 3457347234L;
    CreateMode[] modes = CreateMode.values();
    for (int i = 0; i < modes.length; i++) {
        CreateMode mode = modes[i];
        state += i;
        String pathInZooKeeper = "/testAddWithCreateMode" + mode.name();
        // Test
        store.add(pathInZooKeeper, state, mode);
        if (mode.isSequential()) {
            // Figure out the sequential ID
            List<String> paths = ZooKeeper.getClient().getChildren().forPath("/");
            for (String p : paths) {
                if (p.startsWith("testAddWithCreateMode" + mode.name())) {
                    pathInZooKeeper = "/" + p;
                    break;
                }
            }
        }
        // Verify
        // State handle created
        assertEquals(i + 1, store.getAll().size());
        assertEquals(state, longStateStorage.getStateHandles().get(i).retrieveState());
        // Path created
        Stat stat = ZooKeeper.getClient().checkExists().forPath(pathInZooKeeper);
        assertNotNull(stat);
        // Is ephemeral or persistent
        if (mode.isEphemeral()) {
            assertTrue(stat.getEphemeralOwner() != 0);
        } else {
            assertEquals(0, stat.getEphemeralOwner());
        }
        // Data is equal
        @SuppressWarnings("unchecked") Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(ZooKeeper.getClient().getData().forPath(pathInZooKeeper), ClassLoader.getSystemClassLoader())).retrieveState();
        assertEquals(state, actual);
    }
}

Example 2 with RetrievableStateHandle

use of org.apache.flink.runtime.state.RetrievableStateHandle in project flink by apache.

the class ZooKeeperCompletedCheckpointStoreTest method testRecoverFailsIfDownloadFails.

@Test
public void testRecoverFailsIfDownloadFails() {
    final Configuration configuration = new Configuration();
    configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_QUORUM, zooKeeperResource.getConnectString());
    final List<Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String>> checkpointsInZk = new ArrayList<>();
    final ZooKeeperStateHandleStore<CompletedCheckpoint> checkpointsInZooKeeper = new ZooKeeperStateHandleStore<CompletedCheckpoint>(ZooKeeperUtils.startCuratorFramework(configuration, NoOpFatalErrorHandler.INSTANCE).asCuratorFramework(), new TestingRetrievableStateStorageHelper<>()) {

        @Override
        public List<Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String>> getAllAndLock() {
            return checkpointsInZk;
        }
    };
    checkpointsInZk.add(createHandle(1, id -> {
        throw new ExpectedTestException();
    }));
    final Exception exception = assertThrows(Exception.class, () -> DefaultCompletedCheckpointStoreUtils.retrieveCompletedCheckpoints(checkpointsInZooKeeper, zooKeeperCheckpointStoreUtil));
    assertThat(exception, FlinkMatchers.containsCause(ExpectedTestException.class));
}

Example 3 with RetrievableStateHandle

use of org.apache.flink.runtime.state.RetrievableStateHandle in project flink by apache.

the class KubernetesStateHandleStore method releaseAndTryRemoveAll.

/**
 * Remove all the state handle keys in the ConfigMap and discard the states.
 *
 * @throws Exception when removing the keys or discarding the state failed
 */
@Override
public void releaseAndTryRemoveAll() throws Exception {
    final List<RetrievableStateHandle<T>> validStateHandles = new ArrayList<>();
    kubeClient.checkAndUpdateConfigMap(configMapName, c -> {
        if (isValidOperation(c)) {
            final Map<String, String> updateData = new HashMap<>(c.getData());
            c.getData().entrySet().stream().filter(entry -> configMapKeyFilter.test(entry.getKey())).forEach(entry -> {
                try {
                    validStateHandles.add(deserializeObject(entry.getValue()));
                    updateData.remove(entry.getKey());
                } catch (IOException e) {
                    LOG.warn("ConfigMap {} contained corrupted data. Ignoring the key {}.", configMapName, entry.getKey());
                }
            });
            c.getData().clear();
            c.getData().putAll(updateData);
            return Optional.of(c);
        }
        return Optional.empty();
    }).whenComplete((succeed, ignore) -> {
        if (succeed) {
            Exception exception = null;
            for (RetrievableStateHandle<T> stateHandle : validStateHandles) {
                try {
                    stateHandle.discardState();
                } catch (Exception e) {
                    exception = ExceptionUtils.firstOrSuppressed(e, exception);
                }
            }
            if (exception != null) {
                throw new CompletionException(new KubernetesException("Could not properly remove all state handles.", exception));
            }
        }
    }).get();
}

Example 4 with RetrievableStateHandle

use of org.apache.flink.runtime.state.RetrievableStateHandle in project flink by apache.

the class ZooKeeperStateHandleStoreTest method testGetAllSortedByName.

/**
 * Tests that the state is returned sorted.
 */
@Test
public void testGetAllSortedByName() throws Exception {
    // Setup
    final TestingLongStateHandleHelper stateHandleProvider = new TestingLongStateHandleHelper();
    ZooKeeperStateHandleStore<TestingLongStateHandleHelper.LongStateHandle> store = new ZooKeeperStateHandleStore<>(ZOOKEEPER.getClient(), stateHandleProvider);
    // Config
    final String basePath = "/testGetAllSortedByName";
    final Long[] expected = new Long[] { 311222268470898L, 132812888L, 27255442L, 11122233124L };
    // Test
    for (long val : expected) {
        final String pathInZooKeeper = String.format("%s%016d", basePath, val);
        store.addAndLock(pathInZooKeeper, new TestingLongStateHandleHelper.LongStateHandle(val));
    }
    List<Tuple2<RetrievableStateHandle<TestingLongStateHandleHelper.LongStateHandle>, String>> actual = store.getAllAndLock();
    assertEquals(expected.length, actual.size());
    // bring the elements in sort order
    Arrays.sort(expected);
    Collections.sort(actual, Comparator.comparing(o -> o.f1));
    for (int i = 0; i < expected.length; i++) {
        assertEquals(expected[i], (Long) actual.get(i).f0.retrieveState().getValue());
    }
}

Example 5 with RetrievableStateHandle

use of org.apache.flink.runtime.state.RetrievableStateHandle in project flink by apache.

the class DefaultJobGraphStoreTest method testOnAddedJobGraphShouldOnlyProcessUnknownJobGraphs.

@Test
public void testOnAddedJobGraphShouldOnlyProcessUnknownJobGraphs() throws Exception {
    final RetrievableStateHandle<JobGraph> stateHandle = jobGraphStorageHelper.store(testingJobGraph);
    final TestingStateHandleStore<JobGraph> stateHandleStore = builder.setGetFunction(ignore -> stateHandle).setAddFunction((ignore, state) -> jobGraphStorageHelper.store(state)).build();
    final JobGraphStore jobGraphStore = createAndStartJobGraphStore(stateHandleStore);
    jobGraphStore.recoverJobGraph(testingJobGraph.getJobID());
    // Known recovered job
    testingJobGraphStoreWatcher.addJobGraph(testingJobGraph.getJobID());
    // Unknown job
    final JobID unknownJobId = JobID.generate();
    testingJobGraphStoreWatcher.addJobGraph(unknownJobId);
    assertThat(testingJobGraphListener.getAddedJobGraphs().size(), is(1));
    assertThat(testingJobGraphListener.getAddedJobGraphs(), contains(unknownJobId));
}