@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())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats(); KvStateID kvStateId = registry.registerKvState( new JobID(), new JobVertexID(),
Bootstrap bootstrap = null; try { KvStateRegistry registry = new KvStateRegistry(); KvStateRequestStats stats = new AtomicKvStateRequestStats(); numKeyGroups, new KeyGroupRange(0, 0), registry.createTaskRegistry(jobId, new JobVertexID())); registry.registerListener(jobId, registryListener);
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats(); 1, 1, new KvStateRegistry(), new DisabledKvStateRequestStats()); registry.registerListener(dummyEnv.getJobID(), registryListener);
@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(); } }
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()); } }
/** * Registers the KvState instance at the KvStateRegistry. * * @param keyGroupRange Key group range the KvState instance belongs to * @param registrationName The registration name (not necessarily the same * as the KvState name defined in the state * descriptor used to create the KvState instance) * @param kvState The */ public void registerKvState(KeyGroupRange keyGroupRange, String registrationName, InternalKvState<?, ?, ?> kvState) { KvStateID kvStateId = registry.registerKvState(jobId, jobVertexId, keyGroupRange, registrationName, kvState); registeredKvStates.add(new KvStateInfo(keyGroupRange, registrationName, kvStateId)); }
/** * Unregisters all registered KvState instances from the KvStateRegistry. */ public void unregisterAll() { for (KvStateInfo kvState : registeredKvStates) { registry.unregisterKvState(jobId, jobVertexId, kvState.keyGroupRange, kvState.registrationName, kvState.kvStateId); } }
/** * 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); } } }
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); } }
private void unregisterQueryableState(JobManagerConnection jobManagerConnection) { final KvStateRegistry kvStateRegistry = networkEnvironment.getKvStateRegistry(); if (kvStateRegistry != null) { kvStateRegistry.unregisterListener(jobManagerConnection.getJobID()); } final KvStateClientProxy kvStateClientProxy = networkEnvironment.getKvStateProxy(); if (kvStateClientProxy != null) { kvStateClientProxy.updateKvStateLocationOracle(jobManagerConnection.getJobID(), null); } }
/** * 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()); }
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats(); registry.registerListener(dummyEnv.getJobID(), registryListener);
public TaskKvStateRegistry createKvStateTaskRegistry(JobID jobId, JobVertexID jobVertexId) { return kvStateRegistry.createTaskRegistry(jobId, jobVertexId); }
/** * Registers the KvState instance at the KvStateRegistry. * * @param keyGroupRange Key group range the KvState instance belongs to * @param registrationName The registration name (not necessarily the same * as the KvState name defined in the state * descriptor used to create the KvState instance) * @param kvState The */ public void registerKvState(KeyGroupRange keyGroupRange, String registrationName, InternalKvState<?, ?, ?> kvState) { KvStateID kvStateId = registry.registerKvState(jobId, jobVertexId, keyGroupRange, registrationName, kvState); registeredKvStates.add(new KvStateInfo(keyGroupRange, registrationName, kvStateId)); }
/** * Unregisters all registered KvState instances from the KvStateRegistry. */ public void unregisterAll() { for (KvStateInfo kvState : registeredKvStates) { registry.unregisterKvState(jobId, jobVertexId, kvState.keyGroupRange, kvState.registrationName, kvState.kvStateId); } }
/** * 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); } } }
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); } }
private void disassociateFromJobManager(JobManagerConnection jobManagerConnection, Exception cause) throws IOException { checkNotNull(jobManagerConnection); final KvStateRegistry kvStateRegistry = networkEnvironment.getKvStateRegistry(); if (kvStateRegistry != null) { kvStateRegistry.unregisterListener(jobManagerConnection.getJobID()); } final KvStateClientProxy kvStateClientProxy = networkEnvironment.getKvStateProxy(); if (kvStateClientProxy != null) { kvStateClientProxy.updateKvStateLocationOracle(jobManagerConnection.getJobID(), null); } JobMasterGateway jobManagerGateway = jobManagerConnection.getJobManagerGateway(); jobManagerGateway.disconnectTaskManager(getResourceID(), cause); jobManagerConnection.getLibraryCacheManager().shutdown(); }
@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())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));