org.apache.flink.queryablestate.KvStateID Java Examples

private CompletableFuture<KvStateResponse> getState(
		final KvStateRequest request,
		final boolean forceUpdate) {

	return getKvStateLookupInfo(request.getJobId(), request.getStateName(), forceUpdate)
			.thenComposeAsync((Function<KvStateLocation, CompletableFuture<KvStateResponse>>) location -> {
				final int keyGroupIndex = KeyGroupRangeAssignment.computeKeyGroupForKeyHash(
						request.getKeyHashCode(), location.getNumKeyGroups());

				final InetSocketAddress serverAddress = location.getKvStateServerAddress(keyGroupIndex);
				if (serverAddress == null) {
					return FutureUtils.completedExceptionally(new UnknownKvStateKeyGroupLocationException(getServerName()));
				} else {
					// Query server
					final KvStateID kvStateId = location.getKvStateID(keyGroupIndex);
					final KvStateInternalRequest internalRequest = new KvStateInternalRequest(
							kvStateId, request.getSerializedKeyAndNamespace());
					return kvStateClient.sendRequest(serverAddress, internalRequest);
				}
			}, queryExecutor);
}
 

/**
 * Notifies the JobManager about a registered {@link InternalKvState} instance.
 *
 * @param jobId                JobID the KvState instance belongs to
 * @param jobVertexId          JobVertexID the KvState instance belongs to
 * @param keyGroupRange        Key group range the KvState instance belongs to
 * @param registrationName     Name under which the KvState has been registered
 * @param kvStateId            ID of the registered KvState instance
 * @param kvStateServerAddress Server address where to find the KvState instance
 */
public NotifyKvStateRegistered(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId,
		InetSocketAddress kvStateServerAddress) {

	this.jobId = Preconditions.checkNotNull(jobId, "JobID");
	this.jobVertexId = Preconditions.checkNotNull(jobVertexId, "JobVertexID");
	Preconditions.checkArgument(keyGroupRange != KeyGroupRange.EMPTY_KEY_GROUP_RANGE);
	this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
	this.registrationName = Preconditions.checkNotNull(registrationName, "Registration name");
	this.kvStateId = Preconditions.checkNotNull(kvStateId, "KvStateID");
	this.kvStateServerAddress = Preconditions.checkNotNull(kvStateServerAddress, "ServerAddress");
}
 

@Override
public void notifyKvStateRegistered(final JobID jobId, final JobVertexID jobVertexId, final KeyGroupRange keyGroupRange, final String registrationName, final KvStateID kvStateId, final InetSocketAddress kvStateServerAddress) throws FlinkJobNotFoundException {
	mainThreadExecutor.assertRunningInMainThread();

	if (jobGraph.getJobID().equals(jobId)) {
		if (log.isDebugEnabled()) {
			log.debug("Key value state registered for job {} under name {}.",
				jobGraph.getJobID(), registrationName);
		}

		try {
			executionGraph.getKvStateLocationRegistry().notifyKvStateRegistered(
				jobVertexId, keyGroupRange, registrationName, kvStateId, kvStateServerAddress);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	} else {
		throw new FlinkJobNotFoundException(jobId);
	}
}
 

/**
 * Unregisters the KvState instance identified by the given KvStateID.
 *
 * @param jobId     JobId the KvState instance belongs to
 * @param kvStateId KvStateID to identify the KvState instance
 * @param keyGroupRange    Key group range the KvState instance belongs to
 */
public void unregisterKvState(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {

	KvStateEntry<?, ?, ?> entry = registeredKvStates.remove(kvStateId);
	if (entry != null) {
		entry.clear();

		final KvStateRegistryListener listener = getKvStateRegistryListener(jobId);
		if (listener != null) {
			listener.notifyKvStateUnregistered(
					jobId,
					jobVertexId,
					keyGroupRange,
					registrationName);
		}
	}
}
 

@Override
public CompletableFuture<Acknowledge> notifyKvStateRegistered(
		final JobID jobId,
		final JobVertexID jobVertexId,
		final KeyGroupRange keyGroupRange,
		final String registrationName,
		final KvStateID kvStateId,
		final InetSocketAddress kvStateServerAddress) {

	try {
		schedulerNG.notifyKvStateRegistered(jobId, jobVertexId, keyGroupRange, registrationName, kvStateId, kvStateServerAddress);
		return CompletableFuture.completedFuture(Acknowledge.get());
	} catch (FlinkJobNotFoundException e) {
		log.info("Error while receiving notification about key-value state registration", e);
		return FutureUtils.completedExceptionally(e);
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

@Override
public CompletableFuture<Acknowledge> notifyKvStateRegistered(
		final JobID jobId,
		final JobVertexID jobVertexId,
		final KeyGroupRange keyGroupRange,
		final String registrationName,
		final KvStateID kvStateId,
		final InetSocketAddress kvStateServerAddress) {

	try {
		schedulerNG.notifyKvStateRegistered(jobId, jobVertexId, keyGroupRange, registrationName, kvStateId, kvStateServerAddress);
		return CompletableFuture.completedFuture(Acknowledge.get());
	} catch (FlinkJobNotFoundException e) {
		log.info("Error while receiving notification about key-value state registration", e);
		return FutureUtils.completedExceptionally(e);
	}
}
 

/**
 * Tests request serialization.
 */
@Test
public void testRequestSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = randomByteArray(1024);

	final KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

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

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateInternalRequest requestDeser = serializer.deserializeRequest(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(kvStateId, requestDeser.getKvStateId());
	assertArrayEquals(serializedKeyAndNamespace, requestDeser.getSerializedKeyAndNamespace());
}
 

private CompletableFuture<KvStateResponse> getState(
		final KvStateRequest request,
		final boolean forceUpdate) {

	return getKvStateLookupInfo(request.getJobId(), request.getStateName(), forceUpdate)
			.thenComposeAsync((Function<KvStateLocation, CompletableFuture<KvStateResponse>>) location -> {
				final int keyGroupIndex = KeyGroupRangeAssignment.computeKeyGroupForKeyHash(
						request.getKeyHashCode(), location.getNumKeyGroups());

				final InetSocketAddress serverAddress = location.getKvStateServerAddress(keyGroupIndex);
				if (serverAddress == null) {
					return FutureUtils.completedExceptionally(new UnknownKvStateKeyGroupLocationException(getServerName()));
				} else {
					// Query server
					final KvStateID kvStateId = location.getKvStateID(keyGroupIndex);
					final KvStateInternalRequest internalRequest = new KvStateInternalRequest(
							kvStateId, request.getSerializedKeyAndNamespace());
					return kvStateClient.sendRequest(serverAddress, internalRequest);
				}
			}, queryExecutor);
}
 

/**
 * Tests request serialization.
 */
@Test
public void testRequestSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = randomByteArray(1024);

	final KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

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

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateInternalRequest requestDeser = serializer.deserializeRequest(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(kvStateId, requestDeser.getKvStateId());
	assertArrayEquals(serializedKeyAndNamespace, requestDeser.getSerializedKeyAndNamespace());
}
 

/**
 * Tests request serialization with zero-length serialized key and namespace.
 */
@Test
public void testRequestSerializationWithZeroLengthKeyAndNamespace() throws Exception {

	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = new byte[0];

	final KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

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

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateInternalRequest requestDeser = serializer.deserializeRequest(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(kvStateId, requestDeser.getKvStateId());
	assertArrayEquals(serializedKeyAndNamespace, requestDeser.getSerializedKeyAndNamespace());
}
 

/**
 * Unregisters the KvState instance identified by the given KvStateID.
 *
 * @param jobId     JobId the KvState instance belongs to
 * @param kvStateId KvStateID to identify the KvState instance
 * @param keyGroupRange    Key group range the KvState instance belongs to
 */
public void unregisterKvState(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {

	KvStateEntry<?, ?, ?> entry = registeredKvStates.remove(kvStateId);
	if (entry != null) {
		entry.clear();

		final KvStateRegistryListener listener = getKvStateRegistryListener(jobId);
		if (listener != null) {
			listener.notifyKvStateUnregistered(
					jobId,
					jobVertexId,
					keyGroupRange,
					registrationName);
		}
	}
}
 

/**
 * Unregisters the KvState instance identified by the given KvStateID.
 *
 * @param jobId     JobId the KvState instance belongs to
 * @param kvStateId KvStateID to identify the KvState instance
 * @param keyGroupRange    Key group range the KvState instance belongs to
 */
public void unregisterKvState(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {

	KvStateEntry<?, ?, ?> entry = registeredKvStates.remove(kvStateId);
	if (entry != null) {
		entry.clear();

		final KvStateRegistryListener listener = getKvStateRegistryListener(jobId);
		if (listener != null) {
			listener.notifyKvStateUnregistered(
					jobId,
					jobVertexId,
					keyGroupRange,
					registrationName);
		}
	}
}
 

/**
 * Tests request serialization with zero-length serialized key and namespace.
 */
@Test
public void testRequestSerializationWithZeroLengthKeyAndNamespace() throws Exception {

	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = new byte[0];

	final KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

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

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateInternalRequest requestDeser = serializer.deserializeRequest(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(kvStateId, requestDeser.getKvStateId());
	assertArrayEquals(serializedKeyAndNamespace, requestDeser.getSerializedKeyAndNamespace());
}
 

/**
 * Notifies the JobManager about a registered {@link InternalKvState} instance.
 *
 * @param jobId                JobID the KvState instance belongs to
 * @param jobVertexId          JobVertexID the KvState instance belongs to
 * @param keyGroupRange        Key group range the KvState instance belongs to
 * @param registrationName     Name under which the KvState has been registered
 * @param kvStateId            ID of the registered KvState instance
 * @param kvStateServerAddress Server address where to find the KvState instance
 */
public NotifyKvStateRegistered(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId,
		InetSocketAddress kvStateServerAddress) {

	this.jobId = Preconditions.checkNotNull(jobId, "JobID");
	this.jobVertexId = Preconditions.checkNotNull(jobVertexId, "JobVertexID");
	Preconditions.checkArgument(keyGroupRange != KeyGroupRange.EMPTY_KEY_GROUP_RANGE);
	this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
	this.registrationName = Preconditions.checkNotNull(registrationName, "Registration name");
	this.kvStateId = Preconditions.checkNotNull(kvStateId, "KvStateID");
	this.kvStateServerAddress = Preconditions.checkNotNull(kvStateServerAddress, "ServerAddress");
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

/**
 * Tests request serialization with zero-length serialized key and namespace.
 */
@Test
public void testRequestSerializationWithZeroLengthKeyAndNamespace() throws Exception {

	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = new byte[0];

	final KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

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

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateInternalRequest requestDeser = serializer.deserializeRequest(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(kvStateId, requestDeser.getKvStateId());
	assertArrayEquals(serializedKeyAndNamespace, requestDeser.getSerializedKeyAndNamespace());
}
 

@Override
public void notifyKvStateRegistered(JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {
	this.jobVertexID = jobVertexId;
	this.keyGroupIndex = keyGroupRange;
	this.registrationName = registrationName;
	this.kvStateId = kvStateId;
}
 

public KvStateInternalRequest(
		final KvStateID stateId,
		final byte[] serializedKeyAndNamespace) {

	this.kvStateId = Preconditions.checkNotNull(stateId);
	this.serializedKeyAndNamespace = Preconditions.checkNotNull(serializedKeyAndNamespace);
}
 

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

/**
 * Tests that registrations with duplicate names throw an Exception.
 */
@Test
public void testRegisterDuplicateName() throws Exception {
	ExecutionJobVertex[] vertices = new ExecutionJobVertex[] {
			createJobVertex(32),
			createJobVertex(13) };

	Map<JobVertexID, ExecutionJobVertex> vertexMap = createVertexMap(vertices);

	String registrationName = "duplicated-name";
	KvStateLocationRegistry registry = new KvStateLocationRegistry(new JobID(), vertexMap);

	// First operator registers
	registry.notifyKvStateRegistered(
			vertices[0].getJobVertexId(),
			new KeyGroupRange(0, 0),
			registrationName,
			new KvStateID(),
			new InetSocketAddress(InetAddress.getLocalHost(), 12328));

	try {
		// Second operator registers same name
		registry.notifyKvStateRegistered(
				vertices[1].getJobVertexId(),
				new KeyGroupRange(0, 0),
				registrationName,
				new KvStateID(),
				new InetSocketAddress(InetAddress.getLocalHost(), 12032));

		fail("Did not throw expected Exception after duplicated name");
	} catch (IllegalStateException ignored) {
		// Expected
	}
}
 

@Override
public void notifyKvStateRegistered(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {
	kvStateRegistryGateway.notifyKvStateRegistered(
		jobId,
		jobVertexId,
		keyGroupRange,
		registrationName,
		kvStateId,
		kvStateServerAddress);
}
 

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

/**
 * Creates the location information.
 *
 * @param jobId            JobID the KvState instances belong to
 * @param jobVertexId      JobVertexID the KvState instances belong to
 * @param numKeyGroups     Number of key groups of the operator
 * @param registrationName Name under which the KvState instances have been registered
 */
public KvStateLocation(JobID jobId, JobVertexID jobVertexId, int numKeyGroups, String registrationName) {
	this.jobId = Preconditions.checkNotNull(jobId, "JobID");
	this.jobVertexId = Preconditions.checkNotNull(jobVertexId, "JobVertexID");
	Preconditions.checkArgument(numKeyGroups >= 0, "Negative number of key groups");
	this.numKeyGroups = numKeyGroups;
	this.registrationName = Preconditions.checkNotNull(registrationName, "Registration name");
	this.kvStateIds = new KvStateID[numKeyGroups];
	this.kvStateAddresses = new InetSocketAddress[numKeyGroups];
}
 

@Override
public KvStateInternalRequest deserializeMessage(ByteBuf buf) {
	KvStateID kvStateId = new KvStateID(buf.readLong(), buf.readLong());

	int length = buf.readInt();
	Preconditions.checkArgument(length >= 0,
			"Negative length for key and namespace. " +
					"This indicates a serialization error.");

	byte[] serializedKeyAndNamespace = new byte[length];
	if (length > 0) {
		buf.readBytes(serializedKeyAndNamespace);
	}
	return new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
}
 

/**
 * Creates the location information.
 *
 * @param jobId            JobID the KvState instances belong to
 * @param jobVertexId      JobVertexID the KvState instances belong to
 * @param numKeyGroups     Number of key groups of the operator
 * @param registrationName Name under which the KvState instances have been registered
 */
public KvStateLocation(JobID jobId, JobVertexID jobVertexId, int numKeyGroups, String registrationName) {
	this.jobId = Preconditions.checkNotNull(jobId, "JobID");
	this.jobVertexId = Preconditions.checkNotNull(jobVertexId, "JobVertexID");
	Preconditions.checkArgument(numKeyGroups >= 0, "Negative number of key groups");
	this.numKeyGroups = numKeyGroups;
	this.registrationName = Preconditions.checkNotNull(registrationName, "Registration name");
	this.kvStateIds = new KvStateID[numKeyGroups];
	this.kvStateAddresses = new InetSocketAddress[numKeyGroups];
}
 

@Override
public void notifyKvStateRegistered(JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {
	this.jobVertexID = jobVertexId;
	this.keyGroupIndex = keyGroupRange;
	this.registrationName = registrationName;
	this.kvStateId = kvStateId;
}
 

/**
 * 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);

	assertEquals(requestId, response.getRequestId());

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

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