Examples with StateMap org.apache.nifi.components.state.StateMap used on opensource projects

Example 16 with StateMap

use of org.apache.nifi.components.state.StateMap in project nifi by apache.

the class ITRedisStateProvider method testOnComponentRemoved.

@Test
public void testOnComponentRemoved() throws IOException, InterruptedException {
    final StateProvider provider = getProvider();
    final Map<String, String> newValue = new HashMap<>();
    newValue.put("value", "value");
    provider.setState(newValue, componentId);
    final StateMap stateMap = provider.getState(componentId);
    assertEquals(0L, stateMap.getVersion());
    provider.onComponentRemoved(componentId);
    // wait for the background process to complete
    Thread.sleep(1000L);
    final StateMap stateMapAfterRemoval = provider.getState(componentId);
    // version should be -1 because the state has been removed entirely.
    assertEquals(-1L, stateMapAfterRemoval.getVersion());
}

Example 17 with StateMap

use of org.apache.nifi.components.state.StateMap in project nifi by apache.

the class ITRedisStateProvider method testReplaceWithNonExistingValue.

@Test
public void testReplaceWithNonExistingValue() throws Exception {
    final StateProvider provider = getProvider();
    StateMap stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    final Map<String, String> newValue = new HashMap<>();
    newValue.put("value", "value");
    final boolean replaced = provider.replace(stateMap, newValue, componentId);
    assertFalse(replaced);
}

Example 18 with StateMap

use of org.apache.nifi.components.state.StateMap in project nifi by apache.

the class AttributeRollingWindow method onScheduled.

@OnScheduled
public void onScheduled(final ProcessContext context) throws IOException {
    timeWindow = context.getProperty(TIME_WINDOW).asTimePeriod(TimeUnit.MILLISECONDS);
    microBatchTime = context.getProperty(SUB_WINDOW_LENGTH).asTimePeriod(TimeUnit.MILLISECONDS);
    if (microBatchTime == null || microBatchTime == 0) {
        StateManager stateManager = context.getStateManager();
        StateMap state = stateManager.getState(SCOPE);
        HashMap<String, String> tempMap = new HashMap<>();
        tempMap.putAll(state.toMap());
        if (!tempMap.containsKey(COUNT_KEY)) {
            tempMap.put(COUNT_KEY, "0");
            context.getStateManager().setState(tempMap, SCOPE);
        }
    }
}

Example 19 with StateMap

use of org.apache.nifi.components.state.StateMap in project nifi by apache.

the class UpdateAttribute method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final ComponentLog logger = getLogger();
    final Criteria criteria = criteriaCache.get();
    FlowFile incomingFlowFile = session.get();
    if (incomingFlowFile == null) {
        return;
    }
    // record which rule should be applied to which flow file - when operating
    // in 'use clone' mode, this collection will contain a number of entries
    // that map to single element lists. this is because the original flowfile
    // is cloned for each matching rule. in 'use original' mode, this collection
    // will contain a single entry that maps a list of multiple rules. this is
    // because is the original flowfile is used for all matching rules. in this
    // case the order of the matching rules is preserved in the list
    final Map<FlowFile, List<Rule>> matchedRules = new HashMap<>();
    final Map<String, String> stateInitialAttributes;
    final Map<String, String> stateWorkingAttributes;
    StateMap stateMap = null;
    try {
        if (stateful) {
            stateMap = context.getStateManager().getState(Scope.LOCAL);
            stateInitialAttributes = stateMap.toMap();
            stateWorkingAttributes = new HashMap<>(stateMap.toMap());
        } else {
            stateInitialAttributes = null;
            stateWorkingAttributes = null;
        }
    } catch (IOException e) {
        logger.error("Failed to get the initial state when processing {}; transferring FlowFile back to its incoming queue", new Object[] { incomingFlowFile }, e);
        session.transfer(incomingFlowFile);
        context.yield();
        return;
    }
    Map<String, Action> defaultActions = this.defaultActions;
    List<FlowFile> flowFilesToTransfer = new LinkedList<>();
    // if there is update criteria specified, evaluate it
    if (criteria != null && evaluateCriteria(session, context, criteria, incomingFlowFile, matchedRules, stateInitialAttributes)) {
        // apply the actions for each rule and transfer the flowfile
        for (final Map.Entry<FlowFile, List<Rule>> entry : matchedRules.entrySet()) {
            FlowFile match = entry.getKey();
            final List<Rule> rules = entry.getValue();
            boolean updateWorking = incomingFlowFile.equals(match);
            // execute each matching rule(s)
            match = executeActions(session, context, rules, defaultActions, match, stateInitialAttributes, stateWorkingAttributes);
            if (updateWorking) {
                incomingFlowFile = match;
            }
            if (debugEnabled) {
                logger.debug("Updated attributes for {}; transferring to '{}'", new Object[] { match, REL_SUCCESS.getName() });
            }
            // add the match to the list to transfer
            flowFilesToTransfer.add(match);
        }
    } else {
        // Either we're running without any rules or the FlowFile didn't match any
        incomingFlowFile = executeActions(session, context, null, defaultActions, incomingFlowFile, stateInitialAttributes, stateWorkingAttributes);
        if (debugEnabled) {
            logger.debug("Updated attributes for {}; transferring to '{}'", new Object[] { incomingFlowFile, REL_SUCCESS.getName() });
        }
        // add the flowfile to the list to transfer
        flowFilesToTransfer.add(incomingFlowFile);
    }
    if (stateInitialAttributes != null) {
        try {
            // Able to use "equals()" since we're just checking if the map was modified at all
            if (!stateWorkingAttributes.equals(stateInitialAttributes)) {
                boolean setState = context.getStateManager().replace(stateMap, stateWorkingAttributes, Scope.LOCAL);
                if (!setState) {
                    logger.warn("Failed to update the state after successfully processing {} due to having an old version of the StateMap. This is normally due to multiple threads running at " + "once; transferring to '{}'", new Object[] { incomingFlowFile, REL_FAILED_SET_STATE.getName() });
                    flowFilesToTransfer.remove(incomingFlowFile);
                    if (flowFilesToTransfer.size() > 0) {
                        session.remove(flowFilesToTransfer);
                    }
                    session.transfer(incomingFlowFile, REL_FAILED_SET_STATE);
                    return;
                }
            }
        } catch (IOException e) {
            logger.error("Failed to set the state after successfully processing {} due a failure when setting the state. This is normally due to multiple threads running at " + "once; transferring to '{}'", new Object[] { incomingFlowFile, REL_FAILED_SET_STATE.getName() }, e);
            flowFilesToTransfer.remove(incomingFlowFile);
            if (flowFilesToTransfer.size() > 0) {
                session.remove(flowFilesToTransfer);
            }
            session.transfer(incomingFlowFile, REL_FAILED_SET_STATE);
            context.yield();
            return;
        }
    }
    for (FlowFile toTransfer : flowFilesToTransfer) {
        session.getProvenanceReporter().modifyAttributes(toTransfer);
    }
    session.transfer(flowFilesToTransfer, REL_SUCCESS);
}

Example 20 with StateMap

use of org.apache.nifi.components.state.StateMap in project nifi by apache.

the class TestTailFile method testMigrateFrom100To110.

@Test
public void testMigrateFrom100To110() throws IOException {
    assumeFalse(isWindowsEnvironment());
    runner.setProperty(TailFile.FILENAME, "target/existing-log.txt");
    final MockStateManager stateManager = runner.getStateManager();
    // Before NiFi 1.1.0, TailFile only handles single file
    // and state key doesn't have index in it.
    final Map<String, String> state = new HashMap<>();
    state.put("filename", "target/existing-log.txt");
    // Simulate that it has been tailed up to the 2nd line.
    state.put("checksum", "2279929157");
    state.put("position", "14");
    state.put("timestamp", "1480639134000");
    stateManager.setState(state, Scope.LOCAL);
    runner.run();
    runner.assertAllFlowFilesTransferred(TailFile.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(TailFile.REL_SUCCESS).iterator().next();
    final ByteArrayOutputStream bos = new ByteArrayOutputStream();
    try (final BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(bos))) {
        writer.write("Line 3");
        writer.newLine();
    }
    flowFile.assertContentEquals(bos.toByteArray());
    // The old states should be replaced with new ones.
    final StateMap updatedState = stateManager.getState(Scope.LOCAL);
    assertNull(updatedState.get("filename"));
    assertNull(updatedState.get("checksum"));
    assertNull(updatedState.get("position"));
    assertNull(updatedState.get("timestamp"));
    assertEquals("target/existing-log.txt", updatedState.get("file.0.filename"));
    assertEquals("3380848603", updatedState.get("file.0.checksum"));
    assertEquals("21", updatedState.get("file.0.position"));
    assertNotNull(updatedState.get("file.0.timestamp"));
    // When it runs again, the state is already migrated, so it shouldn't emit any flow files.
    runner.clearTransferState();
    runner.run();
    runner.assertAllFlowFilesTransferred(TailFile.REL_SUCCESS, 0);
}