pvillard31 commented on code in PR #11837:
URL: https://github.com/apache/nifi/pull/11837#discussion_r4401243866
##########
nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java:
##########
@@ -1405,21 +1418,70 @@ public ProvenanceReporter getProvenanceReporter() {
@Override
public void setState(final Map<String, String> state, final Scope scope)
throws IOException {
stateManager.setState(state, scope);
+ pendingStatePatches.remove(scope);
}
+ /**
+ * Returns the stored state with any patches that have not yet been
committed applied. Each patch gives the returned state a new version,
+ * as the framework does, so that a state obtained before a patch no
longer matches in {@link #replaceState(StateMap, Map, Scope)}.
+ */
@Override
public StateMap getState(final Scope scope) throws IOException {
- return stateManager.getState(scope);
+ final StateMap storedState = stateManager.getState(scope);
+ final PendingStatePatches patches = pendingStatePatches.get(scope);
+ if (patches == null) {
+ return storedState;
+ }
+
+ final long storedVersion =
storedState.getStateVersion().map(Long::parseLong).orElse(-1L);
+ return new MockStateMap(storedState.toMap(patches.values),
storedVersion + patches.count);
}
+ /**
+ * Replaces the state if the given old value matches the session's view of
the state. When the session holds patches that have not yet
+ * been committed, the old value must match the stored state with those
patches applied, and a successful replacement discards the patches.
+ */
@Override
public boolean replaceState(final StateMap oldValue, final Map<String,
String> newValue, final Scope scope) throws IOException {
- return stateManager.replace(oldValue, newValue, scope);
+ if (!pendingStatePatches.containsKey(scope)) {
+ return stateManager.replace(oldValue, newValue, scope);
+ }
+
+ final StateMap currentState = getState(scope);
+ if (oldValue == null ||
!oldValue.getStateVersion().equals(currentState.getStateVersion()) ||
!oldValue.toMap().equals(currentState.toMap())) {
+ return false;
+ }
+
+ final boolean replaced =
stateManager.replace(stateManager.getState(scope), newValue, scope);
Review Comment:
This race was already present before the latest commit, but looks cheap to
address: a concurrent update between validation and the second state read could
be overwritten. Could we capture the stored state once and use that same
snapshot for both validation and `replace`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
