This is an automated email from the ASF dual-hosted git repository.
exceptionfactory pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi-api.git
The following commit(s) were added to refs/heads/main by this push:
new 4800c6b NIFI-16493 Added patchState to StateManager and
ProcessSession (#112)
4800c6b is described below
commit 4800c6bf51da63b22fa4b1bf85e709131f74dc54
Author: Mark Payne <[email protected]>
AuthorDate: Mon Oct 26 12:43:33 2026 -0400
NIFI-16493 Added patchState to StateManager and ProcessSession (#112)
Co-authored-by: Cursor <[email protected]>
Signed-off-by: David Handermann <[email protected]>
---
.../apache/nifi/components/state/StateManager.java | 29 +++++
.../org/apache/nifi/components/state/StateMap.java | 31 +++++
.../nifi/documentation/init/NopStateManager.java | 4 +
.../org/apache/nifi/processor/ProcessSession.java | 31 +++++
.../nifi/components/state/TestStateManager.java | 145 +++++++++++++++++++++
5 files changed, 240 insertions(+)
diff --git a/src/main/java/org/apache/nifi/components/state/StateManager.java
b/src/main/java/org/apache/nifi/components/state/StateManager.java
index 5d9a64f..ad8aaab 100644
--- a/src/main/java/org/apache/nifi/components/state/StateManager.java
+++ b/src/main/java/org/apache/nifi/components/state/StateManager.java
@@ -22,6 +22,7 @@ import
org.apache.nifi.components.state.exception.StateTooLargeException;
import java.io.IOException;
import java.util.Map;
+import java.util.Objects;
/**
* <p>
@@ -91,6 +92,34 @@ public interface StateManager {
*/
boolean replace(StateMap oldValue, Map<String, String> newValue, Scope
scope) throws IOException;
+ /**
+ * Updates only the given keys of the component's state, leaving all other
keys as they are.
+ * <p>
+ * Each key in the given map with a non-null value is created or
overwritten. Each key with a {@code null} value is removed;
+ * removing a key that is not present has no effect on that key. An empty
string is stored as a value and does not remove the key.
+ * Keys that are not in the given map are left unchanged. If no state has
been stored yet, the result is the same as calling
+ * {@link #setState(Map, Scope)} with the given map after removing its
{@code null} values. An empty map does nothing: no state is
+ * written and the version does not change.
+ * </p>
+ *
+ * @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 scope the scope to use when storing the state
+ *
+ * @throws StateTooLargeException if the state that results from applying
the changes is more than the backing storage mechanism allows
+ * @throws IOException if unable to communicate with the underlying
storage mechanism
+ */
+ default void patchState(final Map<String, String> updatedState, final
Scope scope) throws IOException {
+ Objects.requireNonNull(updatedState, "Updated State required");
+ if (updatedState.isEmpty()) {
+ return;
+ }
+
+ StateMap currentState;
+ do {
+ currentState = getState(scope);
+ } while (!replace(currentState, currentState.toMap(updatedState),
scope));
+ }
+
/**
* Clears all keys and values from the component's state
*
diff --git a/src/main/java/org/apache/nifi/components/state/StateMap.java
b/src/main/java/org/apache/nifi/components/state/StateMap.java
index 95829a7..ecc6ab5 100644
--- a/src/main/java/org/apache/nifi/components/state/StateMap.java
+++ b/src/main/java/org/apache/nifi/components/state/StateMap.java
@@ -17,7 +17,10 @@
package org.apache.nifi.components.state;
+import java.util.Collections;
+import java.util.LinkedHashMap;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
/**
@@ -48,4 +51,32 @@ public interface StateMap {
* @return an immutable Map representation of all keys and values for the
state of a component.
*/
Map<String, String> toMap();
+
+ /**
+ * Returns an immutable Map of all keys and values for the state of a
component, with the given patches applied. This StateMap is not changed.
+ * <p>
+ * Each key in the given patches with a non-null value is added or
overwritten. Each key with a {@code null} value is removed.
+ * Keys that are not in the given patches keep their values from this
StateMap.
+ * </p>
+ *
+ * @param patches a map of state keys to their new values. A key that is
mapped to {@code null} is removed from the result. The map itself must not be
{@code null}.
+ * @return an immutable Map of this state's keys and values with the given
patches applied
+ */
+ default Map<String, String> toMap(final Map<String, String> patches) {
+ Objects.requireNonNull(patches, "Patches required");
+ if (patches.isEmpty()) {
+ return toMap();
+ }
+
+ final Map<String, String> patchedValues = new LinkedHashMap<>(toMap());
+ for (final Map.Entry<String, String> entry : patches.entrySet()) {
+ if (entry.getValue() == null) {
+ patchedValues.remove(entry.getKey());
+ } else {
+ patchedValues.put(entry.getKey(), entry.getValue());
+ }
+ }
+
+ return Collections.unmodifiableMap(patchedValues);
+ }
}
diff --git
a/src/main/java/org/apache/nifi/documentation/init/NopStateManager.java
b/src/main/java/org/apache/nifi/documentation/init/NopStateManager.java
index 5e2c955..9f704ca 100644
--- a/src/main/java/org/apache/nifi/documentation/init/NopStateManager.java
+++ b/src/main/java/org/apache/nifi/documentation/init/NopStateManager.java
@@ -37,6 +37,10 @@ public class NopStateManager implements StateManager {
return false;
}
+ @Override
+ public void patchState(final Map<String, String> updatedState, final Scope
scope) {
+ }
+
@Override
public void clear(final Scope scope) {
}
diff --git a/src/main/java/org/apache/nifi/processor/ProcessSession.java
b/src/main/java/org/apache/nifi/processor/ProcessSession.java
index b239058..daaa687 100644
--- a/src/main/java/org/apache/nifi/processor/ProcessSession.java
+++ b/src/main/java/org/apache/nifi/processor/ProcessSession.java
@@ -930,6 +930,37 @@ public interface ProcessSession {
*/
boolean replaceState(StateMap oldValue, Map<String, String> newValue,
Scope scope) throws IOException;
+ /**
+ * Updates only the given keys of the component's state, leaving all other
keys as they are.
+ * <p>
+ * Each key in the given map with a non-null value is created or
overwritten. Each key with a {@code null} value is removed;
+ * removing a key that is not present has no effect on that key. An empty
string is stored as a value and does not remove the key.
+ * Keys that are not in the given map are left unchanged. If no state has
been stored yet, the result is the same as calling
+ * {@link #setState(Map, Scope)} with the given map after removing its
{@code null} values. An empty map does nothing: no state is
+ * written and the version does not change.
+ * <p>
+ * This method does not update the State Provider immediately but rather
caches the changes until the session is committed.
+ * Unless one of the methods described below was called first, the given
keys are updated in the State Provider when the session is
+ * committed, even if the stored state has been changed since this session
retrieved it. Keys that are not in the given map keep whatever
+ * values are stored at the time of the commit. Until then, {@link
#getState(Scope)} on this session returns the state with the changes
+ * applied, and a later call to this method applies its changes on top of
the earlier ones. Rolling back the session discards the changes.
+ * <p>
+ * If {@link #setState(Map, Scope)}, {@link #clearState(Scope)}, or a
successful {@link #replaceState(StateMap, Map, Scope)} was already
+ * called on this session for the same scope and that update has not yet
been written to the State Provider, the changes become part of
+ * that pending update. The commit then treats the resulting state exactly
as it would treat a call to {@link #setState(Map, Scope)}:
+ * the state is stored as a whole, so keys that are not in the given map
do not keep values stored by other writers, and if the commit
+ * does not store that state, as described for {@link #setState(Map,
Scope)}, neither the earlier update nor these changes are stored. Likewise,
calling one of those methods
+ * after this one replaces the changes that have not yet been committed.
+ *
+ * @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 scope the {@link Scope} to use for fetching the current and
storing the new state
+ * @throws IOException if unable to communicate with the underlying
storage mechanism
+ * @throws UnsupportedOperationException if this session implementation
does not support patching state
+ */
+ default void patchState(final Map<String, String> updatedState, final
Scope scope) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
/**
* Clears all keys and values from the component's state.
* <p>
diff --git
a/src/test/java/org/apache/nifi/components/state/TestStateManager.java
b/src/test/java/org/apache/nifi/components/state/TestStateManager.java
new file mode 100644
index 0000000..d0a0087
--- /dev/null
+++ b/src/test/java/org/apache/nifi/components/state/TestStateManager.java
@@ -0,0 +1,145 @@
+/*
+ * 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.components.state;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class TestStateManager {
+
+ private static final String FIRST_KEY = "first";
+ private static final String SECOND_KEY = "second";
+ private static final String THIRD_KEY = "third";
+ private static final String MISSING_KEY = "missing";
+ private static final String OTHER_WRITER_KEY = "otherWriter";
+
+ private static final String FIRST_VALUE = "1";
+ private static final String SECOND_VALUE = "2";
+ private static final String EMPTY_VALUE = "";
+
+ @Test
+ void testPatchState() throws Exception {
+ final InMemoryStateManager stateManager = new InMemoryStateManager();
+
+ assertThrows(NullPointerException.class, () ->
stateManager.patchState(null, Scope.LOCAL));
+ stateManager.patchState(Map.of(), Scope.LOCAL);
+ assertEquals(0, stateManager.replaceCount);
+
assertTrue(stateManager.getState(Scope.LOCAL).getStateVersion().isEmpty());
+
+ final Map<String, String> firstPatch = new HashMap<>();
+ firstPatch.put(FIRST_KEY, FIRST_VALUE);
+ firstPatch.put(SECOND_KEY, null);
+ firstPatch.put(THIRD_KEY, EMPTY_VALUE);
+ stateManager.patchState(firstPatch, Scope.LOCAL);
+ final StateMap firstState = stateManager.getState(Scope.LOCAL);
+ assertEquals(Map.of(FIRST_KEY, FIRST_VALUE, THIRD_KEY, EMPTY_VALUE),
firstState.toMap());
+ assertTrue(firstState.getStateVersion().isPresent());
+
+ final Map<String, String> secondPatch = new HashMap<>();
+ secondPatch.put(FIRST_KEY, SECOND_VALUE);
+ secondPatch.put(MISSING_KEY, null);
+ stateManager.patchState(secondPatch, Scope.LOCAL);
+ assertEquals(Map.of(FIRST_KEY, SECOND_VALUE, THIRD_KEY, EMPTY_VALUE),
stateManager.getState(Scope.LOCAL).toMap());
+
+ final Map<String, String> thirdPatch = new HashMap<>();
+ thirdPatch.put(THIRD_KEY, null);
+ stateManager.patchState(thirdPatch, Scope.LOCAL);
+ assertEquals(Map.of(FIRST_KEY, SECOND_VALUE),
stateManager.getState(Scope.LOCAL).toMap());
+ assertEquals(3, stateManager.replaceCount);
+ }
+
+ @Test
+ void testPatchStateRetriesAfterFailedReplace() throws Exception {
+ final InMemoryStateManager stateManager = new InMemoryStateManager();
+ stateManager.setState(Map.of(FIRST_KEY, FIRST_VALUE), Scope.LOCAL);
+ stateManager.writeBeforeNextReplace = Map.of(OTHER_WRITER_KEY,
FIRST_VALUE);
+
+ stateManager.patchState(Map.of(SECOND_KEY, SECOND_VALUE), Scope.LOCAL);
+
+ assertEquals(Map.of(FIRST_KEY, FIRST_VALUE, SECOND_KEY, SECOND_VALUE,
OTHER_WRITER_KEY, FIRST_VALUE), stateManager.getState(Scope.LOCAL).toMap());
+ assertEquals(2, stateManager.replaceCount);
+ }
+
+ private record SimpleStateMap(Map<String, String> values, Optional<String>
version) implements StateMap {
+ @Override
+ public Optional<String> getStateVersion() {
+ return version;
+ }
+
+ @Override
+ public String get(final String key) {
+ return values.get(key);
+ }
+
+ @Override
+ public Map<String, String> toMap() {
+ return values;
+ }
+ }
+
+ /**
+ * Holds a single state map for all scopes. When {@code
writeBeforeNextReplace} is set, the next call to
+ * {@link #replace(StateMap, Map, Scope)} first stores those keys, as
another writer would, and then fails because
+ * the state no longer matches the caller's old value.
+ */
+ private static class InMemoryStateManager implements StateManager {
+ private StateMap currentState = new SimpleStateMap(Map.of(),
Optional.empty());
+ private long version = 0;
+ private int replaceCount = 0;
+ private Map<String, String> writeBeforeNextReplace;
+
+ @Override
+ public void setState(final Map<String, String> state, final Scope
scope) {
+ currentState = new SimpleStateMap(Map.copyOf(state),
Optional.of(String.valueOf(++version)));
+ }
+
+ @Override
+ public StateMap getState(final Scope scope) {
+ return currentState;
+ }
+
+ @Override
+ public boolean replace(final StateMap oldValue, final Map<String,
String> newValue, final Scope scope) {
+ replaceCount++;
+ if (writeBeforeNextReplace != null) {
+ final Map<String, String> otherWriterState = new
HashMap<>(currentState.toMap());
+ otherWriterState.putAll(writeBeforeNextReplace);
+ writeBeforeNextReplace = null;
+ setState(otherWriterState, scope);
+ }
+
+ if (oldValue != currentState) {
+ return false;
+ }
+
+ setState(newValue, scope);
+ return true;
+ }
+
+ @Override
+ public void clear(final Scope scope) {
+ setState(Map.of(), scope);
+ }
+ }
+}
