private StatefulProcessorBolt<?, ?> findStatefulProcessorBolt(ProcessorNode updateStateNode) { for (StreamBolt bolt : streamBolts.keySet()) { if (bolt instanceof StatefulProcessorBolt) { StatefulProcessorBolt<?, ?> statefulProcessorBolt = (StatefulProcessorBolt) bolt; if (statefulProcessorBolt.getNodes().contains(updateStateNode)) { return statefulProcessorBolt; } } } throw new IllegalArgumentException("Could not find Stateful bolt for node " + updateStateNode); }
StatefulProcessorBolt(String boltId, DirectedGraph<Node, Edge> graph, List<ProcessorNode> nodes) { delegate = new ProcessorBoltDelegate(boltId, graph, nodes); statefulProcessors = getStatefulProcessors(nodes); }
private StreamBolt addStatefulBolt(TopologyBuilder topologyBuilder, String boltId, Set<ProcessorNode> initialProcessors, List<ProcessorNode> group) { StateQueryProcessor<?, ?> stateQueryProcessor = getStateQueryProcessor(group); StatefulProcessorBolt<?, ?> bolt; if (stateQueryProcessor == null) { bolt = new StatefulProcessorBolt<>(boltId, graph, group); BoltDeclarer boltDeclarer = topologyBuilder.setBolt(boltId, bolt, getParallelism(group)); bolt.setStreamToInitialProcessors(wireBolt(group, boltDeclarer, initialProcessors)); streamBolts.put(bolt, boltDeclarer); } else { // state query is added to the existing stateful bolt ProcessorNode updateStateNode = stateQueryProcessor.getStreamState().getUpdateStateNode(); bolt = findStatefulProcessorBolt(updateStateNode); for (ProcessorNode node : group) { node.setComponentId(bolt.getId()); } bolt.addNodes(group); bolt.addStreamToInitialProcessors(wireBolt(bolt.getNodes(), streamBolts.get(bolt), initialProcessors)); } return bolt; }
private void setUpStatefulProcessorBolt(Processor<?> processor) { ProcessorNode node = new ProcessorNode(processor, "outputstream", new Fields("value")); node.setEmitsPair(true); Mockito.when(mockStreamToProcessors.get(Mockito.anyString())).thenReturn(Collections.singletonList(node)); graph = new DefaultDirectedGraph(new StreamsEdgeFactory()); graph.addVertex(node); bolt = new StatefulProcessorBolt<>("bolt1", graph, Collections.singletonList(node)); bolt.setStreamToInitialProcessors(mockStreamToProcessors); bolt.prepare(new HashMap<>(), mockTopologyContext, mockOutputCollector); bolt.initState(mockKeyValueState); }
@Test public void testEmitAndAck() throws Exception { setUpStatefulProcessorBolt(new UpdateStateByKeyProcessor<>(new StateUpdater<Object, Long>() { @Override public Long init() { return 0L; } @Override public Long apply(Long state, Object value) { return state + 1; } })); bolt.execute(mockTuple1); ArgumentCaptor<Collection> anchor = ArgumentCaptor.forClass(Collection.class); ArgumentCaptor<Values> values = ArgumentCaptor.forClass(Values.class); ArgumentCaptor<String> os = ArgumentCaptor.forClass(String.class); Mockito.verify(mockOutputCollector).emit(os.capture(), anchor.capture(), values.capture()); assertEquals("outputstream", os.getValue()); assertArrayEquals(new Object[]{ mockTuple1 }, anchor.getValue().toArray()); assertEquals(new Values("k", 1L), values.getValue()); Mockito.verify(mockOutputCollector, Mockito.times(1)).ack(mockTuple1); Mockito.verify(mockKeyValueState, Mockito.times(1)).put("k", 1L); }
void addNodes(List<ProcessorNode> nodes) { delegate.addNodes(nodes); statefulProcessors.addAll(getStatefulProcessors(nodes)); }