org.apache.nifi.components.state.StateMap Java Examples

private void persistState(final Map<String, String> state, final ProcessContext context) {
    try {
        StateMap oldState = context.getStateManager().getState(getStateScope(context));
        Map<String, String> updatedState = new HashMap<String, String>();

        for(String key : oldState.toMap().keySet()) {
            // These states are stored by older version of NiFi, and won't be used anymore.
            // New states have 'file.<index>.' prefix.
            if (TailFileState.StateKeys.CHECKSUM.equals(key)
                    || TailFileState.StateKeys.FILENAME.equals(key)
                    || TailFileState.StateKeys.POSITION.equals(key)
                    || TailFileState.StateKeys.TIMESTAMP.equals(key)) {
                getLogger().info("Removed state {}={} stored by older version of NiFi.", new Object[]{key, oldState.get(key)});
                continue;
            }
            updatedState.put(key, oldState.get(key));
        }

        updatedState.putAll(state);
        context.getStateManager().setState(updatedState, getStateScope(context));
    } catch (final IOException e) {
        getLogger().warn("Failed to store state due to {}; some data may be duplicated on restart of NiFi", new Object[]{e});
    }
}
 

@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);
        }
    }
}
 

@Test
public void testClear() throws IOException {
    final StateProvider provider = getProvider();
    StateMap stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    assertEquals(-1L, stateMap.getVersion());
    assertTrue(stateMap.toMap().isEmpty());

    provider.setState(Collections.singletonMap("testClear", "value"), componentId);

    stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    assertEquals(0, stateMap.getVersion());
    assertEquals("value", stateMap.get("testClear"));

    provider.clear(componentId);

    stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    assertEquals(1L, stateMap.getVersion());
    assertTrue(stateMap.toMap().isEmpty());
}
 

@Test
public void testClusterMonitorActive() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new TestableProcessor(TimeUnit.MINUTES.toMillis(120)));
    runner.setClustered(true);
    runner.setPrimaryNode(false);
    runner.setProperty(MonitorActivity.MONITORING_SCOPE, MonitorActivity.SCOPE_CLUSTER);
    // This has to be very small threshold, otherwise, MonitorActivity skip persisting state.
    runner.setProperty(MonitorActivity.THRESHOLD, "1 ms");

    runner.enqueue("Incoming data");

    runner.run();

    runner.assertAllFlowFilesTransferred(MonitorActivity.REL_SUCCESS);

    final StateMap updatedState = runner.getStateManager().getState(Scope.CLUSTER);
    assertNotNull("Latest timestamp should be persisted", updatedState.get(MonitorActivity.STATE_KEY_LATEST_SUCCESS_TRANSFER));
    // Should be null because COPY_ATTRIBUTES is null.
    assertNull(updatedState.get("key1"));
    assertNull(updatedState.get("key2"));
}
 

@Test
public void testReplaceWithNonExistingValueAndVersionGreaterThanNegativeOne() throws Exception {
    final StateProvider provider = getProvider();
    final StateMap stateMap = new StateMap() {
        @Override
        public long getVersion() {
            return 4;
        }

        @Override
        public String get(String key) {
            return null;
        }

        @Override
        public Map<String, String> toMap() {
            return Collections.emptyMap();
        }
    };

    final Map<String, String> newValue = new HashMap<>();
    newValue.put("value", "value");

    final boolean replaced = provider.replace(stateMap, newValue, componentId);
    assertFalse(replaced);
}
 

private StateMap deserialize(final byte[] data, final int recordVersion, final String componentId) throws IOException {
    try (final ByteArrayInputStream bais = new ByteArrayInputStream(data);
         final DataInputStream dis = new DataInputStream(bais)) {

        final byte encodingVersion = dis.readByte();
        if (encodingVersion > ENCODING_VERSION) {
            throw new IOException("Retrieved a response from ZooKeeper when retrieving state for component with ID " + componentId
                + ", but the response was encoded using the ZooKeeperStateProvider Encoding Version of " + encodingVersion
                + " but this instance can only decode versions up to " + ENCODING_VERSION
                + "; it appears that the state was encoded using a newer version of NiFi than is currently running. This information cannot be decoded.");
        }

        final int numEntries = dis.readInt();
        final Map<String, String> stateValues = new HashMap<>(numEntries);
        for (int i = 0; i < numEntries; i++) {
            final boolean hasKey = dis.readBoolean();
            final String key = hasKey ? dis.readUTF() : null;

            final boolean hasValue = dis.readBoolean();
            final String value = hasValue ? dis.readUTF() : null;
            stateValues.put(key, value);
        }

        return new StandardStateMap(stateValues, recordVersion);
    }
}
 

@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());
}
 

@Test
public void testClusterMonitorActive() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new TestableProcessor(100));
    runner.setClustered(true);
    runner.setPrimaryNode(false);
    runner.setProperty(MonitorActivity.MONITORING_SCOPE, MonitorActivity.SCOPE_CLUSTER);
    // This has to be very small threshold, otherwise, MonitorActivity skip persisting state.
    runner.setProperty(MonitorActivity.THRESHOLD, "1 ms");

    runner.enqueue("Incoming data");

    runner.run();

    runner.assertAllFlowFilesTransferred(MonitorActivity.REL_SUCCESS);

    final StateMap updatedState = runner.getStateManager().getState(Scope.CLUSTER);
    assertNotNull("Latest timestamp should be persisted", updatedState.get(MonitorActivity.STATE_KEY_LATEST_SUCCESS_TRANSFER));
    // Should be null because COPY_ATTRIBUTES is null.
    assertNull(updatedState.get("key1"));
    assertNull(updatedState.get("key2"));
}
 

@Override
public boolean replace(final StateMap oldValue, final Map<String, String> newValue, final String componentId) throws IOException {
    verifyEnabled();

    try {
        setState(newValue, (int) oldValue.getVersion(), componentId, false);
        return true;
    } catch (final NoNodeException nne) {
        return false;
    } catch (final IOException ioe) {
        final Throwable cause = ioe.getCause();
        if (cause != null && cause instanceof KeeperException) {
            final KeeperException ke = (KeeperException) cause;
            if (Code.BADVERSION == ke.code()) {
                return false;
            }
        }

        throw ioe;
    }
}
 

private StateMap deserialize(final byte[] data, final int recordVersion, final String componentId) throws IOException {
    try (final ByteArrayInputStream bais = new ByteArrayInputStream(data);
         final DataInputStream dis = new DataInputStream(bais)) {

        final byte encodingVersion = dis.readByte();
        if (encodingVersion > ENCODING_VERSION) {
            throw new IOException("Retrieved a response from ZooKeeper when retrieving state for component with ID " + componentId
                + ", but the response was encoded using the ZooKeeperStateProvider Encoding Version of " + encodingVersion
                + " but this instance can only decode versions up to " + ENCODING_VERSION
                + "; it appears that the state was encoded using a newer version of NiFi than is currently running. This information cannot be decoded.");
        }

        final int numEntries = dis.readInt();
        final Map<String, String> stateValues = new HashMap<>(numEntries);
        for (int i = 0; i < numEntries; i++) {
            final boolean hasKey = dis.readBoolean();
            final String key = hasKey ? dis.readUTF() : null;

            final boolean hasValue = dis.readBoolean();
            final String value = hasValue ? dis.readUTF() : null;
            stateValues.put(key, value);
        }

        return new StandardStateMap(stateValues, recordVersion);
    }
}
 

@Override
public void clear(final String componentId) throws IOException {
    int attempted = 0;
    boolean updated = false;

    while (!updated && attempted < 20) {
        final StateMap currStateMap = getState(componentId);
        updated = replace(currStateMap, Collections.emptyMap(), componentId, true);

        final String result = updated ? "successful" : "unsuccessful";
        logger.debug("Attempt # {} to clear state for component {} was {}", new Object[] { attempted + 1, componentId, result});

        attempted++;
    }

    if (!updated) {
        throw new IOException("Unable to update state due to concurrent modifications");
    }
}
 

private TimeRange loadState(StateManager stateManager) throws IOException {
    final StateMap stateMap = stateManager.getState(Scope.CLUSTER);

    if (stateMap.getVersion() < 0) {
        getLogger().debug("No previous state found");
        return null;
    }

    final String earliest = stateMap.get(EARLIEST_TIME_KEY);
    final String latest = stateMap.get(LATEST_TIME_KEY);
    getLogger().debug("Loaded state with earliestTime of {} and latestTime of {}", new Object[] {earliest, latest});

    if (StringUtils.isBlank(earliest) && StringUtils.isBlank(latest)) {
        return null;
    } else {
        return new TimeRange(earliest, latest);
    }
}
 

private TimeRange loadState(StateManager stateManager) throws IOException {
    final StateMap stateMap = stateManager.getState(Scope.CLUSTER);

    if (stateMap.getVersion() < 0) {
        getLogger().debug("No previous state found");
        return null;
    }

    final String earliest = stateMap.get(EARLIEST_TIME_KEY);
    final String latest = stateMap.get(LATEST_TIME_KEY);
    getLogger().debug("Loaded state with earliestTime of {} and latestTime of {}", new Object[] {earliest, latest});

    if (StringUtils.isBlank(earliest) && StringUtils.isBlank(latest)) {
        return null;
    } else {
        return new TimeRange(earliest, latest);
    }
}
 

@Test
public void testClear() throws IOException {
    final StateProvider provider = getProvider();
    StateMap stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    assertEquals(-1L, stateMap.getVersion());
    assertTrue(stateMap.toMap().isEmpty());

    provider.setState(Collections.singletonMap("testClear", "value"), componentId);

    stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    assertEquals(0, stateMap.getVersion());
    assertEquals("value", stateMap.get("testClear"));

    provider.clear(componentId);

    stateMap = provider.getState(componentId);
    assertNotNull(stateMap);
    assertEquals(1L, stateMap.getVersion());
    assertTrue(stateMap.toMap().isEmpty());
}
 

@Test
public void testReplaceWithNonExistingValueAndVersionGreaterThanNegativeOne() throws Exception {
    final StateProvider provider = getProvider();
    final StateMap stateMap = new StateMap() {
        @Override
        public long getVersion() {
            return 4;
        }

        @Override
        public String get(String key) {
            return null;
        }

        @Override
        public Map<String, String> toMap() {
            return Collections.emptyMap();
        }
    };

    final Map<String, String> newValue = new HashMap<>();
    newValue.put("value", "value");

    final boolean replaced = provider.replace(stateMap, newValue, componentId);
    assertFalse(replaced);
}
 

/**
 * Demonstrates reading and writing processor state values
 * @throws Exception
 */
@Test
public void testStateJavascript() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new ExecuteScript());
    runner.setValidateExpressionUsage(false);
    runner.setProperty(SCRIPT_ENGINE, "ECMAScript");
    runner.setProperty(ScriptingComponentUtils.SCRIPT_FILE, "src/test/resources/executescript/state/state.js");
    runner.setProperty(ScriptingComponentUtils.MODULES, "src/test/resources/executescript");
    runner.assertValid();

    StateManager stateManager = runner.getStateManager();
    stateManager.clear(Scope.CLUSTER);
    Map<String, String> initialStateValues = new HashMap<>();
    initialStateValues.put("some-state", "foo");
    stateManager.setState(initialStateValues, Scope.CLUSTER);

    runner.enqueue("sample text".getBytes(StandardCharsets.UTF_8));
    runner.run();

    runner.assertAllFlowFilesTransferred("success", 1);
    StateMap resultStateValues = stateManager.getState(Scope.CLUSTER);
    Assert.assertEquals("foobar", resultStateValues.get("some-state"));
}
 

@Override
public synchronized boolean replace(final StateMap oldValue, final Map<String, String> newValue, final Scope scope) throws IOException {
    verifyAnnotation(scope);
    if (scope == Scope.CLUSTER) {
        if (oldValue == clusterStateMap) {
            verifyCanSet(scope);
            clusterStateMap = new MockStateMap(newValue, versionIndex.incrementAndGet());
            return true;
        }

        return false;
    } else {
        if (oldValue == localStateMap) {
            verifyCanSet(scope);
            localStateMap = new MockStateMap(newValue, versionIndex.incrementAndGet());
            return true;
        }

        return false;
    }
}
 

@Override
public boolean replace(final StateMap oldValue, final Map<String, String> newValue, final String componentId) throws IOException {
    verifyEnabled();

    try {
        setState(newValue, (int) oldValue.getVersion(), componentId, false);
        return true;
    } catch (final NoNodeException nne) {
        return false;
    } catch (final IOException ioe) {
        final Throwable cause = ioe.getCause();
        if (cause != null && cause instanceof KeeperException) {
            final KeeperException ke = (KeeperException) cause;
            if (Code.BADVERSION == ke.code()) {
                return false;
            }
        }

        throw ioe;
    }
}
 

/**
 * Audits clearing of state from a Controller Service.
 *
 * @param proceedingJoinPoint join point
 * @param controllerService the controller service
 * @throws java.lang.Throwable ex
 */
@Around("within(org.apache.nifi.web.dao.ComponentStateDAO+) && "
    + "execution(void clearState(org.apache.nifi.controller.service.ControllerServiceNode)) && "
    + "args(controllerService)")
public StateMap clearControllerServiceStateAdvice(ProceedingJoinPoint proceedingJoinPoint, ControllerServiceNode controllerService) throws Throwable {

    // update the controller service state
    final StateMap stateMap = (StateMap) proceedingJoinPoint.proceed();

    // if no exception were thrown, add the clear action...

    // get the current user
    NiFiUser user = NiFiUserUtils.getNiFiUser();

    // ensure the user was found
    if (user != null) {
        Collection<Action> actions = new ArrayList<>();

        // create the controller service details
        FlowChangeExtensionDetails controllerServiceDetails = new FlowChangeExtensionDetails();
        controllerServiceDetails.setType(controllerService.getComponentType());

        // create the clear action
        FlowChangeAction configAction = new FlowChangeAction();
        configAction.setUserIdentity(user.getIdentity());
        configAction.setOperation(Operation.ClearState);
        configAction.setTimestamp(new Date());
        configAction.setSourceId(controllerService.getIdentifier());
        configAction.setSourceName(controllerService.getName());
        configAction.setSourceType(Component.ControllerService);
        configAction.setComponentDetails(controllerServiceDetails);
        actions.add(configAction);

        // record the action
        saveActions(actions, logger);
    }

    return stateMap;
}
 

private Set<String> extractKeys(final StateMap stateMap) {
    Set<String> keys = new HashSet<>();
    for (Map.Entry<String, String>  entry : stateMap.toMap().entrySet()) {
        if (entry.getKey().startsWith(CURRENT_KEY_PREFIX)) {
            keys.add(entry.getValue());
        }
    }
    return keys;
}
 

@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);
}
 

@Test
public void testCreateRoundTrip() throws IOException {
    final String componentId = "1234";

    final StateMapSerDe serde = new StateMapSerDe();
    final Map<String, String> stateValues = new HashMap<>();
    stateValues.put("abc", "xyz");
    stateValues.put("cba", "zyx");
    final StateMap stateMap = new StandardStateMap(stateValues, 3L);
    final StateMapUpdate record = new StateMapUpdate(stateMap, componentId, UpdateType.CREATE);

    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try (final DataOutputStream out = new DataOutputStream(baos)) {
        serde.serializeRecord(record, out);
    }

    final StateMapUpdate update;
    final ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
    try (final DataInputStream in = new DataInputStream(bais)) {
        update = serde.deserializeRecord(in, serde.getVersion());
    }

    assertNotNull(update);
    assertEquals(componentId, update.getComponentId());
    assertEquals(UpdateType.CREATE, update.getUpdateType());
    final StateMap recoveredStateMap = update.getStateMap();

    assertEquals(3L, recoveredStateMap.getVersion());
    assertEquals(stateValues, recoveredStateMap.toMap());
}
 

@Test
public void testReplaceSuccessful() throws IOException {
    final String key = "testReplaceSuccessful";
    final StateProvider provider = getProvider();

    StateMap map = provider.getState(componentId);
    assertNotNull(map);
    assertEquals(-1, map.getVersion());

    assertNotNull(map.toMap());
    assertTrue(map.toMap().isEmpty());
    provider.setState(Collections.singletonMap(key, "value1"), componentId);

    map = provider.getState(componentId);
    assertNotNull(map);
    assertEquals(0, map.getVersion());
    assertEquals("value1", map.get(key));
    assertEquals("value1", map.toMap().get(key));

    final Map<String, String> newMap = new HashMap<>(map.toMap());
    newMap.put(key, "value2");
    assertTrue(provider.replace(map, newMap, componentId));

    map = provider.getState(componentId);
    assertEquals("value2", map.get(key));
    assertEquals(1L, map.getVersion());
}
 

@Override
public ComponentStateDTO getReportingTaskState(final String reportingTaskId) {
    final StateMap clusterState = isClustered() ? reportingTaskDAO.getState(reportingTaskId, Scope.CLUSTER) : null;
    final StateMap localState = reportingTaskDAO.getState(reportingTaskId, Scope.LOCAL);

    // reporting task will be non null as it was already found when getting the state
    final ReportingTaskNode reportingTask = reportingTaskDAO.getReportingTask(reportingTaskId);
    return dtoFactory.createComponentStateDTO(reportingTaskId, reportingTask.getReportingTask().getClass(), localState, clusterState);
}
 

@OnScheduled
public final void updateState(final ProcessContext context) throws IOException {
    final String path = getPath(context);
    final DistributedMapCacheClient client = context.getProperty(DISTRIBUTED_CACHE_SERVICE).asControllerService(DistributedMapCacheClient.class);

    // Check if state already exists for this path. If so, we have already migrated the state.
    final StateMap stateMap = context.getStateManager().getState(getStateScope(context));
    if (stateMap.getVersion() == -1L) {
        try {
            // Migrate state from the old way of managing state (distributed cache service and local file)
            // to the new mechanism (State Manager).
            migrateState(path, client, context.getStateManager(), getStateScope(context));
        } catch (final IOException ioe) {
            throw new IOException("Failed to properly migrate state to State Manager", ioe);
        }
    }

    // When scheduled to run, check if the associated timestamp is null, signifying a clearing of state and reset the internal timestamp
    if (lastListedLatestEntryTimestampMillis != null && stateMap.get(LATEST_LISTED_ENTRY_TIMESTAMP_KEY) == null) {
        getLogger().info("Detected that state was cleared for this component.  Resetting internal values.");
        resetTimeStates();
    }

    if (resetState) {
        context.getStateManager().clear(getStateScope(context));
        resetState = false;
    }
}
 

private StateMap getState(final String componentId, final Scope scope) {
    try {
        final StateManager manager = stateManagerProvider.getStateManager(componentId);
        if (manager == null) {
            throw new ResourceNotFoundException(String.format("State for the specified component %s could not be found.", componentId));
        }

        return manager.getState(scope);
    } catch (final IOException ioe) {
        throw new IllegalStateException(String.format("Unable to get the state for the specified component %s: %s", componentId, ioe), ioe);
    }
}
 

/**
 * Audits clearing of state from a Processor.
 *
 * @param proceedingJoinPoint join point
 * @param reportingTask the reporting task
 * @throws java.lang.Throwable ex
 */
@Around("within(org.apache.nifi.web.dao.ComponentStateDAO+) && "
    + "execution(void clearState(org.apache.nifi.controller.ReportingTaskNode)) && "
    + "args(reportingTask)")
public StateMap clearReportingTaskStateAdvice(ProceedingJoinPoint proceedingJoinPoint, ReportingTaskNode reportingTask) throws Throwable {

    // update the reporting task state
    final StateMap stateMap = (StateMap) proceedingJoinPoint.proceed();

    // if no exception were thrown, add the clear action...

    // get the current user
    NiFiUser user = NiFiUserUtils.getNiFiUser();

    // ensure the user was found
    if (user != null) {
        Collection<Action> actions = new ArrayList<>();

        // create the reporting task details
        FlowChangeExtensionDetails reportingTaskDetails = new FlowChangeExtensionDetails();
        reportingTaskDetails.setType(reportingTask.getReportingTask().getClass().getSimpleName());

        // create the clear action
        FlowChangeAction configAction = new FlowChangeAction();
        configAction.setUserIdentity(user.getIdentity());
        configAction.setOperation(Operation.ClearState);
        configAction.setTimestamp(new Date());
        configAction.setSourceId(reportingTask.getIdentifier());
        configAction.setSourceName(reportingTask.getName());
        configAction.setSourceType(Component.ReportingTask);
        configAction.setComponentDetails(reportingTaskDetails);
        actions.add(configAction);

        // record the action
        saveActions(actions, logger);
    }

    return stateMap;
}
 

private synchronized StateMap retrieveState(final Scope scope) {
    verifyAnnotation(scope);
    if (scope == Scope.CLUSTER) {
        return clusterStateMap;
    } else {
        return localStateMap;
    }
}
 

/**
 * Creates a ComponentStateDTO for the given component and state's.
 *
 * @param componentId component id
 * @param localState local state
 * @param clusterState cluster state
 * @return dto
 */
public ComponentStateDTO createComponentStateDTO(final String componentId, final Class<?> componentClass, final StateMap localState, final StateMap clusterState) {
    final ComponentStateDTO dto = new ComponentStateDTO();
    dto.setComponentId(componentId);
    dto.setStateDescription(getStateDescription(componentClass));
    dto.setLocalState(createStateMapDTO(Scope.LOCAL, localState));
    dto.setClusterState(createStateMapDTO(Scope.CLUSTER, clusterState));
    return dto;
}
 

@Override
public void serializeRecord(final StateMapUpdate record, final DataOutputStream out) throws IOException {
    out.writeUTF(record.getComponentId());
    out.writeUTF(record.getUpdateType().name());
    if (record.getUpdateType() == UpdateType.DELETE) {
        return;
    }

    final StateMap stateMap = record.getStateMap();
    final long recordVersion = stateMap.getVersion();
    out.writeLong(recordVersion);

    final Map<String, String> map = stateMap.toMap();
    out.writeInt(map.size());
    for (final Map.Entry<String, String> entry : map.entrySet()) {
        final boolean hasKey = entry.getKey() != null;
        final boolean hasValue = entry.getValue() != null;
        out.writeBoolean(hasKey);
        if (hasKey) {
            out.writeUTF(entry.getKey());
        }

        out.writeBoolean(hasValue);
        if (hasValue) {
            out.writeUTF(entry.getValue());
        }
    }
}