private static <K, N, V> byte[] getSerializedValue( final KvStateEntry<K, N, V> entry, final byte[] serializedKeyAndNamespace) throws Exception { final InternalKvState<K, N, V> state = entry.getState(); final KvStateInfo<K, N, V> infoForCurrentThread = entry.getInfoForCurrentThread(); return state.getSerializedValue( serializedKeyAndNamespace, infoForCurrentThread.getKeySerializer(), infoForCurrentThread.getNamespaceSerializer(), infoForCurrentThread.getStateValueSerializer() ); }
final KvStateLocationOracle kvStateLocationOracle = proxy.getKvStateLocationOracle(jobId);
@Override public CompletableFuture<KvStateResponse> handleRequest(final long requestId, final KvStateInternalRequest request) { final CompletableFuture<KvStateResponse> responseFuture = new CompletableFuture<>(); try { final KvStateEntry<?, ?, ?> kvState = registry.getKvState(request.getKvStateId()); if (kvState == null) { responseFuture.completeExceptionally(new UnknownKvStateIdException(getServerName(), request.getKvStateId())); } else { byte[] serializedKeyAndNamespace = request.getSerializedKeyAndNamespace(); byte[] serializedResult = getSerializedValue(kvState, serializedKeyAndNamespace); if (serializedResult != null) { responseFuture.complete(new KvStateResponse(serializedResult)); } else { responseFuture.completeExceptionally(new UnknownKeyOrNamespaceException(getServerName())); } } return responseFuture; } catch (Throwable t) { String errMsg = "Error while processing request with ID " + requestId + ". Caused by: " + ExceptionUtils.stringifyException(t); responseFuture.completeExceptionally(new RuntimeException(errMsg)); return responseFuture; } }
@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(); 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(); } }
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats(); KvStateID kvStateId = registry.registerKvState( new JobID(), new JobVertexID(),
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()); } }
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); }
@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));
KvStateRegistry registry = new KvStateRegistry(); KvStateRequestStats stats = new AtomicKvStateRequestStats(); registry.registerListener(dummyEnv.getJobID(), registryListener);
new ResultPartitionManager(), new TaskEventDispatcher(), new KvStateRegistry(), null, null,
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())); 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; } }); }
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats(); registry.registerListener(dummyEnv.getJobID(), registryListener);
/** * 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()); }
this.accumulatorRegistry = new AccumulatorRegistry(jobID, getExecutionId()); KvStateRegistry registry = new KvStateRegistry(); this.kvStateRegistry = registry.createTaskRegistry(jobID, getJobVertexId());
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats(); registry.registerListener(dummyEnv.getJobID(), registryListener);
KvStateRegistry registry = new KvStateRegistry(); AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
/** * 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()); }