This is an automated email from the ASF dual-hosted git repository.
exceptionfactory pushed a commit to branch support/nifi-1.x
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/support/nifi-1.x by this push:
new 878d2e025f NIFI-11595 StateProvider.replace() supports creating the
initial state
878d2e025f is described below
commit 878d2e025f2464dad5ff95e2a095ea22036a0f2e
Author: Peter Turcsanyi <[email protected]>
AuthorDate: Sat May 27 14:24:13 2023 -0500
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]>
(cherry picked from commit ac80b9d900221fedb59142e4a4802bf4ab2a17f4)
---
.../local/WriteAheadLocalStateProvider.java | 5 -
.../zookeeper/ZooKeeperStateProvider.java | 188 +++++++++------------
.../state/providers/AbstractTestStateProvider.java | 55 +++++-
.../StandardOidcAuthorizedClientRepository.java | 11 +-
...StandardOidcAuthorizedClientRepositoryTest.java | 4 +-
.../nifi/redis/state/RedisStateProvider.java | 14 +-
.../nifi/redis/state/ITRedisStateProvider.java | 55 +++++-
7 files changed, 196 insertions(+), 136 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 e341c4cb47..9e7a365bda 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
@@ -255,11 +255,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.getVersion() == -1L) {
- // 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 e4c3be49a8..747aa56ae6 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;
@@ -72,6 +75,11 @@ import java.util.stream.Collectors;
* consistency across configuration interactions.
*/
public class ZooKeeperStateProvider extends AbstractStateProvider {
+
+ private static final String COMPONENTS_RELATIVE_PATH = "/components";
+
+ private static final String COMPONENTS_PATH_FORMAT = "%s%s/%s";
+
private static final Logger logger =
LoggerFactory.getLogger(ZooKeeperStateProvider.class);
private NiFiProperties nifiProperties;
@@ -293,11 +301,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();
@@ -348,108 +351,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;
}
}
@@ -484,29 +483,6 @@ public class ZooKeeperStateProvider extends
AbstractStateProvider {
}
}
- @Override
- public boolean replace(final StateMap oldValue, final Map<String, String>
newValue, final String componentId) throws IOException {
- verifyEnabled();
-
- try {
- setState(newValue, (int) oldValue.getVersion(), 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();
@@ -529,4 +505,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 ccff20597c..2010f8e8cf 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
@@ -26,6 +26,7 @@ import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
+import java.util.Optional;
import org.apache.nifi.components.state.StateMap;
import org.apache.nifi.components.state.StateProvider;
@@ -150,15 +151,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
@@ -188,6 +195,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-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 07a36a8557..ba71dc488b 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
@@ -187,7 +187,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++;
}
@@ -213,10 +213,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;
@@ -231,12 +227,6 @@ public class RedisStateProvider extends
AbstractConfigurableComponent implements
final RedisStateMap currStateMap = serDe.deserialize(currValue);
final long currVersion = currStateMap == null ? -1L :
currStateMap.getVersion();
- // the replace API expects that you can't call replace on a
non-existing value, so unwatch and return
- if (!allowReplaceMissing && currVersion == -1) {
- redisConnection.unwatch();
- return false;
- }
-
// start a transaction
redisConnection.multi();
@@ -274,7 +264,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
{}", new Object[] { 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 8b28cd4df6..be660e89ed 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
@@ -39,6 +39,7 @@ import java.util.Collections;
import java.util.HashMap;
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;
@@ -187,15 +188,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
@@ -225,6 +232,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();
