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 ac80b9d900 NIFI-11595 StateProvider.replace() supports creating the
initial state
ac80b9d900 is described below
commit ac80b9d900221fedb59142e4a4802bf4ab2a17f4
Author: Peter Turcsanyi <[email protected]>
AuthorDate: Sat May 27 21:24:13 2023 +0200
NIFI-11595 StateProvider.replace() supports creating the initial state
- Extracted common logic from setState() and replace() into modifyState()
- Removed redundant code from createNode() because exceptions are handled
on the caller side
- NodeExistsException and InterruptedException are handled in setState()
and replace()
- Also used KeeperException's subclasses instead of KeeperException.code()
This closes #7324
Signed-off-by: David Handermann <[email protected]>
---
.../local/WriteAheadLocalStateProvider.java | 5 -
.../zookeeper/ZooKeeperStateProvider.java | 185 +++++++++------------
.../state/providers/AbstractTestStateProvider.java | 55 +++++-
.../provider/KubernetesConfigMapStateProvider.java | 12 +-
.../KubernetesConfigMapStateProviderTest.java | 53 +++++-
.../StandardOidcAuthorizedClientRepository.java | 11 +-
...StandardOidcAuthorizedClientRepositoryTest.java | 4 +-
.../nifi/redis/state/RedisStateProvider.java | 14 +-
.../nifi/redis/state/ITRedisStateProvider.java | 55 +++++-
9 files changed, 252 insertions(+), 142 deletions(-)
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
index 69276abbd4..5461ba4575 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/local/WriteAheadLocalStateProvider.java
@@ -271,11 +271,6 @@ public class WriteAheadLocalStateProvider extends
AbstractStateProvider {
// see above explanation as to why this method is synchronized.
public synchronized boolean replace(final StateMap oldValue, final
Map<String, String> newValue) throws IOException {
- if (!stateMap.getStateVersion().isPresent()) {
- // state has never been set so return false
- return false;
- }
-
if (stateMap != oldValue) {
return false;
}
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
index 9f02b6b88a..b28b31c74d 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
@@ -36,8 +36,11 @@ import org.apache.nifi.processor.util.StandardValidators;
import org.apache.nifi.util.NiFiProperties;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.KeeperException.BadVersionException;
import org.apache.zookeeper.KeeperException.Code;
+import org.apache.zookeeper.KeeperException.NodeExistsException;
import org.apache.zookeeper.KeeperException.NoNodeException;
+import org.apache.zookeeper.KeeperException.SessionExpiredException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZKUtil;
@@ -74,7 +77,6 @@ import java.util.stream.Collectors;
* consistency across configuration interactions.
*/
public class ZooKeeperStateProvider extends AbstractStateProvider {
- private static final int EMPTY_VERSION = -1;
private static final String COMPONENTS_RELATIVE_PATH = "/components";
@@ -301,11 +303,6 @@ public class ZooKeeperStateProvider extends
AbstractStateProvider {
return new Scope[]{Scope.CLUSTER};
}
- @Override
- public void setState(final Map<String, String> state, final String
componentId) throws IOException {
- setState(state, -1, componentId);
- }
-
private byte[] serialize(final Map<String, String> stateValues) throws
IOException {
try (final ByteArrayOutputStream baos = new ByteArrayOutputStream();
@@ -357,108 +354,104 @@ public class ZooKeeperStateProvider extends
AbstractStateProvider {
}
}
- private void setState(final Map<String, String> stateValues, final int
version, final String componentId) throws IOException {
- try {
- setState(stateValues, version, componentId, true);
- } catch (final NoNodeException nne) {
- // should never happen because we are passing 'true' for
allowNodeCreation
- throw new IOException("Unable to create Node in ZooKeeper to set
state for component with ID " + componentId, nne);
- }
+ /**
+ * Sets the component state to the given stateValues unconditionally..
+ *
+ * @param stateValues the new values to set
+ * @param componentId the ID of the component whose state is being updated
+ *
+ * @throws IOException if unable to communicate with ZooKeeper
+ * @throws StateTooLargeException if the state to be stored exceeds the
maximum size allowed by ZooKeeper (Based on jute.maxbuffer property, after
serialization)
+ */
+ @Override
+ public void setState(final Map<String, String> stateValues, final String
componentId) throws IOException {
+ final StateModifier stateModifier = (keeper, path, data) -> {
+ try {
+ keeper.setData(path, data, -1);
+ } catch (final NoNodeException nne) {
+ try {
+ createNode(path, data, acl);
+ } catch (NodeExistsException nee) {
+ setState(stateValues, componentId);
+ }
+ }
+ return true;
+ };
+
+ modifyState(stateValues, componentId, stateModifier);
}
/**
- * Sets the component state to the given stateValues if and only if the
version is equal to the version currently
- * tracked by ZooKeeper (or if the version is -1, in which case the state
will be updated regardless of the version).
+ * Sets the component state to the given stateValues if and only if
oldState's version is equal to the version currently
+ * tracked by ZooKeeper or if oldState's version is not present (no old
state) and the Zookeeper node does not exist.
*
+ * @param oldState the old state whose version is used to compare against
* @param stateValues the new values to set
- * @param version the expected version of the ZNode
* @param componentId the ID of the component whose state is being updated
- * @param allowNodeCreation if <code>true</code> and the corresponding
ZNode does not exist in ZooKeeper, it will be created; if <code>false</code>
- * and the corresponding node does not exist in ZooKeeper, a
{@link KeeperException.NoNodeException} will be thrown
*
* @throws IOException if unable to communicate with ZooKeeper
- * @throws NoNodeException if the corresponding ZNode does not exist in
ZooKeeper and allowNodeCreation is set to <code>false</code>
* @throws StateTooLargeException if the state to be stored exceeds the
maximum size allowed by ZooKeeper (Based on jute.maxbuffer property, after
serialization)
*/
- private void setState(final Map<String, String> stateValues, final int
version, final String componentId, final boolean allowNodeCreation) throws
IOException, NoNodeException {
+ @Override
+ public boolean replace(final StateMap oldState, final Map<String, String>
stateValues, final String componentId) throws IOException {
+ final Optional<Integer> version =
oldState.getStateVersion().map(Integer::parseInt);
+
+ final StateModifier stateModifier = (keeper, path, data) -> {
+ if (version.isPresent()) {
+ try {
+ keeper.setData(path, data, version.get());
+ } catch (final BadVersionException | NoNodeException e) {
+ return false;
+ }
+ } else {
+ try {
+ createNode(path, data, acl);
+ } catch (final NodeExistsException e) {
+ return false;
+ }
+ }
+ return true;
+ };
+
+ return modifyState(stateValues, componentId, stateModifier);
+ }
+
+ private boolean modifyState(final Map<String, String> stateValues, final
String componentId, final StateModifier stateModifier) throws IOException {
verifyEnabled();
try {
+ final ZooKeeper keeper = getZooKeeper();
final String path = getComponentPath(componentId);
final byte[] data = serialize(stateValues);
- final ZooKeeper keeper = getZooKeeper();
validateDataSize(keeper.getClientConfig(), data, componentId,
stateValues.size());
- try {
- keeper.setData(path, data, version);
- } catch (final NoNodeException nne) {
- if (allowNodeCreation) {
- createNode(path, data, componentId, stateValues, acl);
- return;
- } else {
- throw nne;
- }
- }
- } catch (final InterruptedException e) {
- Thread.currentThread().interrupt();
- throw new IOException("Failed to set cluster-wide state in
ZooKeeper for component with ID " + componentId + " due to interruption", e);
- } catch (final NoNodeException nne) {
- throw nne;
- } catch (final KeeperException ke) {
- if (Code.SESSIONEXPIRED == ke.code()) {
- invalidateClient();
- setState(stateValues, version, componentId, allowNodeCreation);
- return;
- }
- if (Code.NODEEXISTS == ke.code()) {
- setState(stateValues, version, componentId, allowNodeCreation);
- return;
- }
- throw new IOException("Failed to set cluster-wide state in
ZooKeeper for component with ID " + componentId, ke);
+ return stateModifier.apply(keeper, path, data);
+ } catch (final InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ throw new IOException("Failed to set cluster-wide state in
ZooKeeper for component with ID " + componentId + " due to interruption", ie);
+ } catch (final SessionExpiredException see) {
+ invalidateClient();
+ return modifyState(stateValues, componentId, stateModifier);
} catch (final StateTooLargeException stle) {
throw stle;
- } catch (final IOException ioe) {
- throw new IOException("Failed to set cluster-wide state in
ZooKeeper for component with ID " + componentId, ioe);
+ } catch (final Exception e) {
+ throw new IOException("Failed to set cluster-wide state in
ZooKeeper for component with ID " + componentId, e);
}
}
-
- private void createNode(final String path, final byte[] data, final String
componentId, final Map<String, String> stateValues, final List<ACL> acls)
throws IOException, KeeperException {
+ private void createNode(final String path, final byte[] data, final
List<ACL> acls) throws IOException, KeeperException, InterruptedException {
try {
final ZooKeeper zooKeeper = getZooKeeper();
zooKeeper.create(path, data, acls, CreateMode.PERSISTENT);
- } catch (final InterruptedException ie) {
- throw new IOException("Failed to update cluster-wide state due to
interruption", ie);
- } catch (final KeeperException ke) {
- final Code exceptionCode = ke.code();
- if (Code.NONODE == exceptionCode) {
- final String parentPath =
StringUtils.substringBeforeLast(path, "/");
- createNode(parentPath, null, componentId, stateValues,
Ids.OPEN_ACL_UNSAFE);
- createNode(path, data, componentId, stateValues, acls);
- return;
- }
- if (Code.SESSIONEXPIRED == exceptionCode) {
- invalidateClient();
- createNode(path, data, componentId, stateValues, acls);
- return;
- }
-
- // Node already exists. Node must have been created by "someone
else". Just set the data.
- if (Code.NODEEXISTS == exceptionCode) {
- try {
- getZooKeeper().setData(path, data, -1);
- return;
- } catch (final KeeperException ke1) {
- // Node no longer exists -- it was removed by someone
else. Go recreate the node.
- if (ke1.code() == Code.NONODE) {
- createNode(path, data, componentId, stateValues, acls);
- return;
- }
- } catch (final InterruptedException ie) {
- throw new IOException("Failed to update cluster-wide state
due to interruption", ie);
- }
+ } catch (final NoNodeException nne) {
+ final String parentPath = StringUtils.substringBeforeLast(path,
"/");
+ createNode(parentPath, null, Ids.OPEN_ACL_UNSAFE);
+ createNode(path, data, acls);
+ } catch (final NodeExistsException nee) {
+ // Fail only if it is the data/leaf node (not a parent node)
+ if (data != null) {
+ throw nee;
}
- throw ke;
}
}
@@ -493,30 +486,6 @@ public class ZooKeeperStateProvider extends
AbstractStateProvider {
}
}
- @Override
- public boolean replace(final StateMap oldValue, final Map<String, String>
newValue, final String componentId) throws IOException {
- verifyEnabled();
-
- final int version =
oldValue.getStateVersion().map(Integer::parseInt).orElse(EMPTY_VERSION);
- try {
- setState(newValue, version, componentId, false);
- return true;
- } catch (final NoNodeException nne) {
- return false;
- } catch (final IOException ioe) {
- final Throwable cause = ioe.getCause();
- if (cause instanceof KeeperException) {
- final KeeperException ke = (KeeperException) cause;
- if (Code.BADVERSION == ke.code()) {
- return false;
- }
- }
-
- throw ioe;
- }
- }
-
-
@Override
public void clear(final String componentId) throws IOException {
verifyEnabled();
@@ -563,4 +532,8 @@ public class ZooKeeperStateProvider extends
AbstractStateProvider {
}
}
+
+ private interface StateModifier {
+ boolean apply(ZooKeeper keeper, String path, byte[] data) throws
InterruptedException, KeeperException, IOException;
+ }
}
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
index d07929cff4..d74b0adabb 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/state/providers/AbstractTestStateProvider.java
@@ -28,6 +28,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
+import java.util.Optional;
import org.apache.nifi.components.state.StateMap;
import org.apache.nifi.components.state.StateProvider;
@@ -160,15 +161,21 @@ public abstract class AbstractTestStateProvider {
@Test
public void testReplaceWithNonExistingValue() throws Exception {
+ final String key = "testReplaceWithNonExistingValue";
+ final String value = "value";
final StateProvider provider = getProvider();
StateMap stateMap = provider.getState(componentId);
assertNotNull(stateMap);
final Map<String, String> newValue = new HashMap<>();
- newValue.put("value", "value");
+ newValue.put(key, value);
final boolean replaced = provider.replace(stateMap, newValue,
componentId);
- assertFalse(replaced);
+ assertTrue(replaced);
+
+ StateMap map = provider.getState(componentId);
+ assertEquals(value, map.get(key));
+ assertTrue(map.getStateVersion().isPresent());
}
@Test
@@ -198,6 +205,50 @@ public abstract class AbstractTestStateProvider {
assertFalse(replaced);
}
+ @Test
+ void testReplaceConcurrentCreate() throws IOException {
+ final StateProvider provider = getProvider();
+
+ final StateMap stateMap1 = provider.getState(componentId);
+ final StateMap stateMap2 = provider.getState(componentId);
+
+ final boolean replaced1 = provider.replace(stateMap1,
Collections.emptyMap(), componentId);
+
+ assertTrue(replaced1);
+
+ final StateMap replacedStateMap = provider.getState(componentId);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals("0", replacedVersion.get());
+
+ final boolean replaced2 = provider.replace(stateMap2,
Collections.emptyMap(), componentId);
+
+ assertFalse(replaced2);
+ }
+
+ @Test
+ void testReplaceConcurrentUpdate() throws IOException {
+ final StateProvider provider = getProvider();
+
+ provider.setState(Collections.singletonMap("key", "0"), componentId);
+
+ final StateMap stateMap1 = provider.getState(componentId);
+ final StateMap stateMap2 = provider.getState(componentId);
+
+ final boolean replaced1 = provider.replace(stateMap1,
Collections.singletonMap("key", "1"), componentId);
+
+ assertTrue(replaced1);
+
+ final StateMap replacedStateMap = provider.getState(componentId);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals("1", replacedVersion.get());
+
+ final boolean replaced2 = provider.replace(stateMap2,
Collections.singletonMap("key", "2"), componentId);
+
+ assertFalse(replaced2);
+ }
+
@Test
public void testOnComponentRemoved() throws IOException,
InterruptedException {
final StateProvider provider = getProvider();
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/main/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProvider.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/main/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProvider.java
index b2b516fe6e..72a96000d9 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/main/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProvider.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/main/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProvider.java
@@ -168,13 +168,19 @@ public class KubernetesConfigMapStateProvider extends
AbstractConfigurableCompon
final Optional<String> stateVersion = currentState.getStateVersion();
if (stateVersion.isPresent()) {
final String resourceVersion = stateVersion.get();
-
configMapBuilder.editOrNewMetadata().withResourceVersion(resourceVersion);
+
configMapBuilder.editOrNewMetadata().withResourceVersion(resourceVersion).endMetadata();
}
final ConfigMap configMap = configMapBuilder.build();
try {
- final ConfigMap configMapReplaced =
kubernetesClient.configMaps().resource(configMap).replace();
- final Optional<String> version = getVersion(configMapReplaced);
+ Resource<ConfigMap> configMapResource =
kubernetesClient.configMaps().resource(configMap);
+ final ConfigMap newConfigMap;
+ if (stateVersion.isPresent()) {
+ newConfigMap = configMapResource.update();
+ } else {
+ newConfigMap = configMapResource.create();
+ }
+ final Optional<String> version = getVersion(newConfigMap);
logger.debug("Replaced State Component ID [{}] Version [{}]",
componentId, version);
return true;
} catch (final KubernetesClientException e) {
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
index f50834ba67..d0adbcd70f 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-kubernetes-bundle/nifi-framework-kubernetes-state-provider/src/test/java/org/apache/nifi/kubernetes/state/provider/KubernetesConfigMapStateProviderTest.java
@@ -187,7 +187,12 @@ class KubernetesConfigMapStateProviderTest {
final StateMap stateMap = new StandardStateMap(Collections.emptyMap(),
Optional.empty());
final boolean replaced = provider.replace(stateMap,
Collections.emptyMap(), COMPONENT_ID);
- assertFalse(replaced);
+ assertTrue(replaced);
+
+ final StateMap replacedStateMap = provider.getState(COMPONENT_ID);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals(FIRST_VERSION, replacedVersion.get());
}
@Test
@@ -242,6 +247,52 @@ class KubernetesConfigMapStateProviderTest {
assertEquals(COMPONENT_ID, componentIds.next());
}
+ @Test
+ void testReplaceConcurrentCreate() throws IOException {
+ setContext();
+ provider.initialize(context);
+
+ final StateMap stateMap1 = provider.getState(COMPONENT_ID);
+ final StateMap stateMap2 = provider.getState(COMPONENT_ID);
+
+ final boolean replaced1 = provider.replace(stateMap1,
Collections.emptyMap(), COMPONENT_ID);
+
+ assertTrue(replaced1);
+
+ final StateMap replacedStateMap = provider.getState(COMPONENT_ID);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals(FIRST_VERSION, replacedVersion.get());
+
+ final boolean replaced2 = provider.replace(stateMap2,
Collections.emptyMap(), COMPONENT_ID);
+
+ assertFalse(replaced2);
+ }
+
+ @Test
+ void testReplaceConcurrentUpdate() throws IOException {
+ setContext();
+ provider.initialize(context);
+
+ provider.setState(Collections.singletonMap("key", "0"), COMPONENT_ID);
+
+ final StateMap stateMap1 = provider.getState(COMPONENT_ID);
+ final StateMap stateMap2 = provider.getState(COMPONENT_ID);
+
+ final boolean replaced1 = provider.replace(stateMap1,
Collections.singletonMap("key", "1"), COMPONENT_ID);
+
+ assertTrue(replaced1);
+
+ final StateMap replacedStateMap = provider.getState(COMPONENT_ID);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals(SECOND_VERSION, replacedVersion.get());
+
+ final boolean replaced2 = provider.replace(stateMap2,
Collections.singletonMap("key", "2"), COMPONENT_ID);
+
+ assertFalse(replaced2);
+ }
+
private void setContext() {
when(context.getIdentifier()).thenReturn(IDENTIFIER);
when(context.getLogger()).thenReturn(logger);
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepository.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepository.java
index 35bea43be2..29c965025c 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepository.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepository.java
@@ -236,17 +236,10 @@ public class StandardOidcAuthorizedClientRepository
implements OAuth2AuthorizedC
private synchronized void updateState(final String principalId, final
Consumer<Map<String, String>> stateConsumer) {
try {
final StateMap stateMap = getStateMap();
- final Map<String, String> currentStateMap = stateMap.toMap();
- final Map<String, String> updated = new
LinkedHashMap<>(currentStateMap);
+ final Map<String, String> updated = new
LinkedHashMap<>(stateMap.toMap());
stateConsumer.accept(updated);
- final boolean completed;
- if (currentStateMap.isEmpty()) {
- stateManager.setState(updated, SCOPE);
- completed = true;
- } else {
- completed = stateManager.replace(stateMap, updated, SCOPE);
- }
+ final boolean completed = stateManager.replace(stateMap, updated,
SCOPE);
if (completed) {
logger.info("Identity [{}] OIDC Authorized Client update
completed", principalId);
diff --git
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepositoryTest.java
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepositoryTest.java
index 6d9257ffbe..7e05b64362 100644
---
a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepositoryTest.java
+++
b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/oidc/client/web/StandardOidcAuthorizedClientRepositoryTest.java
@@ -149,7 +149,7 @@ class StandardOidcAuthorizedClientRepositoryTest {
repository.saveAuthorizedClient(authorizedClient, principal, request,
response);
verify(authorizedClientConverter).getEncoded(isA(OidcAuthorizedClient.class));
- verify(stateManager).setState(stateMapCaptor.capture(), eq(SCOPE));
+ verify(stateManager).replace(eq(stateMap), stateMapCaptor.capture(),
eq(SCOPE));
final Map<String, String> updatedStateMap = stateMapCaptor.getValue();
final String encodedClient = updatedStateMap.get(IDENTITY);
@@ -165,7 +165,7 @@ class StandardOidcAuthorizedClientRepositoryTest {
repository.removeAuthorizedClient(REGISTRATION_ID, principal, request,
response);
- verify(stateManager).setState(stateMapCaptor.capture(), eq(SCOPE));
+ verify(stateManager).replace(eq(stateMap), stateMapCaptor.capture(),
eq(SCOPE));
final Map<String, String> updatedStateMap = stateMapCaptor.getValue();
assertTrue(updatedStateMap.isEmpty());
}
diff --git
a/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/main/java/org/apache/nifi/redis/state/RedisStateProvider.java
b/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/main/java/org/apache/nifi/redis/state/RedisStateProvider.java
index 0f581a5558..b684a77d5d 100644
---
a/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/main/java/org/apache/nifi/redis/state/RedisStateProvider.java
+++
b/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/main/java/org/apache/nifi/redis/state/RedisStateProvider.java
@@ -198,7 +198,7 @@ public class RedisStateProvider extends
AbstractConfigurableComponent implements
boolean updated = false;
while (!updated && attempted < this.maxAttempts) {
- updated = replace(currStateMap, state, componentId, true);
+ updated = replace(currStateMap, state, componentId);
attempted++;
}
@@ -224,10 +224,6 @@ public class RedisStateProvider extends
AbstractConfigurableComponent implements
@Override
public boolean replace(final StateMap oldValue, final Map<String, String>
newValue, final String componentId) throws IOException {
- return replace(oldValue, newValue, componentId, false);
- }
-
- private boolean replace(final StateMap oldValue, final Map<String, String>
newValue, final String componentId, final boolean allowReplaceMissing) throws
IOException {
return withConnection(redisConnection -> {
boolean replaced = false;
@@ -242,12 +238,6 @@ public class RedisStateProvider extends
AbstractConfigurableComponent implements
final RedisStateMap currStateMap = serDe.deserialize(currValue);
final Optional<String> currentVersion = currStateMap == null ?
Optional.empty() : currStateMap.getStateVersion();
- // the replace API expects that you can't call replace on a
non-existing value, so unwatch and return
- if (!allowReplaceMissing && !currentVersion.isPresent()) {
- redisConnection.unwatch();
- return false;
- }
-
// start a transaction
redisConnection.multi();
@@ -286,7 +276,7 @@ public class RedisStateProvider extends
AbstractConfigurableComponent implements
while (!updated && attempted < this.maxAttempts) {
final StateMap currStateMap = getState(componentId);
- updated = replace(currStateMap, Collections.emptyMap(),
componentId, true);
+ updated = replace(currStateMap, Collections.emptyMap(),
componentId);
final String result = updated ? "successful" : "unsuccessful";
logger.debug("Attempt # {} to clear state for component {} was
{}", attempted + 1, componentId, result);
diff --git
a/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/test/java/org/apache/nifi/redis/state/ITRedisStateProvider.java
b/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/test/java/org/apache/nifi/redis/state/ITRedisStateProvider.java
index 5315439f1f..808163c7c6 100644
---
a/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/test/java/org/apache/nifi/redis/state/ITRedisStateProvider.java
+++
b/nifi-nar-bundles/nifi-redis-bundle/nifi-redis-extensions/src/test/java/org/apache/nifi/redis/state/ITRedisStateProvider.java
@@ -41,6 +41,7 @@ import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;
+import java.util.Optional;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -202,15 +203,21 @@ public class ITRedisStateProvider {
@Test
public void testReplaceWithNonExistingValue() throws Exception {
+ final String key = "testReplaceWithNonExistingValue";
+ final String value = "value";
final StateProvider provider = getProvider();
StateMap stateMap = provider.getState(componentId);
assertNotNull(stateMap);
final Map<String, String> newValue = new HashMap<>();
- newValue.put("value", "value");
+ newValue.put(key, value);
final boolean replaced = provider.replace(stateMap, newValue,
componentId);
- assertFalse(replaced);
+ assertTrue(replaced);
+
+ StateMap map = provider.getState(componentId);
+ assertEquals(value, map.get(key));
+ assertTrue(map.getStateVersion().isPresent());
}
@Test
@@ -240,6 +247,50 @@ public class ITRedisStateProvider {
assertFalse(replaced);
}
+ @Test
+ void testReplaceConcurrentCreate() throws IOException {
+ final StateProvider provider = getProvider();
+
+ final StateMap stateMap1 = provider.getState(componentId);
+ final StateMap stateMap2 = provider.getState(componentId);
+
+ final boolean replaced1 = provider.replace(stateMap1,
Collections.emptyMap(), componentId);
+
+ assertTrue(replaced1);
+
+ final StateMap replacedStateMap = provider.getState(componentId);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals("0", replacedVersion.get());
+
+ final boolean replaced2 = provider.replace(stateMap2,
Collections.emptyMap(), componentId);
+
+ assertFalse(replaced2);
+ }
+
+ @Test
+ void testReplaceConcurrentUpdate() throws IOException {
+ final StateProvider provider = getProvider();
+
+ provider.setState(Collections.singletonMap("key", "0"), componentId);
+
+ final StateMap stateMap1 = provider.getState(componentId);
+ final StateMap stateMap2 = provider.getState(componentId);
+
+ final boolean replaced1 = provider.replace(stateMap1,
Collections.singletonMap("key", "1"), componentId);
+
+ assertTrue(replaced1);
+
+ final StateMap replacedStateMap = provider.getState(componentId);
+ final Optional<String> replacedVersion =
replacedStateMap.getStateVersion();
+ assertTrue(replacedVersion.isPresent());
+ assertEquals("1", replacedVersion.get());
+
+ final boolean replaced2 = provider.replace(stateMap2,
Collections.singletonMap("key", "2"), componentId);
+
+ assertFalse(replaced2);
+ }
+
@Test
public void testOnComponentRemoved() throws IOException,
InterruptedException {
final StateProvider provider = getProvider();