diff --git a/flow-server/src/main/java/com/vaadin/flow/internal/StateNode.java b/flow-server/src/main/java/com/vaadin/flow/internal/StateNode.java
index 84209ea622c..f63896bf389 100644
--- a/flow-server/src/main/java/com/vaadin/flow/internal/StateNode.java
+++ b/flow-server/src/main/java/com/vaadin/flow/internal/StateNode.java
@@ -559,6 +559,12 @@ boolean isClientSideInitialized() {
* {@link #collectChanges(Consumer)} has been called. If the node is
* recently attached, then the reported changes will be relative to a newly
* created node.
+ *
+ * WARNING: this is in fact an internal (private method) which is
+ * expected to be called from {@link StateTree#collectChanges(Consumer)}
+ * method only (which is effectively private itself). Don't call this method
+ * from any other place because it will break the expected {@link UI} state.
+ *
*
* @param collector
* a consumer accepting node changes
diff --git a/flow-server/src/main/java/com/vaadin/flow/internal/StateTree.java b/flow-server/src/main/java/com/vaadin/flow/internal/StateTree.java
index 5eddf66c85d..1fb220ea062 100644
--- a/flow-server/src/main/java/com/vaadin/flow/internal/StateTree.java
+++ b/flow-server/src/main/java/com/vaadin/flow/internal/StateTree.java
@@ -34,6 +34,7 @@
import com.vaadin.flow.internal.change.NodeChange;
import com.vaadin.flow.internal.nodefeature.NodeFeature;
import com.vaadin.flow.server.VaadinSession;
+import com.vaadin.flow.server.communication.UidlWriter;
import com.vaadin.flow.shared.Registration;
/**
@@ -245,6 +246,12 @@ public StateNode getNodeById(int id) {
/**
* Collects all changes made to this tree since the last time
* {@link #collectChanges(Consumer)} has been called.
+ *
+ *
+ * WARNING: This is an internal method which is not intended to be
+ * used outside. The only proper caller of this method is {@link UidlWriter}
+ * class (the {@code UidlWriter::encodeChanges} method). Any call of this
+ * method in any other place will break the expected {@link UI} state.
*
* @param collector
* a consumer accepting node changes
@@ -256,7 +263,7 @@ public void collectChanges(Consumer collector) {
// The updateActiveState method can create new dirty nodes, so they need
// to be collected as well
while (evaluateNewDirtyNodes) {
- Set dirtyNodesSet = collectDirtyNodes();
+ Set dirtyNodesSet = doCollectDirtyNodes(true);
dirtyNodesSet.forEach(StateNode::updateActiveState);
evaluateNewDirtyNodes = allDirtyNodes.addAll(dirtyNodesSet);
}
@@ -275,15 +282,12 @@ public void markAsDirty(StateNode node) {
}
/**
- * Gets all the nodes that have been marked as dirty since the last time
- * this method was invoked.
+ * Gets all the nodes that have been marked.
*
* @return a set of dirty nodes, in the order they were marked dirty
*/
public Set collectDirtyNodes() {
- Set collectedNodes = dirtyNodes;
- dirtyNodes = new LinkedHashSet<>();
- return collectedNodes;
+ return doCollectDirtyNodes(false);
}
/**
@@ -404,4 +408,22 @@ private void checkHasLock() {
session.checkHasLock();
}
}
+
+ /**
+ * Gets all the nodes that have been marked as dirty.
+ *
+ * If {@code reset} is {@code true} then dirty nodes collection is reset.
+ *
+ * @return a set of dirty nodes, in the order they were marked dirty
+ */
+ private Set doCollectDirtyNodes(boolean reset) {
+ if (reset) {
+ Set collectedNodes = dirtyNodes;
+ dirtyNodes = new LinkedHashSet<>();
+ return collectedNodes;
+ } else {
+ return Collections.unmodifiableSet(dirtyNodes);
+ }
+
+ }
}
diff --git a/flow-server/src/test/java/com/vaadin/flow/internal/StateTreeTest.java b/flow-server/src/test/java/com/vaadin/flow/internal/StateTreeTest.java
index a2db7c1e0da..4d31da38ac4 100644
--- a/flow-server/src/test/java/com/vaadin/flow/internal/StateTreeTest.java
+++ b/flow-server/src/test/java/com/vaadin/flow/internal/StateTreeTest.java
@@ -194,15 +194,14 @@ public void testDirtyNodeCollection() {
Assert.assertEquals("Both nodes should initially be empty",
new HashSet<>(Arrays.asList(node1, node2)), initialDirty);
- Set emptyCollection = tree.collectDirtyNodes();
- Assert.assertTrue("Dirty nodes should be empty after collection",
- emptyCollection.isEmpty());
+ tree.collectChanges(change -> {
+ });
node2.markAsDirty();
Set collectAfterOneMarked = tree.collectDirtyNodes();
- Assert.assertEquals("Marked node should be in collect result",
- Collections.singleton(node2), collectAfterOneMarked);
+ Assert.assertTrue("Marked node should be in collect result",
+ collectAfterOneMarked.contains(node2));
}
@Test
@@ -223,6 +222,9 @@ public void testDirtyNodeCollectionOrder() {
Assert.assertArrayEquals(expected.toArray(),
tree.collectDirtyNodes().toArray());
+ tree.collectChanges(change -> {
+ });
+
nodes.forEach(StateNode::markAsDirty);
expected = new ArrayList<>(nodes);
Assert.assertArrayEquals(expected.toArray(),