@Override public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable { ListStateDescriptor<String> descr = (ListStateDescriptor<String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));
protected <N, S extends State, T> S getOrCreateKeyedState( TypeSerializer<N> namespaceSerializer, StateDescriptor<S, T> stateDescriptor) throws Exception { if (keyedStateStore != null) { return keyedStateBackend.getOrCreateKeyedState(namespaceSerializer, stateDescriptor); } else { throw new IllegalStateException("Cannot create partitioned state. " + "The keyed state backend has not been set." + "This indicates that the operator is not partitioned/keyed."); } }
keyedStateBackend.getKeyGroupRange() : KeyGroupRange.EMPTY_KEY_GROUP_RANGE; keyedStateBackend.snapshot(checkpointId, timestamp, factory, checkpointOptions));
/** * @see KeyedStateBackend */ @Override public <N, S extends State, T> void applyToAllKeys( final N namespace, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, T> stateDescriptor, final KeyedStateFunction<K, S> function) throws Exception { try (Stream<K> keyStream = getKeys(stateDescriptor.getName(), namespace)) { final S state = getPartitionedState( namespace, namespaceSerializer, stateDescriptor); keyStream.forEach((K key) -> { setCurrentKey(key); try { function.process(key, state); } catch (Throwable e) { // we wrap the checked exception in an unchecked // one and catch it (and re-throw it) later. throw new RuntimeException(e); } }); } }
@SuppressWarnings({"unchecked", "rawtypes"}) public void setCurrentKey(Object key) { if (keyedStateBackend != null) { try { // need to work around type restrictions @SuppressWarnings("unchecked,rawtypes") AbstractKeyedStateBackend rawBackend = (AbstractKeyedStateBackend) keyedStateBackend; rawBackend.setCurrentKey(key); } catch (Exception e) { throw new RuntimeException("Exception occurred while setting the current key context.", e); } } }
keyedStatedBackend.dispose();
return keyedStateBackend.getPartitionedState(namespace, namespaceSerializer, stateDescriptor); } else { throw new RuntimeException("Cannot create partitioned state. The keyed state " +
public AbstractKeyedStateBackend( TaskKvStateRegistry kvStateRegistry, TypeSerializer<K> keySerializer, ClassLoader userCodeClassLoader, int numberOfKeyGroups, KeyGroupRange keyGroupRange, ExecutionConfig executionConfig) { this.kvStateRegistry = kvStateRegistry; this.keySerializer = Preconditions.checkNotNull(keySerializer); this.numberOfKeyGroups = Preconditions.checkNotNull(numberOfKeyGroups); this.userCodeClassLoader = Preconditions.checkNotNull(userCodeClassLoader); this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange); this.cancelStreamRegistry = new CloseableRegistry(); this.keyValueStatesByName = new HashMap<>(); this.executionConfig = executionConfig; this.keyGroupCompressionDecorator = determineStreamCompression(executionConfig); }
/** * @see KeyedStateBackend */ @Override public <N, S extends State, T> void applyToAllKeys( final N namespace, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, T> stateDescriptor, final KeyedStateFunction<K, S> function) throws Exception { try (Stream<K> keyStream = getKeys(stateDescriptor.getName(), namespace)) { final S state = getPartitionedState( namespace, namespaceSerializer, stateDescriptor); keyStream.forEach((K key) -> { setCurrentKey(key); try { function.process(key, state); } catch (Throwable e) { // we wrap the checked exception in an unchecked // one and catch it (and re-throw it) later. throw new RuntimeException(e); } }); } }
@Override public void setCurrentKey(K newKey) { super.setCurrentKey(newKey); sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(), getCurrentKeyGroupIndex()); }
super.dispose();
return keyedStateBackend.getPartitionedState(namespace, namespaceSerializer, stateDescriptor); } else { throw new RuntimeException("Cannot create partitioned state. The keyed state " +
keyedStateBackend.getKeyGroupRange() : KeyGroupRange.EMPTY_KEY_GROUP_RANGE; keyedStateBackend.snapshot(checkpointId, timestamp, factory, checkpointOptions));
public AbstractKeyedStateBackend( TaskKvStateRegistry kvStateRegistry, TypeSerializer<K> keySerializer, ClassLoader userCodeClassLoader, int numberOfKeyGroups, KeyGroupRange keyGroupRange, ExecutionConfig executionConfig, TtlTimeProvider ttlTimeProvider) { Preconditions.checkArgument(numberOfKeyGroups >= 1, "NumberOfKeyGroups must be a positive number"); Preconditions.checkArgument(numberOfKeyGroups >= keyGroupRange.getNumberOfKeyGroups(), "The total number of key groups must be at least the number in the key group range assigned to this backend"); this.kvStateRegistry = kvStateRegistry; this.keySerializer = Preconditions.checkNotNull(keySerializer); this.numberOfKeyGroups = numberOfKeyGroups; this.userCodeClassLoader = Preconditions.checkNotNull(userCodeClassLoader); this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange); this.cancelStreamRegistry = new CloseableRegistry(); this.keyValueStatesByName = new HashMap<>(); this.executionConfig = executionConfig; this.keyGroupCompressionDecorator = determineStreamCompression(executionConfig); this.ttlTimeProvider = Preconditions.checkNotNull(ttlTimeProvider); }
@Override public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable { MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));
/** * @see KeyedStateBackend */ @Override public <N, S extends State, T> void applyToAllKeys( final N namespace, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, T> stateDescriptor, final KeyedStateFunction<K, S> function) throws Exception { try (Stream<K> keyStream = getKeys(stateDescriptor.getName(), namespace)) { final S state = getPartitionedState( namespace, namespaceSerializer, stateDescriptor); keyStream.forEach((K key) -> { setCurrentKey(key); try { function.process(key, state); } catch (Throwable e) { // we wrap the checked exception in an unchecked // one and catch it (and re-throw it) later. throw new RuntimeException(e); } }); } catch (RuntimeException e) { throw e; } }
@SuppressWarnings({"unchecked", "rawtypes"}) public void setCurrentKey(Object key) { if (keyedStateBackend != null) { try { // need to work around type restrictions @SuppressWarnings("unchecked,rawtypes") AbstractKeyedStateBackend rawBackend = (AbstractKeyedStateBackend) keyedStateBackend; rawBackend.setCurrentKey(key); } catch (Exception e) { throw new RuntimeException("Exception occurred while setting the current key context.", e); } } }
@Test public void testStateBackendLoadingAndClosing() throws Exception { Configuration taskManagerConfig = new Configuration(); taskManagerConfig.setString(CheckpointingOptions.STATE_BACKEND, TestMemoryStateBackendFactory.class.getName()); StreamConfig cfg = new StreamConfig(new Configuration()); cfg.setStateKeySerializer(mock(TypeSerializer.class)); cfg.setOperatorID(new OperatorID(4711L, 42L)); TestStreamSource<Long, MockSourceFunction> streamSource = new TestStreamSource<>(new MockSourceFunction()); cfg.setStreamOperator(streamSource); cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime); Task task = createTask(StateBackendTestSource.class, cfg, taskManagerConfig); StateBackendTestSource.fail = false; task.startTaskThread(); // wait for clean termination task.getExecutingThread().join(); // ensure that the state backends and stream iterables are closed ... verify(TestStreamSource.operatorStateBackend).close(); verify(TestStreamSource.keyedStateBackend).close(); verify(TestStreamSource.rawOperatorStateInputs).close(); verify(TestStreamSource.rawKeyedStateInputs).close(); // ... and disposed verify(TestStreamSource.operatorStateBackend).dispose(); verify(TestStreamSource.keyedStateBackend).dispose(); assertEquals(ExecutionState.FINISHED, task.getExecutionState()); }