fix: Let DataCommunicator#flush keep track of which StateTree is supp…

…osed to execute its flushRequest (#14068) (#14193)

* DataCommunicator#flush can be stuck when moving between UI instances, for instance when the @PreserveOnRefresh annotation is present.

Co-authored-by: Marco Collovati <mcollovati@gmail.com>
Co-authored-by: Willem Verstraeten <willem.verstraeten@gmail.com>

(cherry picked from commit c55bdcc)

flow-data/src/main/java/com/vaadin/flow/data/provider/DataCommunicator.java

@@ -41,6 +41,8 @@
 import com.vaadin.flow.function.SerializableSupplier;
 import com.vaadin.flow.internal.ExecutionContext;
 import com.vaadin.flow.internal.JsonUtils;
+import com.vaadin.flow.internal.NodeOwner;
+import com.vaadin.flow.internal.NullOwner;
 import com.vaadin.flow.internal.Range;
 import com.vaadin.flow.internal.StateNode;
 import com.vaadin.flow.shared.Registration;
@@ -109,9 +111,8 @@ public class DataCommunicator<T> implements Serializable {
 
     private Registration dataProviderUpdateRegistration;
     private HashSet<T> updatedData = new HashSet<>();
-
-    private SerializableConsumer<ExecutionContext> flushRequest;
-    private SerializableConsumer<ExecutionContext> flushUpdatedDataRequest;
+    private FlushRequest flushRequest;
+    private FlushRequest flushUpdatedDataRequest;
 
     private CallbackDataProvider.CountCallback<T, ?> countCallback;
     private int itemCountEstimate = -1;
@@ -1060,28 +1061,27 @@ private void requestFlush() {
     }
 
     private void requestFlush(boolean forced) {
-        if ((flushRequest == null || forced) && fetchEnabled) {
-            flushRequest = context -> {
+        if ((flushRequest == null || !flushRequest.canExecute(stateNode)
+                || forced) && fetchEnabled) {
+            flushRequest = FlushRequest.register(stateNode, context -> {
                 if (!context.isClientSideInitialized()) {
                     reset();
                     arrayUpdater.initialize();
                 }
                 flush();
                 flushRequest = null;
-            };
-            stateNode.runWhenAttached(ui -> ui.getInternals().getStateTree()
-                    .beforeClientResponse(stateNode, flushRequest));
+            });
         }
     }
 
     private void requestFlushUpdatedData() {
-        if (flushUpdatedDataRequest == null) {
-            flushUpdatedDataRequest = context -> {
-                flushUpdatedData();
-                flushUpdatedDataRequest = null;
-            };
-            stateNode.runWhenAttached(ui -> ui.getInternals().getStateTree()
-                    .beforeClientResponse(stateNode, flushUpdatedDataRequest));
+        if (flushUpdatedDataRequest == null
+                || !flushUpdatedDataRequest.canExecute(stateNode)) {
+            flushUpdatedDataRequest = FlushRequest.register(stateNode,
+                    context -> {
+                        flushUpdatedData();
+                        flushUpdatedDataRequest = null;
+                    });
         }
     }
 
@@ -1402,6 +1402,27 @@ public static Activation empty() {
         }
     }
 
+    private static class FlushRequest implements Serializable {
+
+        private NodeOwner owner;
+
+        static FlushRequest register(StateNode stateNode,
+                SerializableConsumer<ExecutionContext> action) {
+            FlushRequest request = new FlushRequest();
+            request.owner = stateNode.getOwner();
+            stateNode.runWhenAttached(ui -> {
+                request.owner = stateNode.getOwner();
+                ui.getInternals().getStateTree().beforeClientResponse(stateNode,
+                        action);
+            });
+            return request;
+        }
+
+        boolean canExecute(StateNode stateNode) {
+            return owner instanceof NullOwner || owner == stateNode.getOwner();
+        }
+    }
+
     private static Logger getLogger() {
         return LoggerFactory.getLogger(DataCommunicator.class);
     }

flow-data/src/test/java/com/vaadin/flow/data/provider/DataCommunicatorTest.java

@@ -1567,6 +1567,26 @@ public void setRequestedRange_customPageSize_customPageSizeConsidered_itemsReque
                 600, queryCaptor.getValue().getLimit());
     }
 
+    // Simulates a flush request enqueued during a page reload with
+    // @PreserveOnRefresh
+    // see https://github.com/vaadin/flow/issues/14067
+    @Test
+    public void reattach_differentUI_requestFlushExecuted() {
+        dataCommunicator.setDataProvider(createDataProvider(), null);
+        dataCommunicator.setRequestedRange(0, 50);
+
+        MockUI newUI = new MockUI();
+        // simulates preserve on refresh
+        // DataCommunicator has a flushRequest pending
+        // that should be rescheduled on the new state tree
+        newUI.getInternals().moveElementsFrom(ui);
+        ui = newUI;
+        fakeClientCommunication();
+
+        Assert.assertEquals("Expected initial full reset.",
+                Range.withLength(0, 50), lastSet);
+    }
+
     @Tag("test-component")
     private static class TestComponent extends Component {