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

Example 6 with RetrievableStateHandle

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

the class ZooKeeperStateHandleStoreITCase method testAdd.

/**
	 * Tests add operation with default {@link CreateMode}.
	 */
@Test
public void testAdd() throws Exception {
    LongStateStorage longStateStorage = new LongStateStorage();
    ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<Long>(ZooKeeper.getClient(), longStateStorage, Executors.directExecutor());
    // Config
    final String pathInZooKeeper = "/testAdd";
    final Long state = 1239712317L;
    // Test
    store.add(pathInZooKeeper, state);
    // Verify
    // State handle created
    assertEquals(1, store.getAll().size());
    assertEquals(state, store.get(pathInZooKeeper).retrieveState());
    // Path created and is persistent
    Stat stat = ZooKeeper.getClient().checkExists().forPath(pathInZooKeeper);
    assertNotNull(stat);
    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 7 with RetrievableStateHandle

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

the class ZooKeeperStateHandleStoreITCase method testReplaceDiscardStateHandleAfterFailure.

/**
	 * Tests that the replace state handle is discarded if ZooKeeper setData fails.
	 */
@Test
public void testReplaceDiscardStateHandleAfterFailure() throws Exception {
    // Setup
    LongStateStorage stateHandleProvider = new LongStateStorage();
    CuratorFramework client = spy(ZooKeeper.getClient());
    when(client.setData()).thenThrow(new RuntimeException("Expected test Exception."));
    ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(client, stateHandleProvider, Executors.directExecutor());
    // Config
    final String pathInZooKeeper = "/testReplaceDiscardStateHandleAfterFailure";
    final Long initialState = 30968470898L;
    final Long replaceState = 88383776661L;
    // Test
    store.add(pathInZooKeeper, initialState);
    try {
        store.replace(pathInZooKeeper, 0, replaceState);
        fail("Did not throw expected exception");
    } catch (Exception ignored) {
    }
    // Verify
    // State handle created and discarded
    assertEquals(2, stateHandleProvider.getStateHandles().size());
    assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).retrieveState());
    assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).retrieveState());
    assertEquals(1, stateHandleProvider.getStateHandles().get(1).getNumberOfDiscardCalls());
    // Initial value
    @SuppressWarnings("unchecked") Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(ZooKeeper.getClient().getData().forPath(pathInZooKeeper), ClassLoader.getSystemClassLoader())).retrieveState();
    assertEquals(initialState, actual);
}

Example 8 with RetrievableStateHandle

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

the class ZooKeeperStateHandleStoreITCase method testReplace.

/**
	 * Tests that a state handle is replaced.
	 */
@Test
public void testReplace() throws Exception {
    // Setup
    LongStateStorage stateHandleProvider = new LongStateStorage();
    ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(ZooKeeper.getClient(), stateHandleProvider, Executors.directExecutor());
    // Config
    final String pathInZooKeeper = "/testReplace";
    final Long initialState = 30968470898L;
    final Long replaceState = 88383776661L;
    // Test
    store.add(pathInZooKeeper, initialState);
    store.replace(pathInZooKeeper, 0, replaceState);
    // Verify
    // State handles created
    assertEquals(2, stateHandleProvider.getStateHandles().size());
    assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).retrieveState());
    assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).retrieveState());
    // Path created and is persistent
    Stat stat = ZooKeeper.getClient().checkExists().forPath(pathInZooKeeper);
    assertNotNull(stat);
    assertEquals(0, stat.getEphemeralOwner());
    // Data is equal
    @SuppressWarnings("unchecked") Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(ZooKeeper.getClient().getData().forPath(pathInZooKeeper), ClassLoader.getSystemClassLoader())).retrieveState();
    assertEquals(replaceState, actual);
}

Example 9 with RetrievableStateHandle

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

the class KubernetesStateHandleStore method addAndLock.

/**
 * Creates a state handle, stores it in ConfigMap. We could guarantee that only the leader could
 * update the ConfigMap. Since “Get(check the leader)-and-Update(write back to the ConfigMap)”
 * is a transactional operation.
 *
 * @param key Key in ConfigMap
 * @param state State to be added
 * @throws AlreadyExistException if the name already exists
 * @throws PossibleInconsistentStateException if the write-to-Kubernetes operation failed. This
 *     indicates that it's not clear whether the new state was successfully written to
 *     Kubernetes or not. No state was discarded. Proper error handling has to be applied on the
 *     caller's side.
 * @throws Exception if persisting state or writing state handle failed
 */
@Override
public RetrievableStateHandle<T> addAndLock(String key, T state) throws PossibleInconsistentStateException, Exception {
    checkNotNull(key, "Key in ConfigMap.");
    checkNotNull(state, "State.");
    final RetrievableStateHandle<T> storeHandle = storage.store(state);
    final byte[] serializedStoreHandle = serializeOrDiscard(storeHandle);
    // initialize flag to serve the failure case
    boolean discardState = true;
    try {
        // a successful operation will result in the state not being discarded
        discardState = !kubeClient.checkAndUpdateConfigMap(configMapName, c -> {
            if (isValidOperation(c)) {
                if (!c.getData().containsKey(key)) {
                    c.getData().put(key, encodeStateHandle(serializedStoreHandle));
                    return Optional.of(c);
                } else {
                    throw new CompletionException(getKeyAlreadyExistException(key));
                }
            }
            return Optional.empty();
        }).get();
        return storeHandle;
    } catch (Exception ex) {
        final Optional<PossibleInconsistentStateException> possibleInconsistentStateException = ExceptionUtils.findThrowable(ex, PossibleInconsistentStateException.class);
        if (possibleInconsistentStateException.isPresent()) {
            // it's unclear whether the state handle metadata was written to the ConfigMap -
            // hence, we don't discard the data
            discardState = false;
            throw possibleInconsistentStateException.get();
        }
        throw ExceptionUtils.findThrowable(ex, AlreadyExistException.class).orElseThrow(() -> ex);
    } finally {
        if (discardState) {
            storeHandle.discardState();
        }
    }
}

Example 10 with RetrievableStateHandle

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

the class KubernetesStateHandleStore method replace.

/**
 * Replaces a state handle in ConfigMap and discards the old state handle. Wo do not lock
 * resource version and then replace in Kubernetes. Since the ConfigMap is periodically updated
 * by leader, the resource version changes very fast. We use a "check-existence and update"
 * transactional operation instead.
 *
 * @param key Key in ConfigMap
 * @param resourceVersion resource version when checking existence via {@link #exists}.
 * @param state State to be added
 * @throws NotExistException if the name does not exist
 * @throws PossibleInconsistentStateException if a failure occurred during the update operation.
 *     It's unclear whether the operation actually succeeded or not. No state was discarded. The
 *     method's caller should handle this case properly.
 * @throws Exception if persisting state or writing state handle failed
 */
@Override
public void replace(String key, StringResourceVersion resourceVersion, T state) throws Exception {
    checkNotNull(key, "Key in ConfigMap.");
    checkNotNull(state, "State.");
    final RetrievableStateHandle<T> oldStateHandle = getAndLock(key);
    final RetrievableStateHandle<T> newStateHandle = storage.store(state);
    final byte[] serializedStateHandle = serializeOrDiscard(newStateHandle);
    // initialize flags to serve the failure case
    boolean discardOldState = false;
    boolean discardNewState = true;
    try {
        boolean success = kubeClient.checkAndUpdateConfigMap(configMapName, c -> {
            if (isValidOperation(c)) {
                // Check the existence
                if (c.getData().containsKey(key)) {
                    c.getData().put(key, encodeStateHandle(serializedStateHandle));
                } else {
                    throw new CompletionException(getKeyNotExistException(key));
                }
                return Optional.of(c);
            }
            return Optional.empty();
        }).get();
        // swap subject for deletion in case of success
        discardOldState = success;
        discardNewState = !success;
    } catch (Exception ex) {
        final Optional<PossibleInconsistentStateException> possibleInconsistentStateException = ExceptionUtils.findThrowable(ex, PossibleInconsistentStateException.class);
        if (possibleInconsistentStateException.isPresent()) {
            // it's unclear whether the state handle metadata was written to the ConfigMap -
            // hence, we don't discard any data
            discardNewState = false;
            throw possibleInconsistentStateException.get();
        }
        throw ExceptionUtils.findThrowable(ex, NotExistException.class).orElseThrow(() -> ex);
    } finally {
        if (discardNewState) {
            newStateHandle.discardState();
        }
        if (discardOldState) {
            oldStateHandle.discardState();
        }
    }
}