use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project flink by apache.
the class KvStateServerTest method testSimpleRequest.
/**
* Tests a simple successful query via a SocketChannel.
*/
@Test
public void testSimpleRequest() throws Exception {
KvStateServer server = null;
Bootstrap bootstrap = null;
try {
KvStateRegistry registry = new KvStateRegistry();
KvStateRequestStats stats = new AtomicKvStateRequestStats();
server = new KvStateServer(InetAddress.getLocalHost(), 0, 1, 1, registry, stats);
server.start();
KvStateServerAddress serverAddress = server.getAddress();
int numKeyGroups = 1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
AbstractKeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), registry.createTaskRegistry(new JobID(), new JobVertexID()));
final KvStateServerHandlerTest.TestRegistryListener registryListener = new KvStateServerHandlerTest.TestRegistryListener();
registry.registerListener(registryListener);
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
desc.setQueryable("vanilla");
ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
// Update KvState
int expectedValue = 712828289;
int key = 99812822;
backend.setCurrentKey(key);
state.update(expectedValue);
// Request
byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(key, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
// Connect to the server
final BlockingQueue<ByteBuf> responses = new LinkedBlockingQueue<>();
bootstrap = createBootstrap(new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE, 0, 4, 0, 4), new ChannelInboundHandlerAdapter() {
@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
responses.add((ByteBuf) msg);
}
});
Channel channel = bootstrap.connect(serverAddress.getHost(), serverAddress.getPort()).sync().channel();
long requestId = Integer.MAX_VALUE + 182828L;
assertTrue(registryListener.registrationName.equals("vanilla"));
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(channel.alloc(), requestId, registryListener.kvStateId, serializedKeyAndNamespace);
channel.writeAndFlush(request);
ByteBuf buf = responses.poll(TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
assertEquals(KvStateRequestType.REQUEST_RESULT, KvStateRequestSerializer.deserializeHeader(buf));
KvStateRequestResult response = KvStateRequestSerializer.deserializeKvStateRequestResult(buf);
assertEquals(requestId, response.getRequestId());
int actualValue = KvStateRequestSerializer.deserializeValue(response.getSerializedResult(), IntSerializer.INSTANCE);
assertEquals(expectedValue, actualValue);
} finally {
if (server != null) {
server.shutDown();
}
if (bootstrap != null) {
EventLoopGroup group = bootstrap.group();
if (group != null) {
group.shutdownGracefully();
}
}
}
}
use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project flink by apache.
the class QueryableStateClientTest method testIntegrationWithKvStateServer.
/**
* Tests queries against multiple servers.
*
* <p>The servers are populated with different keys and the client queries
* all available keys from all servers.
*/
@Test
public void testIntegrationWithKvStateServer() throws Exception {
// Config
int numServers = 2;
int numKeys = 1024;
int numKeyGroups = 1;
JobID jobId = new JobID();
JobVertexID jobVertexId = new JobVertexID();
KvStateServer[] servers = new KvStateServer[numServers];
AtomicKvStateRequestStats[] serverStats = new AtomicKvStateRequestStats[numServers];
QueryableStateClient client = null;
KvStateClient networkClient = null;
AtomicKvStateRequestStats networkClientStats = new AtomicKvStateRequestStats();
MemoryStateBackend backend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
AbstractKeyedStateBackend<Integer> keyedStateBackend = backend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
try {
KvStateRegistry[] registries = new KvStateRegistry[numServers];
KvStateID[] kvStateIds = new KvStateID[numServers];
List<HeapValueState<Integer, VoidNamespace, Integer>> kvStates = new ArrayList<>();
// Start the servers
for (int i = 0; i < numServers; i++) {
registries[i] = new KvStateRegistry();
serverStats[i] = new AtomicKvStateRequestStats();
servers[i] = new KvStateServer(InetAddress.getLocalHost(), 0, 1, 1, registries[i], serverStats[i]);
servers[i].start();
ValueStateDescriptor<Integer> descriptor = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
RegisteredBackendStateMetaInfo<VoidNamespace, Integer> registeredBackendStateMetaInfo = new RegisteredBackendStateMetaInfo<>(descriptor.getType(), descriptor.getName(), VoidNamespaceSerializer.INSTANCE, IntSerializer.INSTANCE);
// Register state
HeapValueState<Integer, VoidNamespace, Integer> kvState = new HeapValueState<>(descriptor, new NestedMapsStateTable<Integer, VoidNamespace, Integer>(keyedStateBackend, registeredBackendStateMetaInfo), IntSerializer.INSTANCE, VoidNamespaceSerializer.INSTANCE);
kvStates.add(kvState);
kvStateIds[i] = registries[i].registerKvState(jobId, new JobVertexID(), new KeyGroupRange(i, i), "choco", kvState);
}
int[] expectedRequests = new int[numServers];
for (int key = 0; key < numKeys; key++) {
int targetKeyGroupIndex = MathUtils.murmurHash(key) % numServers;
expectedRequests[targetKeyGroupIndex]++;
HeapValueState<Integer, VoidNamespace, Integer> kvState = kvStates.get(targetKeyGroupIndex);
keyedStateBackend.setCurrentKey(key);
kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
kvState.update(1337 + key);
}
// Location lookup service
KvStateLocation location = new KvStateLocation(jobId, jobVertexId, numServers, "choco");
for (int keyGroupIndex = 0; keyGroupIndex < numServers; keyGroupIndex++) {
location.registerKvState(new KeyGroupRange(keyGroupIndex, keyGroupIndex), kvStateIds[keyGroupIndex], servers[keyGroupIndex].getAddress());
}
KvStateLocationLookupService lookupService = mock(KvStateLocationLookupService.class);
when(lookupService.getKvStateLookupInfo(eq(jobId), eq("choco"))).thenReturn(Futures.successful(location));
// The client
networkClient = new KvStateClient(1, networkClientStats);
client = new QueryableStateClient(lookupService, networkClient, testActorSystem.dispatcher());
// Send all queries
List<Future<byte[]>> futures = new ArrayList<>(numKeys);
for (int key = 0; key < numKeys; key++) {
byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(key, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
futures.add(client.getKvState(jobId, "choco", key, serializedKeyAndNamespace));
}
// Verify results
Future<Iterable<byte[]>> future = Futures.sequence(futures, testActorSystem.dispatcher());
Iterable<byte[]> results = Await.result(future, timeout);
int index = 0;
for (byte[] buffer : results) {
int deserializedValue = KvStateRequestSerializer.deserializeValue(buffer, IntSerializer.INSTANCE);
assertEquals(1337 + index, deserializedValue);
index++;
}
// Verify requests
for (int i = 0; i < numServers; i++) {
int numRetries = 10;
for (int retry = 0; retry < numRetries; retry++) {
try {
assertEquals("Unexpected number of requests", expectedRequests[i], serverStats[i].getNumRequests());
assertEquals("Unexpected success requests", expectedRequests[i], serverStats[i].getNumSuccessful());
assertEquals("Unexpected failed requests", 0, serverStats[i].getNumFailed());
break;
} catch (Throwable t) {
// Retry
if (retry == numRetries - 1) {
throw t;
} else {
Thread.sleep(100);
}
}
}
}
} finally {
if (client != null) {
client.shutDown();
}
if (networkClient != null) {
networkClient.shutDown();
}
for (KvStateServer server : servers) {
if (server != null) {
server.shutDown();
}
}
}
}
use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project beam by apache.
the class FlinkBroadcastStateInternalsTest method initStateInternals.
@Before
public void initStateInternals() {
MemoryStateBackend backend = new MemoryStateBackend();
try {
OperatorStateBackend operatorStateBackend = backend.createOperatorStateBackend(new DummyEnvironment("test", 1, 0), "");
underTest = new FlinkBroadcastStateInternals<>(1, operatorStateBackend);
} catch (Exception e) {
throw new RuntimeException(e);
}
}
use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project beam by apache.
the class FlinkKeyGroupStateInternalsTest method getKeyedStateBackend.
private KeyedStateBackend getKeyedStateBackend(int numberOfKeyGroups, KeyGroupRange keyGroupRange) {
MemoryStateBackend backend = new MemoryStateBackend();
try {
AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend(new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), numberOfKeyGroups, keyGroupRange, new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
keyedStateBackend.setCurrentKey(ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1")));
return keyedStateBackend;
} catch (Exception e) {
throw new RuntimeException(e);
}
}
use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project beam by apache.
the class FlinkSplitStateInternalsTest method initStateInternals.
@Before
public void initStateInternals() {
MemoryStateBackend backend = new MemoryStateBackend();
try {
OperatorStateBackend operatorStateBackend = backend.createOperatorStateBackend(new DummyEnvironment("test", 1, 0), "");
underTest = new FlinkSplitStateInternals<>(operatorStateBackend);
} catch (Exception e) {
throw new RuntimeException(e);
}
}
Aggregations