org.apache.flink.runtime.query.KvStateRegistry Java Examples

@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();
}
 

/**
 * Tests that the channel is closed if an Exception reaches the channel handler.
 */
@Test
public void testCloseChannelOnExceptionCaught() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(handler);

	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable response = MessageSerializer.deserializeServerFailure(buf);

	assertTrue(response.getMessage().contains("Expected test Exception"));

	channel.closeFuture().await(READ_TIMEOUT_MILLIS);
	assertFalse(channel.isActive());
}
 

private SavepointEnvironment(RuntimeContext ctx, Configuration configuration, int maxParallelism, int indexOfSubtask, PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	this.jobID = new JobID();
	this.vertexID = new JobVertexID();
	this.attemptID = new ExecutionAttemptID();
	this.ctx = Preconditions.checkNotNull(ctx);
	this.configuration = Preconditions.checkNotNull(configuration);

	Preconditions.checkArgument(maxParallelism > 0 && indexOfSubtask < maxParallelism);
	this.maxParallelism = maxParallelism;
	this.indexOfSubtask = indexOfSubtask;

	this.registry = new KvStateRegistry().createTaskRegistry(jobID, vertexID);
	this.taskStateManager = new SavepointTaskStateManager(prioritizedOperatorSubtaskState);
	this.ioManager = new IOManagerAsync(ConfigurationUtils.parseTempDirectories(configuration));
	this.memoryManager = MemoryManager.forDefaultPageSize(64 * 1024 * 1024);
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, attemptID);
}
 

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;
}
 

private SavepointEnvironment(RuntimeContext ctx, Configuration configuration, int maxParallelism, int indexOfSubtask, PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	this.jobID = new JobID();
	this.vertexID = new JobVertexID();
	this.attemptID = new ExecutionAttemptID();
	this.ctx = Preconditions.checkNotNull(ctx);
	this.configuration = Preconditions.checkNotNull(configuration);

	Preconditions.checkArgument(maxParallelism > 0 && indexOfSubtask < maxParallelism);
	this.maxParallelism = maxParallelism;
	this.indexOfSubtask = indexOfSubtask;

	this.registry = new KvStateRegistry().createTaskRegistry(jobID, vertexID);
	this.taskStateManager = new SavepointTaskStateManager(prioritizedOperatorSubtaskState);
	this.ioManager = new IOManagerAsync();
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, attemptID);
}
 

@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();
}
 

private void registerQueryableState(JobID jobId, JobMasterGateway jobMasterGateway) {
	final KvStateServer kvStateServer = kvStateService.getKvStateServer();
	final KvStateRegistry kvStateRegistry = kvStateService.getKvStateRegistry();

	if (kvStateServer != null && kvStateRegistry != null) {
		kvStateRegistry.registerListener(
			jobId,
			new RpcKvStateRegistryListener(
				jobMasterGateway,
				kvStateServer.getServerAddress()));
	}

	final KvStateClientProxy kvStateProxy = kvStateService.getKvStateClientProxy();

	if (kvStateProxy != null) {
		kvStateProxy.updateKvStateLocationOracle(jobId, jobMasterGateway);
	}
}
 

private void disassociateFromJobManager(JobManagerConnection jobManagerConnection, Exception cause) throws IOException {
	checkNotNull(jobManagerConnection);

	final JobID jobId = jobManagerConnection.getJobID();

	// cleanup remaining partitions once all tasks for this job have completed
	scheduleResultPartitionCleanup(jobId);

	final KvStateRegistry kvStateRegistry = kvStateService.getKvStateRegistry();

	if (kvStateRegistry != null) {
		kvStateRegistry.unregisterListener(jobId);
	}

	final KvStateClientProxy kvStateClientProxy = kvStateService.getKvStateClientProxy();

	if (kvStateClientProxy != null) {
		kvStateClientProxy.updateKvStateLocationOracle(jobManagerConnection.getJobID(), null);
	}

	JobMasterGateway jobManagerGateway = jobManagerConnection.getJobManagerGateway();
	jobManagerGateway.disconnectTaskManager(getResourceID(), cause);
	jobManagerConnection.getLibraryCacheManager().shutdown();
}
 

private void registerQueryableState(JobID jobId, JobMasterGateway jobMasterGateway) {
	final KvStateServer kvStateServer = kvStateService.getKvStateServer();
	final KvStateRegistry kvStateRegistry = kvStateService.getKvStateRegistry();

	if (kvStateServer != null && kvStateRegistry != null) {
		kvStateRegistry.registerListener(
			jobId,
			new RpcKvStateRegistryListener(
				jobMasterGateway,
				kvStateServer.getServerAddress()));
	}

	final KvStateClientProxy kvStateProxy = kvStateService.getKvStateClientProxy();

	if (kvStateProxy != null) {
		kvStateProxy.updateKvStateLocationOracle(jobId, jobMasterGateway);
	}
}
 

private void registerQueryableState(JobID jobId, JobMasterGateway jobMasterGateway) {
	final KvStateServer kvStateServer = networkEnvironment.getKvStateServer();
	final KvStateRegistry kvStateRegistry = networkEnvironment.getKvStateRegistry();

	if (kvStateServer != null && kvStateRegistry != null) {
		kvStateRegistry.registerListener(
			jobId,
			new RpcKvStateRegistryListener(
				jobMasterGateway,
				kvStateServer.getServerAddress()));
	}

	final KvStateClientProxy kvStateProxy = networkEnvironment.getKvStateProxy();

	if (kvStateProxy != null) {
		kvStateProxy.updateKvStateLocationOracle(jobId, jobMasterGateway);
	}
}
 

@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();
}
 

static Environment getMockEnvironment(File... tempDirs) {
	final String[] tempDirStrings = new String[tempDirs.length];
	for (int i = 0; i < tempDirs.length; i++) {
		tempDirStrings[i] = tempDirs[i].getAbsolutePath();
	}

	IOManager ioMan = mock(IOManager.class);
	when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);

	Environment env = mock(Environment.class);
	when(env.getJobID()).thenReturn(new JobID());
	when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader());
	when(env.getIOManager()).thenReturn(ioMan);
	when(env.getTaskKvStateRegistry()).thenReturn(new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));

	TaskInfo taskInfo = mock(TaskInfo.class);
	when(env.getTaskInfo()).thenReturn(taskInfo);
	when(taskInfo.getIndexOfThisSubtask()).thenReturn(0);

	TaskManagerRuntimeInfo tmInfo = new TestingTaskManagerRuntimeInfo(new Configuration(), tempDirStrings);
	when(env.getTaskManagerInfo()).thenReturn(tmInfo);

	TestTaskStateManager taskStateManager = new TestTaskStateManager();
	when(env.getTaskStateManager()).thenReturn(taskStateManager);

	return env;
}
 

@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;
}
 

@BeforeClass
public static void setup() {
	try {
		testServer = new KvStateServerImpl(
				InetAddress.getLocalHost().getHostName(),
				Collections.singletonList(0).iterator(),
				1,
				1,
				new KvStateRegistry(),
				new DisabledKvStateRequestStats());
		testServer.start();
	} catch (Throwable e) {
		e.printStackTrace();
	}
}
 

@Test
public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
	final File targetDir1 = tempFolder.newFolder();
	final File targetDir2 = tempFolder.newFolder();
	Assume.assumeTrue("Cannot mark directory non-writable", targetDir1.setWritable(false, false));

	String checkpointPath = tempFolder.newFolder().toURI().toString();
	RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

	try (MockEnvironment env = getMockEnvironment(tempFolder.newFolder())) {
		rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());

		try {
			AbstractKeyedStateBackend<Integer> keyedStateBackend = rocksDbBackend.createKeyedStateBackend(
				env,
				env.getJobID(),
				"foobar",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());

			IOUtils.closeQuietly(keyedStateBackend);
			keyedStateBackend.dispose();
		}
		catch (Exception e) {
			e.printStackTrace();
			fail("Backend initialization failed even though some paths were available");
		}
	} finally {
		//noinspection ResultOfMethodCallIgnored
		targetDir1.setWritable(true, false);
	}
}
 

@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;
}
 

public NetworkEnvironment(
	NetworkBufferPool networkBufferPool,
	ConnectionManager connectionManager,
	ResultPartitionManager resultPartitionManager,
	TaskEventDispatcher taskEventDispatcher,
	KvStateRegistry kvStateRegistry,
	KvStateServer kvStateServer,
	KvStateClientProxy kvStateClientProxy,
	IOMode defaultIOMode,
	int partitionRequestInitialBackoff,
	int partitionRequestMaxBackoff,
	int networkBuffersPerChannel,
	int extraNetworkBuffersPerGate,
	boolean enableCreditBased) {

	this.networkBufferPool = checkNotNull(networkBufferPool);
	this.connectionManager = checkNotNull(connectionManager);
	this.resultPartitionManager = checkNotNull(resultPartitionManager);
	this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
	this.kvStateRegistry = checkNotNull(kvStateRegistry);

	this.kvStateServer = kvStateServer;
	this.kvStateProxy = kvStateClientProxy;

	this.defaultIOMode = defaultIOMode;

	this.partitionRequestInitialBackoff = partitionRequestInitialBackoff;
	this.partitionRequestMaxBackoff = partitionRequestMaxBackoff;

	isShutdown = false;
	this.networkBuffersPerChannel = networkBuffersPerChannel;
	this.extraNetworkBuffersPerGate = extraNetworkBuffersPerGate;

	this.enableCreditBased = enableCreditBased;
}
 

@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;
}
 

/**
 * Create the handler used by the {@link KvStateServerImpl}.
 *
 * @param server the {@link KvStateServerImpl} using the handler.
 * @param kvStateRegistry registry to query.
 * @param serializer the {@link MessageSerializer} used to (de-) serialize the different messages.
 * @param stats server statistics collector.
 */
public KvStateServerHandler(
		final KvStateServerImpl server,
		final KvStateRegistry kvStateRegistry,
		final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer,
		final KvStateRequestStats stats) {

	super(server, serializer, stats);
	this.registry = Preconditions.checkNotNull(kvStateRegistry);
}
 

private AbstractKeyedStateBackend<Integer> createKeyedStateBackend(KvStateRegistry registry, int numKeyGroups, AbstractStateBackend abstractBackend, DummyEnvironment dummyEnv) throws java.io.IOException {
	return abstractBackend.createKeyedStateBackend(
		dummyEnv,
		dummyEnv.getJobID(),
		"test_op",
		IntSerializer.INSTANCE,
		numKeyGroups,
		new KeyGroupRange(0, 0),
		registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());
}
 

/**
 * Tests that incoming buffer instances are recycled.
 */
@Test
public void testIncomingBufferIsRecycled() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), 282872L, request);

	assertEquals(1L, serRequest.refCnt());

	// Write regular request
	channel.writeInbound(serRequest);
	assertEquals("Buffer not recycled", 0L, serRequest.refCnt());

	// Write unexpected msg
	ByteBuf unexpected = channel.alloc().buffer(8);
	unexpected.writeInt(4);
	unexpected.writeInt(4);

	assertEquals(1L, unexpected.refCnt());

	channel.writeInbound(unexpected);
	assertEquals("Buffer not recycled", 0L, unexpected.refCnt());
}
 

/**
 * Tests the failure response with {@link UnknownKvStateIdException} as cause on
 * queries for unregistered KvStateIDs.
 */
@Test
public void testQueryUnknownKvStateID() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	long requestId = Integer.MAX_VALUE + 182828L;

	KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);

	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.REQUEST_FAILURE, MessageSerializer.deserializeHeader(buf));
	RequestFailure response = MessageSerializer.deserializeRequestFailure(buf);
	buf.release();

	assertEquals(requestId, response.getRequestId());

	assertTrue("Did not respond with expected failure cause", response.getCause() instanceof UnknownKvStateIdException);

	assertEquals(1L, stats.getNumRequests());
	assertEquals(1L, stats.getNumFailed());
}
 

private AbstractKeyedStateBackend<Integer> createKeyedStateBackend(KvStateRegistry registry, int numKeyGroups, AbstractStateBackend abstractBackend, DummyEnvironment dummyEnv) throws java.io.IOException {
	return abstractBackend.createKeyedStateBackend(
		dummyEnv,
		dummyEnv.getJobID(),
		"test_op",
		IntSerializer.INSTANCE,
		numKeyGroups,
		new KeyGroupRange(0, 0),
		registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());
}
 

/**
 * Tests that incoming buffer instances are recycled.
 */
@Test
public void testIncomingBufferIsRecycled() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), 282872L, request);

	assertEquals(1L, serRequest.refCnt());

	// Write regular request
	channel.writeInbound(serRequest);
	assertEquals("Buffer not recycled", 0L, serRequest.refCnt());

	// Write unexpected msg
	ByteBuf unexpected = channel.alloc().buffer(8);
	unexpected.writeInt(4);
	unexpected.writeInt(4);

	assertEquals(1L, unexpected.refCnt());

	channel.writeInbound(unexpected);
	assertEquals("Buffer not recycled", 0L, unexpected.refCnt());
}
 

/**
 * Tests that the channel is closed if an Exception reaches the channel handler.
 */
@Test
public void testCloseChannelOnExceptionCaught() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(handler);

	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable response = MessageSerializer.deserializeServerFailure(buf);
	buf.release();

	assertTrue(response.getMessage().contains("Expected test Exception"));

	channel.closeFuture().await(READ_TIMEOUT_MILLIS);
	assertFalse(channel.isActive());
}
 

public TaskManagerServicesBuilder() {
	unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
	ioManager = mock(IOManager.class);
	shuffleEnvironment = mock(ShuffleEnvironment.class);
	kvStateService = new KvStateService(new KvStateRegistry(), null, null);
	broadcastVariableManager = new BroadcastVariableManager();
	taskEventDispatcher = new TaskEventDispatcher();
	taskSlotTable = TestingTaskSlotTable.<Task>newBuilder().closeAsyncReturns(CompletableFuture.completedFuture(null)).build();
	jobTable = DefaultJobTable.create();
	jobLeaderService = new DefaultJobLeaderService(unresolvedTaskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());
	taskStateManager = mock(TaskExecutorLocalStateStoresManager.class);
	ioExecutor = TestingUtils.defaultExecutor();
	libraryCacheManager = TestingLibraryCacheManager.newBuilder().build();
}
 

public TriggerTestHarness(
		Trigger<T, W> trigger,
		TypeSerializer<W> windowSerializer) throws Exception {
	this.trigger = trigger;
	this.windowSerializer = windowSerializer;

	// we only ever use one key, other tests make sure that windows work across different
	// keys
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	MemoryStateBackend backend = new MemoryStateBackend();

	@SuppressWarnings("unchecked")
	HeapKeyedStateBackend<Integer> stateBackend = (HeapKeyedStateBackend<Integer>) backend.createKeyedStateBackend(
		dummyEnv,
		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());
	this.stateBackend = stateBackend;

	this.stateBackend.setCurrentKey(KEY);

	this.internalTimerService = new TestInternalTimerService<>(new KeyContext() {
		@Override
		public void setCurrentKey(Object key) {
			// ignore
		}

		@Override
		public Object getCurrentKey() {
			return KEY;
		}
	});
}
 

/**
 * Creates and returns the KvState service.
 *
 * @param taskManagerServicesConfiguration task manager configuration
 * @return service for kvState related components
 */
public static KvStateService fromConfiguration(TaskManagerServicesConfiguration taskManagerServicesConfiguration) {
	KvStateRegistry kvStateRegistry = new KvStateRegistry();

	QueryableStateConfiguration qsConfig = taskManagerServicesConfiguration.getQueryableStateConfig();

	KvStateClientProxy kvClientProxy = null;
	KvStateServer kvStateServer = null;

	if (qsConfig != null) {
		int numProxyServerNetworkThreads = qsConfig.numProxyServerThreads() == 0 ?
			taskManagerServicesConfiguration.getNumberOfSlots() : qsConfig.numProxyServerThreads();
		int numProxyServerQueryThreads = qsConfig.numProxyQueryThreads() == 0 ?
			taskManagerServicesConfiguration.getNumberOfSlots() : qsConfig.numProxyQueryThreads();
		kvClientProxy = QueryableStateUtils.createKvStateClientProxy(
			taskManagerServicesConfiguration.getTaskManagerAddress(),
			qsConfig.getProxyPortRange(),
			numProxyServerNetworkThreads,
			numProxyServerQueryThreads,
			new DisabledKvStateRequestStats());

		int numStateServerNetworkThreads = qsConfig.numStateServerThreads() == 0 ?
			taskManagerServicesConfiguration.getNumberOfSlots() : qsConfig.numStateServerThreads();
		int numStateServerQueryThreads = qsConfig.numStateQueryThreads() == 0 ?
			taskManagerServicesConfiguration.getNumberOfSlots() : qsConfig.numStateQueryThreads();
		kvStateServer = QueryableStateUtils.createKvStateServer(
			taskManagerServicesConfiguration.getTaskManagerAddress(),
			qsConfig.getStateServerPortRange(),
			numStateServerNetworkThreads,
			numStateServerQueryThreads,
			kvStateRegistry,
			new DisabledKvStateRequestStats());
	}

	return new KvStateService(kvStateRegistry, kvStateServer, kvClientProxy);
}
 

/**
 * Tests the failure response with {@link UnknownKvStateIdException} as cause on
 * queries for unregistered KvStateIDs.
 */
@Test
public void testQueryUnknownKvStateID() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	long requestId = Integer.MAX_VALUE + 182828L;

	KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);

	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.REQUEST_FAILURE, MessageSerializer.deserializeHeader(buf));
	RequestFailure response = MessageSerializer.deserializeRequestFailure(buf);

	assertEquals(requestId, response.getRequestId());

	assertTrue("Did not respond with expected failure cause", response.getCause() instanceof UnknownKvStateIdException);

	assertEquals(1L, stats.getNumRequests());
	assertEquals(1L, stats.getNumFailed());
}
 

@BeforeClass
public static void setup() {
	try {
		testServer = new KvStateServerImpl(
				InetAddress.getLocalHost(),
				Collections.singletonList(0).iterator(),
				1,
				1,
				new KvStateRegistry(),
				new DisabledKvStateRequestStats());
		testServer.start();
	} catch (Throwable e) {
		e.printStackTrace();
	}
}