Remove objects from ActiveDataHandler that are no longer available in used DataProvider (#11167)

Remove objects from DataCommunicator.ActiveDataHandler that are no longer available in used DataProvider.

2 files changed, 90 insertions, 1 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
index 5a99cacfb9..74818f9f42 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -140,6 +140,15 @@ public class DataCommunicator<T> extends AbstractExtension {
         }
 
         /**
+         * Marks all currently active data objects to be dropped.
+         *
+         * @since
+         */
+        public void dropAllActiveData() {
+            activeData.forEach(this::dropActiveData);
+        }
+
+        /**
          * Marks a data object identified by given key string to be dropped.
          *
          * @param key
@@ -152,6 +161,19 @@ public class DataCommunicator<T> extends AbstractExtension {
         }
 
         /**
+         * Returns all dropped data mapped by their id from DataProvider.
+         *
+         * @return map of ids to dropped data objects
+         *
+         * @since
+         */
+        protected Map<Object, T> getDroppedData() {
+            Function<T, Object> getId = getDataProvider()::getId;
+            return droppedData.stream().map(getKeyMapper()::get)
+                    .collect(Collectors.toMap(getId, i -> i));
+        }
+
+        /**
          * Returns all currently active data mapped by their id from
          * DataProvider.
          *
@@ -329,6 +351,10 @@ public class DataCommunicator<T> extends AbstractExtension {
         }
 
         if (initial || reset) {
+            if (reset) {
+                handler.dropAllActiveData();
+            }
+
             rpc.reset(getDataProviderSize());
         }
 
diff --git a/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java b/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java
index d07fa27eb0..c187c91471 100644
--- a/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java
@@ -4,18 +4,21 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.concurrent.Future;
 
 import org.junit.Test;
 import org.mockito.Mockito;
 
+import com.vaadin.data.provider.DataCommunicator.ActiveDataHandler;
 import com.vaadin.server.MockVaadinSession;
 import com.vaadin.server.SerializableConsumer;
 import com.vaadin.server.SerializablePredicate;
 import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinService;
 import com.vaadin.server.VaadinSession;
+import com.vaadin.shared.Range;
 import com.vaadin.shared.Registration;
 import com.vaadin.ui.UI;
 
@@ -30,6 +33,7 @@ import elemental.json.JsonObject;
 public class DataCommunicatorTest {
 
     private static final Object TEST_OBJECT = new Object();
+    private static final Object TEST_OBJECT_TWO = new Object();
 
     public static class TestUI extends UI {
 
@@ -61,7 +65,8 @@ public class DataCommunicatorTest {
         private Registration registration;
 
         public TestDataProvider() {
-            super(Collections.singleton(TEST_OBJECT));
+            super(new ArrayList());
+            addItem(TEST_OBJECT);
         }
 
         @Override
@@ -81,6 +86,18 @@ public class DataCommunicatorTest {
             return registration != null;
         }
 
+        public void setItem(Object item) {
+            clear();
+            addItem(item);
+        }
+
+        public void clear() {
+            getItems().clear();
+        }
+
+        public void addItem(Object item) {
+            getItems().add(item);
+        }
     }
 
     private static class TestDataCommunicator extends DataCommunicator<Object> {
@@ -184,6 +201,52 @@ public class DataCommunicatorTest {
     }
 
     @Test
+    public void refreshDataProviderRemovesOldObjectsFromActiveDataHandler() {
+        session.lock();
+
+        UI ui = new TestUI(session);
+
+        TestDataProvider dataProvider = new TestDataProvider();
+
+        TestDataCommunicator communicator = new TestDataCommunicator();
+        communicator.setDataProvider(dataProvider, null);
+
+        communicator.extend(ui);
+        communicator.beforeClientResponse(true);
+
+        DataKeyMapper keyMapper = communicator.getKeyMapper();
+        assertTrue("Object not mapped by key mapper",
+                keyMapper.has(TEST_OBJECT));
+
+        ActiveDataHandler handler = communicator.getActiveDataHandler();
+        assertTrue("Object not amongst active data",
+                handler.getActiveData().containsKey(TEST_OBJECT));
+
+        dataProvider.setItem(TEST_OBJECT_TWO);
+        dataProvider.refreshAll();
+
+        communicator.beforeClientResponse(false);
+
+        // assert that test object is marked as removed
+        assertTrue("Object not marked as dropped",
+                handler.getDroppedData().containsKey(TEST_OBJECT));
+
+        communicator
+                .setPushRows(Range.between(0, communicator.getMinPushSize()));
+        communicator.beforeClientResponse(false);
+
+        assertFalse("Object still mapped by key mapper",
+                keyMapper.has(TEST_OBJECT));
+        assertTrue("Object not mapped by key mapper",
+                keyMapper.has(TEST_OBJECT_TWO));
+
+        assertFalse("Object still amongst active data",
+                handler.getActiveData().containsKey(TEST_OBJECT));
+        assertTrue("Object not amongst active data",
+                handler.getActiveData().containsKey(TEST_OBJECT_TWO));
+    }
+
+    @Test
     public void testDestroyData() {
         session.lock();
         UI ui = new TestUI(session);