Examples with KvStateID org.apache.flink.queryablestate.KvStateID used on opensource projects

Example 16 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class JobMasterQueryableStateTest method testRegisterKvState.

@Test
public void testRegisterKvState() throws Exception {
    final JobMaster jobMaster = new JobMasterBuilder(JOB_GRAPH, rpcService).createJobMaster();
    jobMaster.start();
    final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);
    registerSlotsRequiredForJobExecution(jobMasterGateway, JOB_GRAPH.getJobID());
    try {
        final String registrationName = "register-me";
        final KvStateID kvStateID = new KvStateID();
        final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 0);
        final InetSocketAddress address = new InetSocketAddress(InetAddress.getLocalHost(), 1029);
        jobMasterGateway.notifyKvStateRegistered(JOB_GRAPH.getJobID(), JOB_VERTEX_1.getID(), keyGroupRange, registrationName, kvStateID, address).get();
        final KvStateLocation location = jobMasterGateway.requestKvStateLocation(JOB_GRAPH.getJobID(), registrationName).get();
        assertEquals(JOB_GRAPH.getJobID(), location.getJobId());
        assertEquals(JOB_VERTEX_1.getID(), location.getJobVertexId());
        assertEquals(JOB_VERTEX_1.getMaxParallelism(), location.getNumKeyGroups());
        assertEquals(1, location.getNumRegisteredKeyGroups());
        assertEquals(1, keyGroupRange.getNumberOfKeyGroups());
        assertEquals(kvStateID, location.getKvStateID(keyGroupRange.getStartKeyGroup()));
        assertEquals(address, location.getKvStateServerAddress(keyGroupRange.getStartKeyGroup()));
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}

Example 17 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateLocationTest method testRegisterAndLookup.

/**
 * Simple test registering/unregistering state and looking it up again.
 */
@Test
public void testRegisterAndLookup() throws Exception {
    JobID jobId = new JobID();
    JobVertexID jobVertexId = new JobVertexID();
    int numKeyGroups = 123;
    int numRanges = 10;
    int fract = numKeyGroups / numRanges;
    int remain = numKeyGroups % numRanges;
    List<KeyGroupRange> keyGroupRanges = new ArrayList<>(numRanges);
    int start = 0;
    for (int i = 0; i < numRanges; ++i) {
        int end = start + fract - 1;
        if (remain > 0) {
            --remain;
            ++end;
        }
        KeyGroupRange range = new KeyGroupRange(start, end);
        keyGroupRanges.add(range);
        start = end + 1;
    }
    String registrationName = "asdasdasdasd";
    KvStateLocation location = new KvStateLocation(jobId, jobVertexId, numKeyGroups, registrationName);
    KvStateID[] kvStateIds = new KvStateID[numRanges];
    InetSocketAddress[] serverAddresses = new InetSocketAddress[numRanges];
    InetAddress host = InetAddress.getLocalHost();
    // Register
    int registeredCount = 0;
    for (int rangeIdx = 0; rangeIdx < numRanges; rangeIdx++) {
        kvStateIds[rangeIdx] = new KvStateID();
        serverAddresses[rangeIdx] = new InetSocketAddress(host, 1024 + rangeIdx);
        KeyGroupRange keyGroupRange = keyGroupRanges.get(rangeIdx);
        location.registerKvState(keyGroupRange, kvStateIds[rangeIdx], serverAddresses[rangeIdx]);
        registeredCount += keyGroupRange.getNumberOfKeyGroups();
        assertEquals(registeredCount, location.getNumRegisteredKeyGroups());
    }
    // Lookup
    for (int rangeIdx = 0; rangeIdx < numRanges; rangeIdx++) {
        KeyGroupRange keyGroupRange = keyGroupRanges.get(rangeIdx);
        for (int keyGroup = keyGroupRange.getStartKeyGroup(); keyGroup <= keyGroupRange.getEndKeyGroup(); ++keyGroup) {
            assertEquals(kvStateIds[rangeIdx], location.getKvStateID(keyGroup));
            assertEquals(serverAddresses[rangeIdx], location.getKvStateServerAddress(keyGroup));
        }
    }
    // Overwrite
    for (int rangeIdx = 0; rangeIdx < numRanges; rangeIdx++) {
        kvStateIds[rangeIdx] = new KvStateID();
        serverAddresses[rangeIdx] = new InetSocketAddress(host, 1024 + rangeIdx);
        location.registerKvState(keyGroupRanges.get(rangeIdx), kvStateIds[rangeIdx], serverAddresses[rangeIdx]);
        assertEquals(registeredCount, location.getNumRegisteredKeyGroups());
    }
    // Lookup
    for (int rangeIdx = 0; rangeIdx < numRanges; rangeIdx++) {
        KeyGroupRange keyGroupRange = keyGroupRanges.get(rangeIdx);
        for (int keyGroup = keyGroupRange.getStartKeyGroup(); keyGroup <= keyGroupRange.getEndKeyGroup(); ++keyGroup) {
            assertEquals(kvStateIds[rangeIdx], location.getKvStateID(keyGroup));
            assertEquals(serverAddresses[rangeIdx], location.getKvStateServerAddress(keyGroup));
        }
    }
    // Unregister
    for (int rangeIdx = 0; rangeIdx < numRanges; rangeIdx++) {
        KeyGroupRange keyGroupRange = keyGroupRanges.get(rangeIdx);
        location.unregisterKvState(keyGroupRange);
        registeredCount -= keyGroupRange.getNumberOfKeyGroups();
        assertEquals(registeredCount, location.getNumRegisteredKeyGroups());
    }
    // Lookup
    for (int rangeIdx = 0; rangeIdx < numRanges; rangeIdx++) {
        KeyGroupRange keyGroupRange = keyGroupRanges.get(rangeIdx);
        for (int keyGroup = keyGroupRange.getStartKeyGroup(); keyGroup <= keyGroupRange.getEndKeyGroup(); ++keyGroup) {
            assertEquals(null, location.getKvStateID(keyGroup));
            assertEquals(null, location.getKvStateServerAddress(keyGroup));
        }
    }
    assertEquals(0, location.getNumRegisteredKeyGroups());
}

Example 18 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateRegistryTest method testKvStateRegistryListenerNotification.

/**
 * Tests that {@link KvStateRegistryListener} only receive the notifications which are destined
 * for them.
 */
@Test
public void testKvStateRegistryListenerNotification() {
    final JobID jobId1 = new JobID();
    final JobID jobId2 = new JobID();
    final KvStateRegistry kvStateRegistry = new KvStateRegistry();
    final ArrayDeque<JobID> registeredNotifications1 = new ArrayDeque<>(2);
    final ArrayDeque<JobID> deregisteredNotifications1 = new ArrayDeque<>(2);
    final TestingKvStateRegistryListener listener1 = new TestingKvStateRegistryListener(registeredNotifications1, deregisteredNotifications1);
    final ArrayDeque<JobID> registeredNotifications2 = new ArrayDeque<>(2);
    final ArrayDeque<JobID> deregisteredNotifications2 = new ArrayDeque<>(2);
    final TestingKvStateRegistryListener listener2 = new TestingKvStateRegistryListener(registeredNotifications2, deregisteredNotifications2);
    kvStateRegistry.registerListener(jobId1, listener1);
    kvStateRegistry.registerListener(jobId2, listener2);
    final JobVertexID jobVertexId = new JobVertexID();
    final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
    final String registrationName = "foobar";
    final KvStateID kvStateID = kvStateRegistry.registerKvState(jobId1, jobVertexId, keyGroupRange, registrationName, new DummyKvState(), getClass().getClassLoader());
    assertThat(registeredNotifications1.poll(), equalTo(jobId1));
    assertThat(registeredNotifications2.isEmpty(), is(true));
    final JobVertexID jobVertexId2 = new JobVertexID();
    final KeyGroupRange keyGroupRange2 = new KeyGroupRange(0, 1);
    final String registrationName2 = "barfoo";
    final KvStateID kvStateID2 = kvStateRegistry.registerKvState(jobId2, jobVertexId2, keyGroupRange2, registrationName2, new DummyKvState(), getClass().getClassLoader());
    assertThat(registeredNotifications2.poll(), equalTo(jobId2));
    assertThat(registeredNotifications1.isEmpty(), is(true));
    kvStateRegistry.unregisterKvState(jobId1, jobVertexId, keyGroupRange, registrationName, kvStateID);
    assertThat(deregisteredNotifications1.poll(), equalTo(jobId1));
    assertThat(deregisteredNotifications2.isEmpty(), is(true));
    kvStateRegistry.unregisterKvState(jobId2, jobVertexId2, keyGroupRange2, registrationName2, kvStateID2);
    assertThat(deregisteredNotifications2.poll(), equalTo(jobId2));
    assertThat(deregisteredNotifications1.isEmpty(), is(true));
}

Example 19 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateLocationRegistryTest method testRegisterAndLookup.

/**
 * Simple test registering/unregistering state and looking it up again.
 */
@Test
public void testRegisterAndLookup() throws Exception {
    String[] registrationNames = new String[] { "TAsIrGnc7MULwVupNKZ0", "086133IrGn0Ii2853237" };
    ExecutionJobVertex[] vertices = new ExecutionJobVertex[] { createJobVertex(32), createJobVertex(13) };
    // IDs for each key group of each vertex
    KvStateID[][] ids = new KvStateID[vertices.length][];
    for (int i = 0; i < ids.length; i++) {
        ids[i] = new KvStateID[vertices[i].getMaxParallelism()];
        for (int j = 0; j < vertices[i].getMaxParallelism(); j++) {
            ids[i][j] = new KvStateID();
        }
    }
    InetSocketAddress server = new InetSocketAddress(InetAddress.getLocalHost(), 12032);
    // Create registry
    Map<JobVertexID, ExecutionJobVertex> vertexMap = createVertexMap(vertices);
    KvStateLocationRegistry registry = new KvStateLocationRegistry(new JobID(), vertexMap);
    // Register
    for (int i = 0; i < vertices.length; i++) {
        int numKeyGroups = vertices[i].getMaxParallelism();
        for (int keyGroupIndex = 0; keyGroupIndex < numKeyGroups; keyGroupIndex++) {
            // Register
            registry.notifyKvStateRegistered(vertices[i].getJobVertexId(), new KeyGroupRange(keyGroupIndex, keyGroupIndex), registrationNames[i], ids[i][keyGroupIndex], server);
        }
    }
    // Verify all registrations
    for (int i = 0; i < vertices.length; i++) {
        KvStateLocation location = registry.getKvStateLocation(registrationNames[i]);
        assertNotNull(location);
        int maxParallelism = vertices[i].getMaxParallelism();
        for (int keyGroupIndex = 0; keyGroupIndex < maxParallelism; keyGroupIndex++) {
            assertEquals(ids[i][keyGroupIndex], location.getKvStateID(keyGroupIndex));
            assertEquals(server, location.getKvStateServerAddress(keyGroupIndex));
        }
    }
    // Unregister
    for (int i = 0; i < vertices.length; i++) {
        int numKeyGroups = vertices[i].getMaxParallelism();
        JobVertexID jobVertexId = vertices[i].getJobVertexId();
        for (int keyGroupIndex = 0; keyGroupIndex < numKeyGroups; keyGroupIndex++) {
            registry.notifyKvStateUnregistered(jobVertexId, new KeyGroupRange(keyGroupIndex, keyGroupIndex), registrationNames[i]);
        }
    }
    for (int i = 0; i < registrationNames.length; i++) {
        assertNull(registry.getKvStateLocation(registrationNames[i]));
    }
}

Example 20 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateLocationRegistryTest method testUnregisterFailures.

/**
 * Tests failures during unregistration.
 */
@Test
public void testUnregisterFailures() throws Exception {
    String name = "IrGnc73237TAs";
    ExecutionJobVertex[] vertices = new ExecutionJobVertex[] { createJobVertex(32), createJobVertex(13) };
    Map<JobVertexID, ExecutionJobVertex> vertexMap = new HashMap<>();
    for (ExecutionJobVertex vertex : vertices) {
        vertexMap.put(vertex.getJobVertexId(), vertex);
    }
    KvStateLocationRegistry registry = new KvStateLocationRegistry(new JobID(), vertexMap);
    // First operator registers name
    registry.notifyKvStateRegistered(vertices[0].getJobVertexId(), new KeyGroupRange(0, 0), name, new KvStateID(), mock(InetSocketAddress.class));
    try {
        // Unregister not registered keyGroupIndex
        int notRegisteredKeyGroupIndex = 2;
        registry.notifyKvStateUnregistered(vertices[0].getJobVertexId(), new KeyGroupRange(notRegisteredKeyGroupIndex, notRegisteredKeyGroupIndex), name);
        fail("Did not throw expected Exception");
    } catch (IllegalArgumentException expected) {
    }
    try {
        // Wrong operator tries to unregister
        registry.notifyKvStateUnregistered(vertices[1].getJobVertexId(), new KeyGroupRange(0, 0), name);
        fail("Did not throw expected Exception");
    } catch (IllegalArgumentException expected) {
    }
}