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.git
The following commit(s) were added to refs/heads/main by this push:
new 9148f1ba28d NIFI-15917 Fixed potential corruption with
modify-after-write on Local State Provider (#11220)
9148f1ba28d is described below
commit 9148f1ba28d642ce6c070fefbd791c38d9edb4e0
Author: Mark Payne <[email protected]>
AuthorDate: Fri May 8 09:06:06 2026 -0400
NIFI-15917 Fixed potential corruption with modify-after-write on Local
State Provider (#11220)
When creating a StandardStateMap, create a defensive copy of the provided
map to avoid allowing modifications by the extension after state is stored to
alter the internal state.
Signed-off-by: David Handermann <[email protected]>
---
.../nifi/controller/state/StandardStateMap.java | 8 ++-
.../local/WriteAheadLocalStateProvider.java | 7 ++
.../local/TestWriteAheadLocalStateProvider.java | 83 +++++++++++++++++-----
3 files changed, 80 insertions(+), 18 deletions(-)
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateMap.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateMap.java
index 09fa7e234fc..fa615ea09af 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateMap.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/state/StandardStateMap.java
@@ -20,6 +20,7 @@ package org.apache.nifi.controller.state;
import org.apache.nifi.components.state.StateMap;
import java.util.Collections;
+import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
@@ -30,7 +31,12 @@ public class StandardStateMap implements StateMap {
private final Optional<String> stateVersion;
public StandardStateMap(final Map<String, String> stateValues, final
Optional<String> stateVersion) {
- this.stateValues = Collections.unmodifiableMap(stateValues == null ?
Collections.emptyMap() : stateValues);
+ // Defensively copy the caller's Map so that this StateMap's contents
are immutable and
+ // isolated from any subsequent mutations to the caller's reference.
Extensions that retain
+ // and continue to mutate the Map they pass to setState() across
invocations would otherwise
+ // race with the framework's checkpoint thread iterating the StateMap
during snapshot
+ // serialization.
+ this.stateValues = (stateValues == null || stateValues.isEmpty()) ?
Collections.emptyMap() : Collections.unmodifiableMap(new
HashMap<>(stateValues));
this.stateVersion = stateVersion;
}
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 82dea780d9b..83135a8e4b0 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
@@ -168,6 +168,13 @@ public class WriteAheadLocalStateProvider extends
AbstractStateProvider {
return properties;
}
+ // Visible for testing. Forces a synchronous checkpoint of the underlying
Write-Ahead Log so
+ // that tests can deterministically observe the on-disk snapshot without
waiting for the
+ // scheduled CheckpointTask to run.
+ void checkpoint() throws IOException {
+ writeAheadLog.checkpoint();
+ }
+
@Override
public synchronized void shutdown() {
executor.shutdown();
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 1c4cd1db385..af3c32e7933 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
@@ -28,30 +28,91 @@ import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.parameter.ParameterLookup;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
import org.wali.WriteAheadRepository;
import java.io.IOException;
+import java.nio.file.Path;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
-import java.util.UUID;
import javax.net.ssl.SSLContext;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
public class TestWriteAheadLocalStateProvider extends
AbstractTestStateProvider {
+ @TempDir
+ private Path temporaryDirectory;
+
private StateProvider provider;
private WriteAheadRepository<StateMapUpdate> wal;
@BeforeEach
public void setup() throws IOException {
- provider = new WriteAheadLocalStateProvider();
+ provider =
initializeProvider(temporaryDirectory.resolve("setup").toString());
+ }
+
+ @AfterEach
+ public void cleanup() throws IOException {
+ provider.onComponentRemoved(componentId);
+ provider.shutdown();
+
+ if (wal != null) {
+ wal.shutdown();
+ }
+ }
+
+ @Override
+ protected StateProvider getProvider() {
+ return provider;
+ }
+
+ /**
+ * Verifies that recovered state reflects the values that the caller
passed to
+ * {@link StateProvider#setState(Map, String)} at the moment of that call,
and is
+ * not influenced by subsequent mutations the caller performs to the same
Map
+ * instance after returning. This guards against a corruption scenario
where a
+ * caller (such as an extension that holds onto its own ConcurrentHashMap)
keeps
+ * mutating its state map between calls to setState while the framework's
+ * checkpoint thread is concurrently serializing a snapshot. The provider
must
+ * isolate its persisted state from any further mutation of the caller's
map.
+ */
+ @Test
+ public void testRecoveredStateIsolatedFromPostSetStateMapMutations()
throws Exception {
+ final String storageDirectory =
temporaryDirectory.resolve("recovery").toString();
+ final String testComponentId = "test-mutation-isolation-component";
+
+ final WriteAheadLocalStateProvider firstProvider =
initializeProvider(storageDirectory);
+ final Map<String, String> liveStateValues = new HashMap<>();
+ liveStateValues.put("key1", "value1");
+ firstProvider.setState(liveStateValues, testComponentId);
+
+ liveStateValues.put("key2", "value2");
+ liveStateValues.put("key3", "value3");
+
+ firstProvider.checkpoint();
+ firstProvider.shutdown();
+
+ final WriteAheadLocalStateProvider recoveredProvider =
initializeProvider(storageDirectory);
+ try {
+ final Map<String, String> recoveredValues =
recoveredProvider.getState(testComponentId).toMap();
+ assertEquals(Collections.singletonMap("key1", "value1"),
recoveredValues);
+ } finally {
+ recoveredProvider.shutdown();
+ }
+ }
+
+ private WriteAheadLocalStateProvider initializeProvider(final String
storageDirectory) throws IOException {
+ final WriteAheadLocalStateProvider newProvider = new
WriteAheadLocalStateProvider();
final Map<PropertyDescriptor, PropertyValue> properties = new
HashMap<>();
- properties.put(WriteAheadLocalStateProvider.PATH, new
StandardPropertyValue("target/local-state-provider/" +
UUID.randomUUID().toString(), null, ParameterLookup.EMPTY));
+ properties.put(WriteAheadLocalStateProvider.PATH, new
StandardPropertyValue(storageDirectory, null, ParameterLookup.EMPTY));
properties.put(WriteAheadLocalStateProvider.ALWAYS_SYNC, new
StandardPropertyValue("false", null, ParameterLookup.EMPTY));
properties.put(WriteAheadLocalStateProvider.CHECKPOINT_INTERVAL, new
StandardPropertyValue("2 mins", null, ParameterLookup.EMPTY));
properties.put(WriteAheadLocalStateProvider.NUM_PARTITIONS, new
StandardPropertyValue("16", null, ParameterLookup.EMPTY));
- provider.initialize(new StateProviderInitializationContext() {
+ newProvider.initialize(new StateProviderInitializationContext() {
@Override
public String getIdentifier() {
return "Unit Test Provider Initialization Context";
@@ -90,19 +151,7 @@ public class TestWriteAheadLocalStateProvider extends
AbstractTestStateProvider
return null;
}
});
- }
- @AfterEach
- public void cleanup() throws IOException {
- provider.onComponentRemoved(componentId);
-
- if (wal != null) {
- wal.shutdown();
- }
- }
-
- @Override
- protected StateProvider getProvider() {
- return provider;
+ return newProvider;
}
}
