This is an automated email from the ASF dual-hosted git repository.
pvillard31 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new f492db48779 NIFI-16494 Implemented partial component state updates in
framework (#11837)
f492db48779 is described below
commit f492db48779e630118ee1872f9f95ecad135d146
Author: Mark Payne <[email protected]>
AuthorDate: Wed Oct 28 12:43:37 2026 -0400
NIFI-16494 Implemented partial component state updates in framework (#11837)
---
nifi-docs/src/main/asciidoc/developer-guide.adoc | 6 +-
.../groovyx/flow/ProcessSessionWrap.java | 5 +
.../nifi/components/state/StateProvider.java | 23 +++
.../KubernetesConfigMapStateProviderTest.java | 6 +
.../controller/repository/SessionStateUpdate.java | 75 +++++++++
.../repository/StandardProcessSession.java | 179 ++++++++++++---------
.../WeakHashMapProcessSessionFactory.java | 5 +
.../controller/state/StandardStateManager.java | 6 +
.../manager/StandardStateManagerProvider.java | 7 +
.../TaskTerminationAwareStateManager.java | 6 +
.../repository/BatchingSessionFactory.java | 7 +-
.../nifi/controller/state/StateMapSerDe.java | 65 +++++++-
.../nifi/controller/state/StateMapUpdate.java | 21 +++
.../state/providers/ManagedStateProvider.java | 7 +
.../local/WriteAheadLocalStateProvider.java | 40 ++++-
.../repository/StandardProcessSessionIT.java | 103 +++++++++++-
.../nifi/controller/state/TestStateMapSerDe.java | 60 +++++++
.../state/providers/AbstractTestStateProvider.java | 18 +++
.../local/TestWriteAheadLocalStateProvider.java | 64 ++++++++
.../org/apache/nifi/util/MockProcessSession.java | 81 +++++++++-
.../apache/nifi/util/TestMockProcessSession.java | 42 +++++
.../system/VerifyLocalClusterStateService.java | 5 +-
.../processors/tests/system/MultiKeyState.java | 2 +-
.../tests/system/state/ClusterStateKeyDropIT.java | 4 +
24 files changed, 737 insertions(+), 100 deletions(-)
diff --git a/nifi-docs/src/main/asciidoc/developer-guide.adoc
b/nifi-docs/src/main/asciidoc/developer-guide.adoc
index 8e9fc0ea3b6..608df79a20b 100644
--- a/nifi-docs/src/main/asciidoc/developer-guide.adoc
+++ b/nifi-docs/src/main/asciidoc/developer-guide.adoc
@@ -706,8 +706,10 @@ that the instance is clustered and write the code
accordingly.
==== Storing and Retrieving State
-State is stored using the StateManager's `getState`, `setState`, `replace`,
and `clear` methods. All of these methods
-require that a Scope be provided. It should be noted that the state that is
stored with the Local scope is entirely different
+State is stored using the StateManager's `getState`, `setState`, `patchState`,
`replace`, and `clear` methods. All of these methods
+require that a Scope be provided. The `patchState` method changes only the
keys that it is given: a key with a value is created or
+overwritten, a key with a `null` value is removed, and all other keys are left
as they are. The local state provider writes only
+the changed keys to disk, which makes `patchState` much cheaper than
`setState` when a component holds many keys but changes few. It should be noted
that the state that is stored with the Local scope is entirely different
than state stored with a Cluster scope. If a Processor stores a value with the
key of _My Key_ using the `Scope.CLUSTER` scope,
and then attempts to retrieve the value using the `Scope.LOCAL` scope, the
value retrieved will be `null` (unless a value was
also stored with the same key using the `Scope.CLUSTER` scope). Each
Processor's state is stored in isolation from other
diff --git
a/nifi-extension-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/main/java/org/apache/nifi/processors/groovyx/flow/ProcessSessionWrap.java
b/nifi-extension-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/main/java/org/apache/nifi/processors/groovyx/flow/ProcessSessionWrap.java
index b2727fc1cc2..dc4a0fe4d1b 100644
---
a/nifi-extension-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/main/java/org/apache/nifi/processors/groovyx/flow/ProcessSessionWrap.java
+++
b/nifi-extension-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/main/java/org/apache/nifi/processors/groovyx/flow/ProcessSessionWrap.java
@@ -997,6 +997,11 @@ public abstract class ProcessSessionWrap implements
ProcessSession {
return session.replaceState(oldValue, newValue, scope);
}
+ @Override
+ public void patchState(final Map<String, String> updatedState, final Scope
scope) throws IOException {
+ session.patchState(updatedState, scope);
+ }
+
@Override
public void clearState(final Scope scope) throws IOException {
session.clearState(scope);
diff --git
a/nifi-framework-api/src/main/java/org/apache/nifi/components/state/StateProvider.java
b/nifi-framework-api/src/main/java/org/apache/nifi/components/state/StateProvider.java
index 58fedb29762..347b779c74c 100644
---
a/nifi-framework-api/src/main/java/org/apache/nifi/components/state/StateProvider.java
+++
b/nifi-framework-api/src/main/java/org/apache/nifi/components/state/StateProvider.java
@@ -23,6 +23,7 @@ import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
+import java.util.Objects;
/**
* <p>
@@ -92,6 +93,28 @@ public interface StateProvider extends ConfigurableComponent
{
*/
boolean replace(StateMap oldValue, Map<String, String> newValue, String
componentId) throws IOException;
+ /**
+ * Updates only the given keys of the component's state, leaving all other
keys as they are. Each key with a non-null value is
+ * created or overwritten, and each key with a {@code null} value is
removed. An empty map does nothing: no state is written and the
+ * version does not change.
+ *
+ * @param updatedState a map of state keys to their new values. A key that
is mapped to {@code null} is removed from the state. The map itself must not be
{@code null}.
+ * @param componentId the id of the component for which state is being
updated
+ *
+ * @throws IOException if unable to communicate with the underlying
storage mechanism
+ */
+ default void patchState(final Map<String, String> updatedState, final
String componentId) throws IOException {
+ Objects.requireNonNull(updatedState, "Updated State required");
+ if (updatedState.isEmpty()) {
+ return;
+ }
+
+ StateMap currentState;
+ do {
+ currentState = getState(componentId);
+ } while (!replace(currentState, currentState.toMap(updatedState),
componentId));
+ }
+
/**
* Removes all values from the component's state that is stored using the
given scope
*
diff --git
a/nifi-framework-bundle/nifi-framework-extensions/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
b/nifi-framework-bundle/nifi-framework-extensions/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
index 16a13cd87ae..9c9892ab759 100644
---
a/nifi-framework-bundle/nifi-framework-extensions/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
+++
b/nifi-framework-bundle/nifi-framework-extensions/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
@@ -57,6 +57,7 @@ import static
org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -307,6 +308,11 @@ class KubernetesConfigMapStateProviderTest {
final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
assertTrue(replacedVersion.isPresent());
assertEquals(SECOND_VERSION, replacedVersion.get());
+
+ provider.patchState(Collections.singletonMap(STATE_PROPERTY,
STATE_VALUE), COMPONENT_ID);
+ final StateMap patchedStateMap = provider.getState(COMPONENT_ID);
+ assertEquals(state, patchedStateMap.toMap());
+ assertNotEquals(replacedVersion, patchedStateMap.getStateVersion());
}
@Test
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/SessionStateUpdate.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/SessionStateUpdate.java
new file mode 100644
index 00000000000..26747703d8f
--- /dev/null
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/SessionStateUpdate.java
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.controller.repository;
+
+import org.apache.nifi.components.state.StateMap;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * A state update that a Process Session holds until it is committed. The
update always provides the full state as the session sees it.
+ * When the update was made only by patching state, it also provides the
patches, so that only the patches need to be sent to the
+ * State Manager on commit.
+ */
+final class SessionStateUpdate {
+ private final StateMap stateMap;
+ private final Map<String, String> patches;
+
+ private SessionStateUpdate(final StateMap stateMap, final Map<String,
String> patches) {
+ this.stateMap = stateMap;
+ this.patches = patches;
+ }
+
+ /**
+ * Creates an update that replaces the whole state with the given state
+ *
+ * @param stateMap the new state
+ * @return an update that replaces the whole state
+ */
+ static SessionStateUpdate createWholeStateUpdate(final StateMap stateMap) {
+ return new SessionStateUpdate(stateMap, null);
+ }
+
+ /**
+ * Creates an update that changes only the keys in the given patches
+ *
+ * @param stateMap the full state that results from applying the patches
+ * @param patches the keys to change, where a {@code null} value means
that the key is removed
+ * @return an update that changes only the keys in the given patches
+ */
+ static SessionStateUpdate createPatchUpdate(final StateMap stateMap, final
Map<String, String> patches) {
+ // Map.copyOf does not allow null values, which are needed to indicate
removed keys
+ return new SessionStateUpdate(stateMap,
Collections.unmodifiableMap(new HashMap<>(patches)));
+ }
+
+ /**
+ * @return the full state as the session sees it once this update is
applied
+ */
+ StateMap getStateMap() {
+ return stateMap;
+ }
+
+ /**
+ * @return an unmodifiable map of the keys to change, where a {@code null}
value means that the key is removed, or {@code null} if this update
+ * replaces the whole state
+ */
+ Map<String, String> getPatches() {
+ return patches;
+ }
+}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
index 0a6d214fa06..1175b636d44 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
@@ -140,8 +140,6 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
private static final String CONNECTION_COMPONENT_TYPE = "Connection";
private static final AtomicLong idGenerator = new AtomicLong(0L);
private static final AtomicLong enqueuedIndex = new AtomicLong(0L);
- private static final StateMap EMPTY_STATE_MAP = new
StandardStateMap(Collections.emptyMap(), Optional.empty());
-
// determines how many things must be transferred, removed, modified in
order to avoid logging the FlowFile ID's on commit/rollback
public static final int VERBOSE_LOG_THRESHOLD = 10;
public static final String DEFAULT_FLOWFILE_PATH = "./";
@@ -205,8 +203,8 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
private Checkpoint checkpoint = null;
private final ContentClaimWriteCache claimCache;
- private StateMap localState;
- private StateMap clusterState;
+ private SessionStateUpdate localStateUpdate;
+ private SessionStateUpdate clusterStateUpdate;
private final String retryAttribute;
private final FlowFileLinkage flowFileLinkage = new FlowFileLinkage();
@@ -719,40 +717,9 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
LOG.debug("{}", timingInfo);
}
- // Update local state
final StateManager stateManager = context.getStateManager();
- if (checkpoint.localState != null) {
- try {
- final StateMap stateMap =
stateManager.getState(Scope.LOCAL);
- final Optional<String> stateVersion =
stateMap.getStateVersion();
- if
(!stateVersion.equals(checkpoint.localState.getStateVersion())) {
- LOG.debug("Updating State Manager's Local State");
- stateManager.setState(checkpoint.localState.toMap(),
Scope.LOCAL);
- } else {
- LOG.debug("Will not update State Manager's Local State
because the State Manager reports the latest version as {}, which is newer than
the session's known version of {}.",
- stateVersion,
checkpoint.localState.getStateVersion());
- }
- } catch (final Exception e) {
- LOG.warn("Failed to update Local State for {}. If NiFi is
restarted before the state is able to be updated, it could result in data
duplication.", connectableDescription, e);
- }
- }
-
- // Update cluster state
- if (checkpoint.clusterState != null) {
- try {
- final StateMap stateMap =
stateManager.getState(Scope.CLUSTER);
- final Optional<String> stateVersion =
stateMap.getStateVersion();
- if
(!stateVersion.equals(checkpoint.clusterState.getStateVersion())) {
- LOG.debug("Updating State Manager's Cluster State");
- stateManager.setState(checkpoint.clusterState.toMap(),
Scope.CLUSTER);
- } else {
- LOG.debug("Will not update State Manager's Cluster
State because the State Manager reports the latest version as {}, which is
newer than the session's known version of {}.",
- stateVersion,
checkpoint.clusterState.getStateVersion());
- }
- } catch (final Exception e) {
- LOG.warn("Failed to update Cluster State for {}. If NiFi
is restarted before the state is able to be updated, it could result in data
duplication.", connectableDescription, e);
- }
- }
+ commitStateUpdate(stateManager, checkpoint.localStateUpdate,
Scope.LOCAL);
+ commitStateUpdate(stateManager, checkpoint.clusterStateUpdate,
Scope.CLUSTER);
// Acknowledge records in order to update counts for incoming
connections' queues
acknowledgeRecords();
@@ -781,6 +748,40 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
}
}
+ /**
+ * Sends the session's state update to the State Manager. An update that
was made only by patching state always sends its patches, because
+ * the patches change only their own keys and so apply on top of whatever
state is stored at the time of the commit. Any other update sends
+ * the whole state, unless the State Manager already reports the session's
version.
+ */
+ private void commitStateUpdate(final StateManager stateManager, final
SessionStateUpdate stateUpdate, final Scope scope) {
+ if (stateUpdate == null) {
+ return;
+ }
+
+ final String scopeName = scope == Scope.LOCAL ? "Local" : "Cluster";
+ try {
+ final Map<String, String> patches = stateUpdate.getPatches();
+ if (patches != null) {
+ LOG.debug("Patching {} keys in State Manager's {} State",
patches.size(), scopeName);
+ stateManager.patchState(patches, scope);
+ return;
+ }
+
+ final Optional<String> stateVersion =
stateManager.getState(scope).getStateVersion();
+ final Optional<String> sessionStateVersion =
stateUpdate.getStateMap().getStateVersion();
+ if (stateVersion.equals(sessionStateVersion)) {
+ LOG.debug("Will not update State Manager's {} State because
the State Manager reports the latest version as {}, which is newer than the
session's known version of {}.",
+ scopeName, stateVersion, sessionStateVersion);
+ return;
+ }
+
+ LOG.debug("Updating State Manager's {} State", scopeName);
+ stateManager.setState(stateUpdate.getStateMap().toMap(), scope);
+ } catch (final Exception e) {
+ LOG.warn("Failed to update {} State for {}. If NiFi is restarted
before the state is able to be updated, it could result in data duplication.",
scopeName, connectableDescription, e);
+ }
+ }
+
private void updateEventRepository(final Checkpoint checkpoint) {
try {
// update event repository
@@ -1342,10 +1343,10 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
LOG.warn("{} Attempted to close Output Stream for {} due to
session rollback but close failed", this, this.connectableDescription, e1);
}
- if (localState != null || clusterState != null) {
+ if (localStateUpdate != null || clusterStateUpdate != null) {
LOG.debug("Rolling back session that has state stored. This state
will not be updated.");
}
- if (rollbackCheckpoint && checkpoint != null && (checkpoint.localState
!= null || checkpoint.clusterState != null)) {
+ if (rollbackCheckpoint && checkpoint != null &&
(checkpoint.localStateUpdate != null || checkpoint.clusterStateUpdate != null))
{
LOG.debug("Rolling back checkpoint that has state stored. This
state will not be updated.");
}
@@ -1565,8 +1566,8 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
forkEventBuilders.clear();
provenanceReporter.clear();
- localState = null;
- clusterState = null;
+ localStateUpdate = null;
+ clusterStateUpdate = null;
processingStartTime = System.nanoTime();
}
@@ -4026,37 +4027,38 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
}
private void setState(final StateMap stateMap, final Scope scope) {
+ setStateUpdate(SessionStateUpdate.createWholeStateUpdate(stateMap),
scope);
+ }
+
+ private void setStateUpdate(final SessionStateUpdate stateUpdate, final
Scope scope) {
if (scope == Scope.LOCAL) {
- localState = stateMap;
+ localStateUpdate = stateUpdate;
} else {
- clusterState = stateMap;
+ clusterStateUpdate = stateUpdate;
}
}
@Override
public StateMap getState(final Scope scope) throws IOException {
- if (scope == Scope.LOCAL) {
- if (localState != null) {
- return localState;
- }
-
- if (checkpoint != null && checkpoint.localState != null) {
- return checkpoint.localState;
- }
-
- // If no state is held locally, get it from the State Manager.
- return context.getStateManager().getState(scope);
+ final SessionStateUpdate pendingUpdate = getPendingStateUpdate(scope);
+ if (pendingUpdate != null) {
+ return pendingUpdate.getStateMap();
}
- if (clusterState != null) {
- return clusterState;
- }
+ // If no state is held locally, get it from the State Manager.
+ return context.getStateManager().getState(scope);
+ }
- if (checkpoint != null && checkpoint.clusterState != null) {
- return checkpoint.clusterState;
+ /**
+ * @return the state update held by this session for the given scope, or
else the one held by the checkpoint, or {@code null} if neither holds one
+ */
+ private SessionStateUpdate getPendingStateUpdate(final Scope scope) {
+ final SessionStateUpdate sessionUpdate = scope == Scope.LOCAL ?
localStateUpdate : clusterStateUpdate;
+ if (sessionUpdate != null || checkpoint == null) {
+ return sessionUpdate;
}
- return context.getStateManager().getState(scope);
+ return scope == Scope.LOCAL ? checkpoint.localStateUpdate :
checkpoint.clusterStateUpdate;
}
@Override
@@ -4082,9 +4084,42 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
return false;
}
+ /**
+ * Applies the given changes to the session's view of the state. When the
session commits, only the patches are sent to the
+ * State Manager, unless the session or its checkpoint already holds a
whole-state update from {@link #setState(Map, Scope)},
+ * {@link #replaceState(StateMap, Map, Scope)}, or {@link
#clearState(Scope)}, in which case the whole state is sent.
+ */
+ @Override
+ public void patchState(final Map<String, String> updatedState, final Scope
scope) throws IOException {
+ Objects.requireNonNull(updatedState, "Updated State required");
+ if (updatedState.isEmpty()) {
+ return;
+ }
+
+ final StateMap currentState = getState(scope);
+ final Map<String, String> patchedValues =
currentState.toMap(updatedState);
+ final String version =
currentState.getStateVersion().map(this::getIncrementedVersion).orElse(INITIAL_VERSION);
+ final StateMap patchedStateMap = new StandardStateMap(patchedValues,
Optional.of(version));
+
+ final SessionStateUpdate pendingUpdate = getPendingStateUpdate(scope);
+ if (pendingUpdate == null) {
+
setStateUpdate(SessionStateUpdate.createPatchUpdate(patchedStateMap,
updatedState), scope);
+ return;
+ }
+
+ if (pendingUpdate.getPatches() == null) {
+ setState(patchedStateMap, scope);
+ return;
+ }
+
+ final Map<String, String> combinedPatches = new
HashMap<>(pendingUpdate.getPatches());
+ combinedPatches.putAll(updatedState);
+ setStateUpdate(SessionStateUpdate.createPatchUpdate(patchedStateMap,
combinedPatches), scope);
+ }
+
@Override
- public void clearState(final Scope scope) {
- setState(EMPTY_STATE_MAP, scope);
+ public void clearState(final Scope scope) throws IOException {
+ setState(Collections.emptyMap(), scope);
}
@Override
@@ -4153,8 +4188,8 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
private long bytesSent = 0L;
private boolean initialized = false;
- private StateMap localState;
- private StateMap clusterState;
+ private SessionStateUpdate localStateUpdate;
+ private SessionStateUpdate clusterStateUpdate;
private void initializeForCopy() {
if (initialized) {
@@ -4228,11 +4263,11 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
this.flowFilesSent = session.provenanceReporter.getFlowFilesSent();
this.bytesSent = session.provenanceReporter.getBytesSent();
- if (session.localState != null) {
- this.localState = session.localState;
+ if (session.localStateUpdate != null) {
+ this.localStateUpdate = session.localStateUpdate;
}
- if (session.clusterState != null) {
- this.clusterState = session.clusterState;
+ if (session.clusterStateUpdate != null) {
+ this.clusterStateUpdate = session.clusterStateUpdate;
}
}
@@ -4283,11 +4318,11 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
this.flowFilesSent +=
session.provenanceReporter.getFlowFilesSent();
this.bytesSent += session.provenanceReporter.getBytesSent();
- if (session.localState != null) {
- this.localState = session.localState;
+ if (session.localStateUpdate != null) {
+ this.localStateUpdate = session.localStateUpdate;
}
- if (session.clusterState != null) {
- this.clusterState = session.clusterState;
+ if (session.clusterStateUpdate != null) {
+ this.clusterStateUpdate = session.clusterStateUpdate;
}
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/WeakHashMapProcessSessionFactory.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/WeakHashMapProcessSessionFactory.java
index 325d1749035..c12d8d1a7a4 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/WeakHashMapProcessSessionFactory.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/WeakHashMapProcessSessionFactory.java
@@ -345,6 +345,11 @@ public class WeakHashMapProcessSessionFactory implements
ActiveProcessSessionFac
return delegate.replaceState(oldValue, newValue, scope);
}
+ @Override
+ public void patchState(final Map<String, String> updatedState, final
Scope scope) throws IOException {
+ delegate.patchState(updatedState, scope);
+ }
+
@Override
public void clearState(final Scope scope) throws IOException {
delegate.clearState(scope);
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateManager.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateManager.java
index a3ce131a9a8..5ecf0d7e566 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateManager.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateManager.java
@@ -87,6 +87,12 @@ public class StandardStateManager implements StateManager {
getProvider(scope).setState(state, componentId);
}
+ @Override
+ public void patchState(final Map<String, String> updatedState, final Scope
scope) throws IOException {
+ getLogger(componentId).debug("Patching {} State with {}", scope,
updatedState);
+ getProvider(scope).patchState(updatedState, componentId);
+ }
+
@Override
public void clear(final Scope scope) throws IOException {
getLogger(componentId).debug("Clearing {} State", scope);
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/manager/StandardStateManagerProvider.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/manager/StandardStateManagerProvider.java
index ac624224671..e72f1ace837 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/manager/StandardStateManagerProvider.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/manager/StandardStateManagerProvider.java
@@ -459,6 +459,13 @@ public class StandardStateManagerProvider implements
StateManagerProvider {
}
}
+ @Override
+ public void patchState(Map<String, String> updatedState, String
componentId) throws IOException {
+ try (final NarCloseable ignored =
NarCloseable.withNarLoader()) {
+ stateProvider.patchState(updatedState, componentId);
+ }
+ }
+
@Override
public void clear(String componentId) throws IOException {
try (final NarCloseable ignored =
NarCloseable.withNarLoader()) {
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/lifecycle/TaskTerminationAwareStateManager.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/lifecycle/TaskTerminationAwareStateManager.java
index e846628fca4..9948c9a723d 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/lifecycle/TaskTerminationAwareStateManager.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/lifecycle/TaskTerminationAwareStateManager.java
@@ -63,4 +63,10 @@ public class TaskTerminationAwareStateManager implements
StateManager {
verifyNotTerminated();
stateManager.setState(state, scope);
}
+
+ @Override
+ public void patchState(final Map<String, String> updatedState, final Scope
scope) throws IOException {
+ verifyNotTerminated();
+ stateManager.patchState(updatedState, scope);
+ }
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/BatchingSessionFactory.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/BatchingSessionFactory.java
index b3a373d85d3..0a055613a11 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/BatchingSessionFactory.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/BatchingSessionFactory.java
@@ -312,7 +312,12 @@ public class BatchingSessionFactory implements
ProcessSessionFactory {
}
@Override
- public void clearState(final Scope scope) {
+ public void patchState(final Map<String, String> updatedState, final
Scope scope) throws IOException {
+ session.patchState(updatedState, scope);
+ }
+
+ @Override
+ public void clearState(final Scope scope) throws IOException {
session.clearState(scope);
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapSerDe.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapSerDe.java
index 012e7ed9a08..ffa688da1d9 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapSerDe.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapSerDe.java
@@ -28,14 +28,33 @@ import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
+/**
+ * Serializes component state for the write-ahead log. Full records, which are
used for snapshots, always hold the whole state map.
+ * Starting with version 1, an edit in the journal holds either the whole
state map or only the patches. An edit that holds
+ * only the patches is applied on top of the component's previous record when
the journal is recovered. Version 0 edits always
+ * hold the whole state map, in the same layout as full records.
+ */
public class StateMapSerDe implements SerDe<StateMapUpdate> {
private static final long EMPTY_VERSION = -1;
- private static final int VERSION = 0;
+ private static final int PATCHES_VERSION = 1;
+ private static final int VERSION = PATCHES_VERSION;
@Override
public void serializeEdit(final StateMapUpdate previousRecordState, final
StateMapUpdate newRecordState, final DataOutputStream out) throws IOException {
- serializeRecord(newRecordState, out);
+ out.writeUTF(newRecordState.getComponentId());
+ out.writeUTF(newRecordState.getUpdateType().name());
+ if (newRecordState.getUpdateType() == UpdateType.DELETE) {
+ return;
+ }
+
+ // Only the patches can be written if the journal is able to recover
the previous record to apply them to
+ final boolean writePatchesOnly = newRecordState.getPatches() != null
&& previousRecordState != null && previousRecordState.getStateMap() != null;
+ out.writeBoolean(writePatchesOnly);
+
+ final StateMap stateMap = newRecordState.getStateMap();
+
out.writeLong(stateMap.getStateVersion().map(Long::parseLong).orElse(EMPTY_VERSION));
+ writeEntries(writePatchesOnly ? newRecordState.getPatches() :
stateMap.toMap(), out);
}
@Override
@@ -49,8 +68,10 @@ public class StateMapSerDe implements SerDe<StateMapUpdate> {
final StateMap stateMap = record.getStateMap();
final long recordVersion =
stateMap.getStateVersion().map(Long::parseLong).orElse(EMPTY_VERSION);
out.writeLong(recordVersion);
+ writeEntries(stateMap.toMap(), out);
+ }
- final Map<String, String> map = stateMap.toMap();
+ private void writeEntries(final Map<String, String> map, final
DataOutputStream out) throws IOException {
out.writeInt(map.size());
for (final Map.Entry<String, String> entry : map.entrySet()) {
final boolean hasKey = entry.getKey() != null;
@@ -69,7 +90,30 @@ public class StateMapSerDe implements SerDe<StateMapUpdate> {
@Override
public StateMapUpdate deserializeEdit(final DataInputStream in, final
Map<Object, StateMapUpdate> currentRecordStates, final int version) throws
IOException {
- return deserializeRecord(in, version);
+ if (version < PATCHES_VERSION) {
+ return deserializeRecord(in, version);
+ }
+
+ final String componentId = in.readUTF();
+ final UpdateType updateType = UpdateType.valueOf(in.readUTF());
+ if (updateType == UpdateType.DELETE) {
+ return new StateMapUpdate(null, componentId, updateType);
+ }
+
+ final boolean patchesOnly = in.readBoolean();
+ final String stateVersion = String.valueOf(in.readLong());
+ final Map<String, String> entries = readEntries(in);
+ if (!patchesOnly) {
+ return new StateMapUpdate(new StandardStateMap(entries,
Optional.of(stateVersion)), componentId, updateType);
+ }
+
+ final StateMapUpdate previousRecord =
currentRecordStates.get(componentId);
+ if (previousRecord == null || previousRecord.getStateMap() == null) {
+ throw new IOException("Journal holds state patches for component "
+ componentId + " but no previous state was recovered for that component");
+ }
+
+ final Map<String, String> patchedState =
previousRecord.getStateMap().toMap(entries);
+ return new StateMapUpdate(new StandardStateMap(patchedState,
Optional.of(stateVersion)), componentId, updateType);
}
@Override
@@ -82,18 +126,23 @@ public class StateMapSerDe implements
SerDe<StateMapUpdate> {
}
final long recordVersion = in.readLong();
+ final Map<String, String> stateValues = readEntries(in);
+ final String stateVersion = String.valueOf(recordVersion);
+ return new StateMapUpdate(new StandardStateMap(stateValues,
Optional.of(stateVersion)), componentId, updateType);
+ }
+
+ private Map<String, String> readEntries(final DataInputStream in) throws
IOException {
final int numEntries = in.readInt();
- final Map<String, String> stateValues = new HashMap<>(numEntries);
+ final Map<String, String> entries = new HashMap<>(numEntries);
for (int i = 0; i < numEntries; i++) {
final boolean hasKey = in.readBoolean();
final String key = hasKey ? in.readUTF() : null;
final boolean hasValue = in.readBoolean();
final String value = hasValue ? in.readUTF() : null;
- stateValues.put(key, value);
+ entries.put(key, value);
}
- final String stateVersion = String.valueOf(recordVersion);
- return new StateMapUpdate(new StandardStateMap(stateValues,
Optional.of(stateVersion)), componentId, updateType);
+ return entries;
}
@Override
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapUpdate.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapUpdate.java
index 9eb478bc91c..9318ff18ca5 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapUpdate.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/StateMapUpdate.java
@@ -20,15 +20,29 @@ package org.apache.nifi.controller.state;
import org.apache.nifi.components.state.StateMap;
import org.wali.UpdateType;
+import java.util.Map;
+
public class StateMapUpdate {
private final StateMap stateMap;
private final String componentId;
private final UpdateType updateType;
+ private final Map<String, String> patches;
public StateMapUpdate(final StateMap stateMap, final String componentId,
final UpdateType updateType) {
+ this(stateMap, componentId, updateType, null);
+ }
+
+ /**
+ * @param stateMap the full state of the component once the update is
applied
+ * @param componentId the identifier of the component
+ * @param updateType the type of update
+ * @param patches the keys that were changed, where a {@code null} value
means that the key was removed, or {@code null} if the update replaces the
whole state
+ */
+ public StateMapUpdate(final StateMap stateMap, final String componentId,
final UpdateType updateType, final Map<String, String> patches) {
this.stateMap = stateMap;
this.componentId = componentId;
this.updateType = updateType;
+ this.patches = patches;
}
public StateMap getStateMap() {
@@ -42,4 +56,11 @@ public class StateMapUpdate {
public UpdateType getUpdateType() {
return updateType;
}
+
+ /**
+ * @return the keys that were changed, where a {@code null} value means
that the key was removed, or {@code null} if the update replaces the whole state
+ */
+ public Map<String, String> getPatches() {
+ return patches;
+ }
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/ManagedStateProvider.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/ManagedStateProvider.java
index ab196b29e06..3b6bcacad68 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/ManagedStateProvider.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/ManagedStateProvider.java
@@ -140,6 +140,13 @@ public class ManagedStateProvider implements StateProvider
{
}
}
+ @Override
+ public void patchState(final Map<String, String> updatedState, final
String componentId) throws IOException {
+ try (final NarCloseable ignored = NarCloseable.withNarLoader()) {
+ stateProvider.patchState(updatedState, componentId);
+ }
+ }
+
@Override
public void clear(final String componentId) throws IOException {
try (final NarCloseable ignored = NarCloseable.withNarLoader()) {
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
index 1eb68ac2595..1babf959422 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
@@ -44,6 +44,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
@@ -230,6 +231,14 @@ public class WriteAheadLocalStateProvider extends
AbstractStateProvider {
return getProvider(componentId).replace(oldValue, newValue);
}
+ /**
+ * Writes only the patches to the write-ahead log's journal. Snapshots
taken at each checkpoint still hold the whole state.
+ */
+ @Override
+ public void patchState(final Map<String, String> updatedState, final
String componentId) throws IOException {
+ getProvider(componentId).patchState(updatedState);
+ }
+
@Override
public void clear(final String componentId) throws IOException {
getProvider(componentId).clear();
@@ -280,10 +289,11 @@ public class WriteAheadLocalStateProvider extends
AbstractStateProvider {
// Additionally, the implementation of WriteAheadRepository that we
are using requires that only a single thread update the
// repository at a time for a record with the same key. I.e., many
threads can update the repository at once, as long as they
// are not updating the repository with records that have the same
identifier.
+ // The local stateMap variable is updated only after the WAL update
succeeds, so that it always matches the WAL. A journal entry
+ // that holds only patches is applied to the previous state in the WAL
when the journal is recovered.
public synchronized void setState(final Map<String, String> state)
throws IOException {
- stateMap = new StandardStateMap(state,
Optional.of(getIncrementedVersion()));
- final StateMapUpdate updateRecord = new StateMapUpdate(stateMap,
componentId, UpdateType.UPDATE);
- wal.update(Collections.singleton(updateRecord), alwaysSync);
+ final StateMap updatedStateMap = new StandardStateMap(state,
Optional.of(getIncrementedVersion()));
+ writeUpdate(new StateMapUpdate(updatedStateMap, componentId,
UpdateType.UPDATE));
}
// see above explanation as to why this method is synchronized.
@@ -292,16 +302,30 @@ public class WriteAheadLocalStateProvider extends
AbstractStateProvider {
return false;
}
- stateMap = new StandardStateMap(new HashMap<>(newValue),
Optional.of(getIncrementedVersion()));
- final StateMapUpdate updateRecord = new StateMapUpdate(stateMap,
componentId, UpdateType.UPDATE);
- wal.update(Collections.singleton(updateRecord), alwaysSync);
+ final StateMap updatedStateMap = new StandardStateMap(new
HashMap<>(newValue), Optional.of(getIncrementedVersion()));
+ writeUpdate(new StateMapUpdate(updatedStateMap, componentId,
UpdateType.UPDATE));
return true;
}
+ // see above explanation as to why this method is synchronized.
+ public synchronized void patchState(final Map<String, String>
updatedState) throws IOException {
+ Objects.requireNonNull(updatedState, "Updated State required");
+ if (updatedState.isEmpty()) {
+ return;
+ }
+
+ final StateMap updatedStateMap = new
StandardStateMap(stateMap.toMap(updatedState),
Optional.of(getIncrementedVersion()));
+ writeUpdate(new StateMapUpdate(updatedStateMap, componentId,
UpdateType.UPDATE, Collections.unmodifiableMap(new HashMap<>(updatedState))));
+ }
+
public synchronized void clear() throws IOException {
- stateMap = new StandardStateMap(null,
Optional.of(getIncrementedVersion()));
- final StateMapUpdate update = new StateMapUpdate(stateMap,
componentId, UpdateType.UPDATE);
+ final StateMap updatedStateMap = new StandardStateMap(null,
Optional.of(getIncrementedVersion()));
+ writeUpdate(new StateMapUpdate(updatedStateMap, componentId,
UpdateType.UPDATE));
+ }
+
+ private void writeUpdate(final StateMapUpdate update) throws
IOException {
wal.update(Collections.singleton(update), alwaysSync);
+ stateMap = update.getStateMap();
}
private String getIncrementedVersion() {
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
index 4f1129d72ca..25261c982a0 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
@@ -118,9 +118,11 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.notNull;
+import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@@ -229,7 +231,7 @@ public class StandardProcessSessionIT {
contentRepo.initialize(new StandardContentRepositoryContext(new
StandardResourceClaimManager(), EventReporter.NO_OP));
flowFileRepo = new MockFlowFileRepository(contentRepo);
- stateManager = new MockStateManager(connectable);
+ stateManager = spy(new MockStateManager(connectable));
stateManager.setIgnoreAnnotations(true);
context = new StandardRepositoryContext(connectable, new
AtomicLong(0L), contentRepo, flowFileRepo, flowFileEventRepository,
@@ -2494,6 +2496,12 @@ public class StandardProcessSessionIT {
stateManager.assertStateEquals("abc", "123", Scope.LOCAL);
stateManager.assertStateNotSet(Scope.CLUSTER);
+ // Clearing the state and then patching it replaces the stored state
with the patched values
+ session.clearState(Scope.LOCAL);
+ session.patchState(Map.of("new", "2"), Scope.LOCAL);
+ session.commit();
+ stateManager.assertStateEquals(Map.of("new", "2"), Scope.LOCAL);
+
stateManager.reset();
session.setState(Collections.singletonMap("abc", "123"),
Scope.CLUSTER);
@@ -2502,6 +2510,43 @@ public class StandardProcessSessionIT {
session.commit();
stateManager.assertStateEquals("abc", "123", Scope.CLUSTER);
stateManager.assertStateNotSet(Scope.LOCAL);
+
+ stateManager.reset();
+ session.patchState(Collections.emptyMap(), Scope.LOCAL);
+ session.commit();
+ stateManager.assertStateNotSet();
+
+ // Patches made in one session are combined and sent to the State
Manager as a single patch
+ session.setState(Collections.singletonMap("abc", "123"), Scope.LOCAL);
+ session.commit();
+ clearInvocations(stateManager);
+
+ session.patchState(Map.of("def", "456"), Scope.LOCAL);
+ final Map<String, String> secondPatch = new HashMap<>();
+ secondPatch.put("def", "789");
+ secondPatch.put("abc", null);
+ session.patchState(secondPatch, Scope.LOCAL);
+ stateManager.assertStateEquals(Map.of("abc", "123"), Scope.LOCAL);
+
+ session.commit();
+ verify(stateManager).patchState(secondPatch, Scope.LOCAL);
+ verify(stateManager, never()).setState(any(), any());
+ stateManager.assertStateEquals(Map.of("def", "789"), Scope.LOCAL);
+
+ // A patch is applied on commit even if another writer updated the
state after the session retrieved it
+ session.patchState(Map.of("ghi", "1"), Scope.LOCAL);
+ stateManager.setState(Map.of("def", "789", "other", "1"), Scope.LOCAL);
+ session.commit();
+ stateManager.assertStateEquals(Map.of("def", "789", "other", "1",
"ghi", "1"), Scope.LOCAL);
+
+ // A whole-state update made after a patch sends the whole state
+ clearInvocations(stateManager);
+ session.patchState(Map.of("ghi", "1"), Scope.LOCAL);
+ session.setState(Map.of("xyz", "1"), Scope.LOCAL);
+ session.commit();
+ verify(stateManager).setState(Map.of("xyz", "1"), Scope.LOCAL);
+ verify(stateManager, never()).patchState(any(), any());
+ stateManager.assertStateEquals(Map.of("xyz", "1"), Scope.LOCAL);
}
@Test
@@ -2554,16 +2599,33 @@ public class StandardProcessSessionIT {
retrieved = session.getState(Scope.LOCAL);
assertNotNull(retrieved);
assertTrue(retrieved.getStateVersion().isPresent());
+
+ stateManager.reset();
+ final Map<String, String> patch = new HashMap<>();
+ patch.put("x", "1");
+ patch.put("y", null);
+ session.patchState(patch, Scope.LOCAL);
+ retrieved = session.getState(Scope.LOCAL);
+ assertEquals(Map.of("x", "1"), retrieved.toMap());
+ assertTrue(retrieved.getStateVersion().isPresent());
+ stateManager.assertStateNotSet();
+
+ session.patchState(Map.of("z", "2"), Scope.LOCAL);
+ assertEquals(Map.of("x", "1", "z", "2"),
session.getState(Scope.LOCAL).toMap());
+ session.commit();
+ stateManager.assertStateEquals(Map.of("x", "1", "z", "2"),
Scope.LOCAL);
}
@Test
public void testRollbackDoesNotStoreState() throws IOException {
session.setState(Collections.singletonMap("abc", "1"), Scope.LOCAL);
+ session.patchState(Map.of("def", "2"), Scope.CLUSTER);
session.rollback();
stateManager.assertStateNotSet();
session.commit();
stateManager.assertStateNotSet();
+ verify(stateManager, never()).patchState(any(), any());
}
@Test
@@ -2575,6 +2637,17 @@ public class StandardProcessSessionIT {
session.commit();
stateManager.assertStateEquals("abc", "1", Scope.LOCAL);
+
+ // Patches made before and after a checkpoint are sent together as one
patch
+ stateManager.reset();
+ clearInvocations(stateManager);
+ session.patchState(Map.of("abc", "1"), Scope.LOCAL);
+ session.checkpoint();
+ session.patchState(Map.of("def", "2"), Scope.LOCAL);
+ session.commit();
+
+ verify(stateManager).patchState(Map.of("abc", "1", "def", "2"),
Scope.LOCAL);
+ stateManager.assertStateEquals(Map.of("abc", "1", "def", "2"),
Scope.LOCAL);
}
@Test
@@ -2587,6 +2660,18 @@ public class StandardProcessSessionIT {
session.commit();
stateManager.assertStateEquals("abc", "2", Scope.LOCAL);
+
+ // A patch made after a checkpointed whole-state update sends the
whole state
+ stateManager.reset();
+ clearInvocations(stateManager);
+ session.setState(Map.of("abc", "1"), Scope.LOCAL);
+ session.checkpoint();
+ session.patchState(Map.of("def", "2"), Scope.LOCAL);
+ session.commit();
+
+ verify(stateManager).setState(Map.of("abc", "1", "def", "2"),
Scope.LOCAL);
+ verify(stateManager, never()).patchState(any(), any());
+ stateManager.assertStateEquals(Map.of("abc", "1", "def", "2"),
Scope.LOCAL);
}
@Test
@@ -2597,12 +2682,26 @@ public class StandardProcessSessionIT {
session.checkpoint();
session.setState(Collections.singletonMap("abc", "2"), Scope.LOCAL);
+ session.patchState(Map.of("def", "2"), Scope.LOCAL);
session.rollback();
stateManager.assertStateNotSet();
session.commit();
- stateManager.assertStateEquals("abc", "1", Scope.LOCAL);
+ stateManager.assertStateEquals(Map.of("abc", "1"), Scope.LOCAL);
+ verify(stateManager, never()).patchState(any(), any());
+
+ // Rolling back discards a patch made after the checkpoint and keeps
the checkpointed patch
+ stateManager.reset();
+ clearInvocations(stateManager);
+ session.patchState(Map.of("abc", "1"), Scope.LOCAL);
+ session.checkpoint();
+ session.patchState(Map.of("def", "2"), Scope.LOCAL);
+ session.rollback();
+ session.commit();
+
+ verify(stateManager).patchState(Map.of("abc", "1"), Scope.LOCAL);
+ stateManager.assertStateEquals(Map.of("abc", "1"), Scope.LOCAL);
}
@Test
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/TestStateMapSerDe.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/TestStateMapSerDe.java
index 206e3fb1e51..f4ad606d229 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/TestStateMapSerDe.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/TestStateMapSerDe.java
@@ -32,6 +32,7 @@ import java.util.Optional;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
public class TestStateMapSerDe {
@@ -69,5 +70,64 @@ public class TestStateMapSerDe {
assertTrue(stateVersion.isPresent());
assertEquals(version, stateVersion.get());
assertEquals(stateValues, recoveredStateMap.toMap());
+
+ // Journals written with version 0 hold edits in the same layout as
full records
+ final StateMapUpdate versionZeroEdit = deserializeEdit(serde,
baos.toByteArray(), Map.of(), 0);
+ assertEquals(componentId, versionZeroEdit.getComponentId());
+ assertEquals(stateValues, versionZeroEdit.getStateMap().toMap());
+ assertEquals(Optional.of(version),
versionZeroEdit.getStateMap().getStateVersion());
+ }
+
+ @Test
+ public void testPatchEditRoundTrip() throws IOException {
+ final String componentId = "1234";
+ final StateMapSerDe serde = new StateMapSerDe();
+
+ final Map<String, String> largeState = new HashMap<>();
+ for (int i = 0; i < 1000; i++) {
+ largeState.put("key" + i, "value" + i);
+ }
+ final StateMapUpdate previousRecord = new StateMapUpdate(new
StandardStateMap(largeState, Optional.of("4")), componentId, UpdateType.UPDATE);
+
+ final Map<String, String> patches = new HashMap<>();
+ patches.put("key0", "changed");
+ patches.put("key1", null);
+ final Map<String, String> patchedState = new HashMap<>(largeState);
+ patchedState.put("key0", "changed");
+ patchedState.remove("key1");
+ final StandardStateMap patchedStateMap = new
StandardStateMap(patchedState, Optional.of("5"));
+
+ final byte[] patchEdit = serializeEdit(serde, previousRecord, new
StateMapUpdate(patchedStateMap, componentId, UpdateType.UPDATE, patches));
+ final byte[] wholeStateEdit = serializeEdit(serde, previousRecord, new
StateMapUpdate(patchedStateMap, componentId, UpdateType.UPDATE));
+ assertTrue(patchEdit.length * 100 < wholeStateEdit.length);
+
+ // Both kinds of edit recover to the same state when the previous
record is available
+ final Map<Object, StateMapUpdate> recoveredRecords =
Map.of(componentId, previousRecord);
+ for (final byte[] edit : new byte[][] {patchEdit, wholeStateEdit}) {
+ final StateMapUpdate recoveredEdit = deserializeEdit(serde, edit,
recoveredRecords, serde.getVersion());
+ assertEquals(patchedState, recoveredEdit.getStateMap().toMap());
+ assertEquals(Optional.of("5"),
recoveredEdit.getStateMap().getStateVersion());
+ }
+
+ final IOException noPreviousRecord = assertThrows(IOException.class,
() -> deserializeEdit(serde, patchEdit, Map.of(), serde.getVersion()));
+ assertTrue(noPreviousRecord.getMessage().contains(componentId));
+
+ // Without a previous record to apply patches to, the whole state is
written
+ final byte[] firstEdit = serializeEdit(serde, null, new
StateMapUpdate(patchedStateMap, componentId, UpdateType.UPDATE, patches));
+ assertEquals(patchedState, deserializeEdit(serde, firstEdit, Map.of(),
serde.getVersion()).getStateMap().toMap());
+ }
+
+ private byte[] serializeEdit(final StateMapSerDe serde, final
StateMapUpdate previousRecord, final StateMapUpdate record) throws IOException {
+ final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (final DataOutputStream out = new DataOutputStream(baos)) {
+ serde.serializeEdit(previousRecord, record, out);
+ }
+ return baos.toByteArray();
+ }
+
+ private StateMapUpdate deserializeEdit(final StateMapSerDe serde, final
byte[] edit, final Map<Object, StateMapUpdate> recoveredRecords, final int
version) throws IOException {
+ try (final DataInputStream in = new DataInputStream(new
ByteArrayInputStream(edit))) {
+ return serde.deserializeEdit(in, recoveredRecords, version);
+ }
}
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
index ec53d427d8c..1ea6e5fe918 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
@@ -31,6 +31,7 @@ import java.util.Optional;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -87,6 +88,23 @@ public abstract class AbstractTestStateProvider {
map = provider.getState(componentId);
assertEquals("value2", map.get(key));
assertTrue(map.getStateVersion().isPresent());
+
+ final Optional<String> versionBeforeEmptyPatch = map.getStateVersion();
+ provider.patchState(Collections.emptyMap(), componentId);
+ assertEquals(versionBeforeEmptyPatch,
provider.getState(componentId).getStateVersion());
+
+ final Map<String, String> patch = new HashMap<>();
+ patch.put("other", "");
+ patch.put("missing", null);
+ provider.patchState(patch, componentId);
+ map = provider.getState(componentId);
+ assertEquals(Map.of(key, "value2", "other", ""), map.toMap());
+ assertNotEquals(versionBeforeEmptyPatch, map.getStateVersion());
+
+ final Map<String, String> removal = new HashMap<>();
+ removal.put(key, null);
+ provider.patchState(removal, componentId);
+ assertEquals(Map.of("other", ""),
provider.getState(componentId).toMap());
}
@Test
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/local/TestWriteAheadLocalStateProvider.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/local/TestWriteAheadLocalStateProvider.java
index fb11a2eef0f..1112b8104ae 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/local/TestWriteAheadLocalStateProvider.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/local/TestWriteAheadLocalStateProvider.java
@@ -20,6 +20,7 @@ package org.apache.nifi.controller.state.providers.local;
import org.apache.nifi.attribute.expression.language.StandardPropertyValue;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.components.PropertyValue;
+import org.apache.nifi.components.state.StateMap;
import org.apache.nifi.components.state.StateProvider;
import org.apache.nifi.components.state.StateProviderInitializationContext;
import org.apache.nifi.controller.state.StateMapUpdate;
@@ -137,6 +138,69 @@ public class TestWriteAheadLocalStateProvider extends
AbstractTestStateProvider
}
}
+ /**
+ * Verifies that patching state writes only the patches to the journal, so
that many small patches of a large state do not fill
+ * the journal, and that the patched state and its version are recovered
both from the journal and from a journal applied on top of a snapshot.
+ */
+ @Test
+ public void testPatchStateWritesPatchesAndRecovers() throws IOException {
+ final Path storageDirectory =
temporaryDirectory.resolve("patch-journal");
+ final String testComponentId = "test-patch-component";
+ final String largeValue = "v".repeat(100);
+
+ final Map<String, String> expectedState = new HashMap<>();
+ for (int i = 0; i < 2000; i++) {
+ expectedState.put("key" + i, largeValue);
+ }
+
+ final StateMap patchedStateMap;
+ final WriteAheadLocalStateProvider firstProvider =
initializeProvider(storageDirectory.toString(), "1 MB");
+ try {
+ firstProvider.setState(expectedState, testComponentId);
+ for (int i = 0; i < 20; i++) {
+ final Map<String, String> patch = new HashMap<>();
+ patch.put("patched" + i, String.valueOf(i));
+ patch.put("key" + i, null);
+ firstProvider.patchState(patch, testComponentId);
+
+ expectedState.put("patched" + i, String.valueOf(i));
+ expectedState.remove("key" + i);
+ }
+
+ final List<File> journalFiles = getJournalFiles(storageDirectory);
+ assertEquals(1, journalFiles.size());
+ assertEquals("0.journal", journalFiles.getFirst().getName());
+
+ patchedStateMap = firstProvider.getState(testComponentId);
+ assertEquals(expectedState, patchedStateMap.toMap());
+ } finally {
+ firstProvider.shutdown();
+ }
+
+ final StateMap secondPatchedStateMap;
+ final WriteAheadLocalStateProvider secondProvider =
initializeProvider(storageDirectory.toString(), "1 MB");
+ try {
+ final StateMap recoveredStateMap =
secondProvider.getState(testComponentId);
+ assertEquals(expectedState, recoveredStateMap.toMap());
+ assertEquals(patchedStateMap.getStateVersion(),
recoveredStateMap.getStateVersion());
+
+ secondProvider.patchState(Map.of("afterRecovery", "1"),
testComponentId);
+ expectedState.put("afterRecovery", "1");
+ secondPatchedStateMap = secondProvider.getState(testComponentId);
+ } finally {
+ secondProvider.shutdown();
+ }
+
+ final WriteAheadLocalStateProvider thirdProvider =
initializeProvider(storageDirectory.toString(), "1 MB");
+ try {
+ final StateMap recoveredStateMap =
thirdProvider.getState(testComponentId);
+ assertEquals(expectedState, recoveredStateMap.toMap());
+ assertEquals(secondPatchedStateMap.getStateVersion(),
recoveredStateMap.getStateVersion());
+ } finally {
+ thirdProvider.shutdown();
+ }
+ }
+
private void updateStateRepeatedly(final StateProvider stateProvider,
final String stateComponentId, final int iterations) throws IOException {
for (int i = 0; i < iterations; i++) {
stateProvider.setState(Collections.singletonMap("iteration",
String.valueOf(i)), stateComponentId);
diff --git
a/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java
b/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java
index dcdae4e2280..54a8080c8a5 100644
--- a/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java
+++ b/nifi-mock/src/main/java/org/apache/nifi/util/MockProcessSession.java
@@ -35,6 +35,7 @@ import org.apache.nifi.processor.io.StreamCallback;
import org.apache.nifi.processor.metrics.CommitTiming;
import org.apache.nifi.provenance.ProvenanceReporter;
import org.apache.nifi.state.MockStateManager;
+import org.apache.nifi.state.MockStateMap;
import org.junit.jupiter.api.Assertions;
import java.io.ByteArrayInputStream;
@@ -89,6 +90,7 @@ public class MockProcessSession implements ProcessSession {
// A List of OutputStreams that have been created by calls to {@link
#write(FlowFile)} and have not yet been closed.
private final Map<FlowFile, OutputStream> openOutputStreams = new
HashMap<>();
private final StateManager stateManager;
+ private final Map<Scope, PendingStatePatches> pendingStatePatches = new
HashMap<>();
private final boolean allowSynchronousCommits;
private final boolean allowRecursiveReads;
private final boolean failCommit;
@@ -340,6 +342,15 @@ public class MockProcessSession implements ProcessSession {
closeStreams(openInputStreams, enforceStreamsClosed);
closeStreams(openOutputStreams, enforceStreamsClosed);
+ for (final Map.Entry<Scope, PendingStatePatches> entry :
pendingStatePatches.entrySet()) {
+ try {
+ stateManager.patchState(entry.getValue().values,
entry.getKey());
+ } catch (final IOException e) {
+ throw new ProcessException("Failed to patch " + entry.getKey()
+ " state", e);
+ }
+ }
+ pendingStatePatches.clear();
+
committed = true;
beingProcessed.clear();
currentVersions.clear();
@@ -816,6 +827,8 @@ public class MockProcessSession implements ProcessSession {
@Override
public void rollback(final boolean penalize) {
+ pendingStatePatches.clear();
+
//if we've already committed then rollback is basically a no-op
if (committed) {
return;
@@ -1405,21 +1418,77 @@ public class MockProcessSession implements
ProcessSession {
@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;
+ }
+
+ return applyPatches(storedState, patches);
}
+ /**
+ * 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);
+ final PendingStatePatches patches = pendingStatePatches.get(scope);
+ if (patches == null) {
+ return stateManager.replace(oldValue, newValue, scope);
+ }
+
+ // The same stored state is used for the comparison and the
replacement, so that an update stored in between causes the replacement to fail
+ final StateMap storedState = stateManager.getState(scope);
+ final StateMap currentState = applyPatches(storedState, patches);
+ if (oldValue == null ||
!oldValue.getStateVersion().equals(currentState.getStateVersion()) ||
!oldValue.toMap().equals(currentState.toMap())) {
+ return false;
+ }
+
+ final boolean replaced = stateManager.replace(storedState, newValue,
scope);
+ if (replaced) {
+ pendingStatePatches.remove(scope);
+ }
+ return replaced;
+ }
+
+ private static StateMap applyPatches(final StateMap storedState, final
PendingStatePatches patches) {
+ final long storedVersion =
storedState.getStateVersion().map(Long::parseLong).orElse(-1L);
+ return new MockStateMap(storedState.toMap(patches.values),
storedVersion + patches.count);
+ }
+
+ /**
+ * Holds the given changes until the session is committed, so that rolling
back the session discards them, as the framework does.
+ * A later call to {@link #setState(Map, Scope)}, {@link
#replaceState(StateMap, Map, Scope)}, or {@link #clearState(Scope)} discards
+ * the patches that have not yet been committed.
+ */
+ @Override
+ public void patchState(final Map<String, String> updatedState, final Scope
scope) throws IOException {
+ Objects.requireNonNull(updatedState, "Updated State required");
+ // Retrieving the state verifies that the component is allowed to use
state in the given scope
+ stateManager.getState(scope);
+ if (updatedState.isEmpty()) {
+ return;
+ }
+
+ final PendingStatePatches patches =
pendingStatePatches.computeIfAbsent(scope, key -> new PendingStatePatches());
+ patches.values.putAll(updatedState);
+ patches.count++;
}
@Override
public void clearState(final Scope scope) throws IOException {
stateManager.clear(scope);
+ pendingStatePatches.remove(scope);
}
@Override
@@ -1528,4 +1597,12 @@ public class MockProcessSession implements
ProcessSession {
);
}
}
+
+ /**
+ * The state changes that the session has not yet committed for one scope,
along with the number of times that state was patched.
+ */
+ private static final class PendingStatePatches {
+ private final Map<String, String> values = new HashMap<>();
+ private int count;
+ }
}
diff --git
a/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java
b/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java
index 4eba9e0ca43..6772284ea8c 100644
--- a/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java
+++ b/nifi-mock/src/test/java/org/apache/nifi/util/TestMockProcessSession.java
@@ -37,6 +37,7 @@ import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Collection;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -573,6 +574,7 @@ public class TestMockProcessSession {
assertThrows(AssertionError.class, () -> {
session.setState(Map.of("key", "value"), Scope.LOCAL);
}, "Was able to set local state without declaring the being
stateful");
+ assertThrows(AssertionError.class, () ->
session.patchState(Map.of("key", "value"), Scope.LOCAL));
stateManager.assertStateNotSet();
}
@@ -591,6 +593,46 @@ public class TestMockProcessSession {
sessionOfStatefulProcessor.commitAsync();
assertEquals(expectedState,
sessionOfStatefulProcessor.getState(Scope.LOCAL).toMap());
+
+ final Map<String, String> patch = new HashMap<>();
+ patch.put("key", null);
+ patch.put("other", "value");
+ sessionOfStatefulProcessor.patchState(patch, Scope.LOCAL);
+ assertEquals(Map.of("other", "value"),
sessionOfStatefulProcessor.getState(Scope.LOCAL).toMap());
+ stateManagerOfStatefulProcessor.assertStateEquals(expectedState,
Scope.LOCAL);
+
+ sessionOfStatefulProcessor.rollback();
+ assertEquals(expectedState,
sessionOfStatefulProcessor.getState(Scope.LOCAL).toMap());
+
+ final StateMap stateBeforeSameValuePatch =
sessionOfStatefulProcessor.getState(Scope.LOCAL);
+ sessionOfStatefulProcessor.patchState(expectedState, Scope.LOCAL);
+ assertEquals(expectedState,
sessionOfStatefulProcessor.getState(Scope.LOCAL).toMap());
+
assertFalse(sessionOfStatefulProcessor.replaceState(stateBeforeSameValuePatch,
Map.of("replaced", "value"), Scope.LOCAL));
+
+ sessionOfStatefulProcessor.patchState(patch, Scope.LOCAL);
+
assertTrue(sessionOfStatefulProcessor.replaceState(sessionOfStatefulProcessor.getState(Scope.LOCAL),
Map.of("replaced", "value"), Scope.LOCAL));
+ sessionOfStatefulProcessor.patchState(Map.of("patched", "value"),
Scope.LOCAL);
+ sessionOfStatefulProcessor.commitAsync();
+
stateManagerOfStatefulProcessor.assertStateEquals(Map.of("replaced", "value",
"patched", "value"), Scope.LOCAL);
+
+ // An update stored by another party right after the session reads
the state must not be overwritten by a replacement
+ final AtomicBoolean storeUpdateAfterNextRead = new AtomicBoolean();
+ final MockStateManager concurrentlyUpdatedStateManager = new
MockStateManager(statefulProcessor) {
+ @Override
+ public synchronized StateMap getState(final Scope scope)
throws IOException {
+ final StateMap state = super.getState(scope);
+ if (storeUpdateAfterNextRead.getAndSet(false)) {
+ super.setState(Map.of("concurrent", "update"), scope);
+ }
+ return state;
+ }
+ };
+ final MockProcessSession concurrentlyUpdatedSession = new
MockProcessSession(sharedStateOfStatefulProcessor, statefulProcessor,
concurrentlyUpdatedStateManager);
+ concurrentlyUpdatedSession.patchState(patch, Scope.LOCAL);
+ final StateMap patchedState =
concurrentlyUpdatedSession.getState(Scope.LOCAL);
+ storeUpdateAfterNextRead.set(true);
+ assertFalse(concurrentlyUpdatedSession.replaceState(patchedState,
Map.of("replaced", "value"), Scope.LOCAL));
+
concurrentlyUpdatedStateManager.assertStateEquals(Map.of("concurrent",
"update"), Scope.LOCAL);
}
@Test
diff --git
a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions-services/src/main/java/org/apache/nifi/cs/tests/system/VerifyLocalClusterStateService.java
b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions-services/src/main/java/org/apache/nifi/cs/tests/system/VerifyLocalClusterStateService.java
index 3780bb61aac..d035c73d224 100644
---
a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions-services/src/main/java/org/apache/nifi/cs/tests/system/VerifyLocalClusterStateService.java
+++
b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions-services/src/main/java/org/apache/nifi/cs/tests/system/VerifyLocalClusterStateService.java
@@ -29,7 +29,6 @@ import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
-import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
@@ -71,8 +70,6 @@ public class VerifyLocalClusterStateService extends
AbstractControllerService im
@Override
public void setState(final String key, final String value, final Scope
scope) throws IOException {
- final Map<String, String> state = new
HashMap<>(getStateManager().getState(scope).toMap());
- state.put(key, value);
- getStateManager().setState(state, scope);
+ getStateManager().patchState(Map.of(key, value), scope);
}
}
diff --git
a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/MultiKeyState.java
b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/MultiKeyState.java
index 34ea12014de..525093afb9d 100644
---
a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/MultiKeyState.java
+++
b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/MultiKeyState.java
@@ -64,7 +64,7 @@ public class MultiKeyState extends AbstractProcessor {
newState.put("b", String.valueOf(b));
newState.put("c", String.valueOf(c));
try {
- session.setState(newState, Scope.CLUSTER);
+ session.patchState(newState, Scope.CLUSTER);
} catch (IOException e) {
throw new ProcessException(e);
}
diff --git
a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/state/ClusterStateKeyDropIT.java
b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/state/ClusterStateKeyDropIT.java
index 5cb341dbeac..cf0ec7016e4 100644
---
a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/state/ClusterStateKeyDropIT.java
+++
b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/state/ClusterStateKeyDropIT.java
@@ -128,6 +128,10 @@ public class ClusterStateKeyDropIT extends
AbstractStateKeyDropIT {
assertEquals(newState, updatedState);
waitForProcessorState(processorId, Scope.CLUSTER, newState);
+
+ // MultiKeyState patches all three keys, so the dropped key is created
again alongside the updated keys
+ runProcessorOnce(processor);
+ waitForProcessorState(processorId, Scope.CLUSTER, Map.of("a", "1",
"b", "2", "c", "2"));
}
@Test
