Avoid state mutation in a public method, add a warning

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.
+     * <p>
+     * <b>WARNING:</b> 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

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.
+     * <p>
+     *
+     * <b>WARNING</b>: 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<NodeChange> collector) {
         // The updateActiveState method can create new dirty nodes, so they need
         // to be collected as well
         while (evaluateNewDirtyNodes) {
-            Set<StateNode> dirtyNodesSet = collectDirtyNodes();
+            Set<StateNode> 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<StateNode> collectDirtyNodes() {
-        Set<StateNode> 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.
+     * <p>
+     * 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<StateNode> doCollectDirtyNodes(boolean reset) {
+        if (reset) {
+            Set<StateNode> collectedNodes = dirtyNodes;
+            dirtyNodes = new LinkedHashSet<>();
+            return collectedNodes;
+        } else {
+            return Collections.unmodifiableSet(dirtyNodes);
+        }
+
+    }
 }

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<StateNode> emptyCollection = tree.collectDirtyNodes();
-        Assert.assertTrue("Dirty nodes should be empty after collection",
-                emptyCollection.isEmpty());
+        tree.collectChanges(change -> {
+        });
 
         node2.markAsDirty();
 
         Set<StateNode> 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(),