org.apache.flink.runtime.state.AbstractKeyedStateBackend Java Examples

The following examples show how to use org.apache.flink.runtime.state.AbstractKeyedStateBackend. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: MultiStateKeyIteratorTest.java    From flink with Apache License 2.0 6 votes vote down vote up
@Test
public void testIteratorRemovesFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 1);

	MultiStateKeyIterator<Integer> iterator =
		new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	int key = iterator.next();
	Assert.assertEquals("Unexpected keys pulled from state backend", 1, key);

	iterator.remove();
	Assert.assertFalse("Failed to drop key from all descriptors in state backend", iterator.hasNext());

	for (StateDescriptor<?, ?> descriptor : descriptors) {
		Assert.assertEquals(
			"Failed to drop key for state descriptor",
			0,
			keyedStateBackend.getKeys(descriptor.getName(), VoidNamespace.INSTANCE).count());
	}
}
 
Example #2
Source File: MultiStateKeyIteratorTest.java    From flink with Apache License 2.0 6 votes vote down vote up
@Test
public void testIteratorPullsKeyFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 2);

	MultiStateKeyIterator<Integer> iterator = new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	List<Integer> keys = new ArrayList<>();

	while (iterator.hasNext()) {
		keys.add(iterator.next());
	}

	Assert.assertEquals("Unexpected number of keys", 2, keys.size());
	Assert.assertEquals("Unexpected keys found", Arrays.asList(1, 2), keys);
}
 
Example #3
Source File: FlinkStateInternalsTest.java    From beam with Apache License 2.0 6 votes vote down vote up
public static KeyedStateBackend<ByteBuffer> createStateBackend() throws Exception {
  MemoryStateBackend backend = new MemoryStateBackend();

  AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend =
      backend.createKeyedStateBackend(
          new DummyEnvironment("test", 1, 0),
          new JobID(),
          "test_op",
          new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
          2,
          new KeyGroupRange(0, 1),
          new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
          TtlTimeProvider.DEFAULT,
          null,
          Collections.emptyList(),
          new CloseableRegistry());

  changeKey(keyedStateBackend);

  return keyedStateBackend;
}
 
Example #4
Source File: MockStateBackend.java    From flink with Apache License 2.0 6 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 
Example #5
Source File: TestSpyWrapperStateBackend.java    From flink with Apache License 2.0 6 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws IOException {
	return spy(delegate.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry));
}
 
Example #6
Source File: MultiStateKeyIteratorTest.java    From flink with Apache License 2.0 6 votes vote down vote up
private static AbstractKeyedStateBackend<Integer> createKeyedStateBackend() {
	MockStateBackend backend = new MockStateBackend();

	return backend.createKeyedStateBackend(
		new DummyEnvironment(),
		new JobID(),
		"mock-backend",
		IntSerializer.INSTANCE,
		129,
		KeyGroupRange.of(0, 128),
		null,
		TtlTimeProvider.DEFAULT,
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
		Collections.emptyList(),
		new CloseableRegistry());
}
 
Example #7
Source File: MockStateBackend.java    From Flink-CEPplus with Apache License 2.0 6 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 
Example #8
Source File: TestSpyWrapperStateBackend.java    From Flink-CEPplus with Apache License 2.0 6 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws IOException {
	return spy(delegate.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry));
}
 
Example #9
Source File: TestSpyWrapperStateBackend.java    From flink with Apache License 2.0 6 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws IOException {
	return spy(delegate.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry));
}
 
Example #10
Source File: MultiStateKeyIteratorTest.java    From flink with Apache License 2.0 6 votes vote down vote up
@Test
public void testIteratorPullsKeyFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 2);

	MultiStateKeyIterator<Integer> iterator = new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	List<Integer> keys = new ArrayList<>();

	while (iterator.hasNext()) {
		keys.add(iterator.next());
	}

	Assert.assertEquals("Unexpected number of keys", 2, keys.size());
	Assert.assertEquals("Unexpected keys found", Arrays.asList(1, 2), keys);
}
 
Example #11
Source File: MockStateBackend.java    From flink with Apache License 2.0 6 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 
Example #12
Source File: MultiStateKeyIteratorTest.java    From flink with Apache License 2.0 6 votes vote down vote up
@Test
public void testIteratorRemovesFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 1);

	MultiStateKeyIterator<Integer> iterator =
		new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	int key = iterator.next();
	Assert.assertEquals("Unexpected keys pulled from state backend", 1, key);

	iterator.remove();
	Assert.assertFalse("Failed to drop key from all descriptors in state backend", iterator.hasNext());

	for (StateDescriptor<?, ?> descriptor : descriptors) {
		Assert.assertEquals(
			"Failed to drop key for state descriptor",
			0,
			keyedStateBackend.getKeys(descriptor.getName(), VoidNamespace.INSTANCE).count());
	}
}
 
Example #13
Source File: StreamTaskStateInitializerImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
	AbstractKeyedStateBackend<K> keyedStatedBackend,
	KeyContext keyContext, //the operator
	Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {

	if (keyedStatedBackend == null) {
		return null;
	}

	final KeyGroupRange keyGroupRange = keyedStatedBackend.getKeyGroupRange();

	final InternalTimeServiceManager<K> timeServiceManager = new InternalTimeServiceManager<>(
		keyGroupRange,
		keyContext,
		keyedStatedBackend,
		processingTimeService,
		keyedStatedBackend.requiresLegacySynchronousTimerSnapshots());

	// and then initialize the timer services
	for (KeyGroupStatePartitionStreamProvider streamProvider : rawKeyedStates) {
		int keyGroupIdx = streamProvider.getKeyGroupId();

		Preconditions.checkArgument(keyGroupRange.contains(keyGroupIdx),
			"Key Group " + keyGroupIdx + " does not belong to the local range.");

		timeServiceManager.restoreStateForKeyGroup(
			streamProvider.getStream(),
			keyGroupIdx, environment.getUserClassLoader());
	}

	return timeServiceManager;
}
 
Example #14
Source File: MemoryStateBackend.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		taskStateManager.createLocalRecoveryConfig(),
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 
Example #15
Source File: StubStateBackend.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	return backend.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		this.ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry);
}
 
Example #16
Source File: StreamingRuntimeContextTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
	ExecutionConfig config = new ExecutionConfig();

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<ListState<String>>() {

		@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()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	return operatorMock;
}
 
Example #17
Source File: StreamingRuntimeContextTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
	ExecutionConfig config = new ExecutionConfig();

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<ListState<String>>() {

		@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()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	return operatorMock;
}
 
Example #18
Source File: StreamTaskStateInitializerImplTest.java    From flink with Apache License 2.0 5 votes vote down vote up
private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				ProcessingTimeService processingTimeService,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 
Example #19
Source File: StreamTaskTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public void initializeState(StateInitializationContext context) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) context.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) context.getRawKeyedStateInputs();
	super.initializeState(context);
}
 
Example #20
Source File: RocksDBStateBackendTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

		ValueState<String> state =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}
 
Example #21
Source File: StreamingRuntimeContextTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
	ExecutionConfig config = new ExecutionConfig();

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<MapState<Integer, String>>() {

		@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()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	return operatorMock;
}
 
Example #22
Source File: StreamTaskStateInitializerImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
StreamOperatorStateContextImpl(
	boolean restored,
	OperatorStateBackend operatorStateBackend,
	AbstractKeyedStateBackend<?> keyedStateBackend,
	InternalTimeServiceManager<?> internalTimeServiceManager,
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs,
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs) {

	this.restored = restored;
	this.operatorStateBackend = operatorStateBackend;
	this.keyedStateBackend = keyedStateBackend;
	this.internalTimeServiceManager = internalTimeServiceManager;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 
Example #23
Source File: AbstractStreamOperator.java    From flink with Apache License 2.0 5 votes vote down vote up
@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);
		}
	}
}
 
Example #24
Source File: CheckpointSettingsSerializableTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	throw new UnsupportedOperationException();
}
 
Example #25
Source File: RocksDBStateBackendTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

		ValueState<String> state =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}
 
Example #26
Source File: MultiStateKeyIteratorTest.java    From flink with Apache License 2.0 5 votes vote down vote up
private static void setKey(
	AbstractKeyedStateBackend<Integer> backend,
	ValueStateDescriptor<Integer> descriptor,
	int key) throws Exception {
	backend.setCurrentKey(key);
	backend
		.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descriptor)
		.update(0);
}
 
Example #27
Source File: StreamTaskTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public void initializeState(StateInitializationContext controller) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) controller.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) controller.getRawKeyedStateInputs();
	super.initializeState(controller);
}
 
Example #28
Source File: KeyedStateInputFormat.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final Environment environment = new SavepointEnvironment
		.Builder(getRuntimeContext(), split.getNumKeyGroups())
		.setSubtaskIndex(split.getSplitNumber())
		.setPrioritizedOperatorSubtaskState(split.getPrioritizedOperatorSubtaskState())
		.build();

	final StreamOperatorStateContext context = getStreamOperatorStateContext(environment);

	keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();

	final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
	SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);
	FunctionUtils.setFunctionRuntimeContext(userFunction, ctx);

	keys = getKeyIterator(ctx);

	final InternalTimerService<VoidNamespace> timerService = restoreTimerService(context);
	try {
		this.ctx = new Context(keyedStateBackend, timerService);
	} catch (Exception e) {
		throw new IOException("Failed to restore timer state", e);
	}
}
 
Example #29
Source File: StubStateBackend.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	return backend.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		this.ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry);
}
 
Example #30
Source File: StreamTaskStateInitializerImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
	AbstractKeyedStateBackend<K> keyedStatedBackend,
	KeyContext keyContext, //the operator
	ProcessingTimeService processingTimeService,
	Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {

	if (keyedStatedBackend == null) {
		return null;
	}

	final KeyGroupRange keyGroupRange = keyedStatedBackend.getKeyGroupRange();

	final InternalTimeServiceManager<K> timeServiceManager = new InternalTimeServiceManager<>(
		keyGroupRange,
		keyContext,
		keyedStatedBackend,
		processingTimeService,
		keyedStatedBackend.requiresLegacySynchronousTimerSnapshots());

	// and then initialize the timer services
	for (KeyGroupStatePartitionStreamProvider streamProvider : rawKeyedStates) {
		int keyGroupIdx = streamProvider.getKeyGroupId();

		Preconditions.checkArgument(keyGroupRange.contains(keyGroupIdx),
			"Key Group " + keyGroupIdx + " does not belong to the local range.");

		timeServiceManager.restoreStateForKeyGroup(
			streamProvider.getStream(),
			keyGroupIdx, environment.getUserClassLoader());
	}

	return timeServiceManager;
}