Correctly use id to identify data when refreshing (#9398)

This patch refactors the internals of Grid single selection model
implementation.

Fixes #9380

9 files changed, 243 insertions, 65 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 2ab8c2d00e..0730e67904 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -23,8 +23,11 @@ import java.util.Comparator;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -151,20 +154,22 @@ public class DataCommunicator<T> extends AbstractExtension {
         }
 
         /**
-         * Returns the collection of all currently active data.
+         * Returns all currently active data mapped by their id from
+         * DataProvider.
          *
-         * @return collection of active data objects
+         * @return map of ids to active data objects
          */
-        public Collection<T> getActiveData() {
-            HashSet<T> hashSet = new HashSet<>();
-            for (String key : activeData) {
-                hashSet.add(getKeyMapper().get(key));
-            }
-            return hashSet;
+        public Map<Object, T> getActiveData() {
+            Function<T, Object> getId = getDataProvider()::getId;
+            return activeData.stream().map(getKeyMapper()::get)
+                    .collect(Collectors.toMap(getId, i -> i));
         }
 
         @Override
         public void generateData(T data, JsonObject jsonObject) {
+            // Make sure KeyMapper is up to date
+            getKeyMapper().refresh(data, dataProvider::getId);
+
             // Write the key string for given data object
             jsonObject.put(DataCommunicatorConstants.KEY,
                     getKeyMapper().key(data));
@@ -490,19 +495,24 @@ public class DataCommunicator<T> extends AbstractExtension {
      * Informs the DataProvider that a data object has been updated.
      *
      * @param data
-     *            updated data object
+     *            updated data object; not {@code null}
      */
     public void refresh(T data) {
-        if (!handler.getActiveData().contains(data)) {
-            // Item is not currently available at the client-side
-            return;
-        }
+        Objects.requireNonNull(data,
+                "DataCommunicator can not refresh null object");
+        Object id = getDataProvider().getId(data);
 
-        if (updatedData.isEmpty()) {
-            markAsDirty();
-        }
+        // ActiveDataHandler has always the latest data through KeyMapper.
+        Map<Object, T> activeData = getActiveDataHandler().getActiveData();
 
-        updatedData.add(data);
+        if (activeData.containsKey(id)) {
+            // Item is currently available at the client-side
+            if (updatedData.isEmpty()) {
+                markAsDirty();
+            }
+
+            updatedData.add(activeData.get(id));
+        }
     }
 
     /**
@@ -699,8 +709,8 @@ public class DataCommunicator<T> extends AbstractExtension {
                     getUI().access(() -> {
                         if (event instanceof DataRefreshEvent) {
                             T item = ((DataRefreshEvent<T>) event).getItem();
-                            generators.forEach(g -> g.refreshData(item));
                             keyMapper.refresh(item, dataProvider::getId);
+                            generators.forEach(g -> g.refreshData(item));
                             refresh(item);
                         } else {
                             reset();
diff --git a/server/src/main/java/com/vaadin/data/provider/DataProvider.java b/server/src/main/java/com/vaadin/data/provider/DataProvider.java
index e609eaaf03..f24be8ca17 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataProvider.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataProvider.java
@@ -120,6 +120,9 @@ public interface DataProvider<T, F> extends Serializable {
      * Default is to use item itself as its own identifier. If the item has
      * {@link Object#equals(Object)} and {@link Object#hashCode()} implemented
      * in a way that it can be compared to other items, no changes are required.
+     * <p>
+     * <strong>Note:</strong> This method will be called often by the Framework.
+     * It should not do any expensive operations.
      *
      * @param item
      *            the item to get identifier for; not {@code null}
@@ -277,8 +280,8 @@ public interface DataProvider<T, F> extends Serializable {
      * <p>
      * <strong>Using big streams is not recommended, you should instead use a
      * lazy data provider.</strong> See
-     * {@link #fromCallbacks(FetchCallback, CountCallback)}
-     * or {@link BackEndDataProvider} for more info.
+     * {@link #fromCallbacks(FetchCallback, CountCallback)} or
+     * {@link BackEndDataProvider} for more info.
      *
      * @param <T>
      *            the data item type
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
index ca8ccc89dd..715538112e 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -530,7 +530,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         Objects.requireNonNull(provider, "Provider can't be null");
         itemCollapseAllowedProvider = provider;
 
-        getActiveDataHandler().getActiveData().forEach(this::refresh);
+        getActiveDataHandler().getActiveData().values().forEach(this::refresh);
     }
 
     /**
diff --git a/server/src/main/java/com/vaadin/ui/components/grid/AbstractSelectionModel.java b/server/src/main/java/com/vaadin/ui/components/grid/AbstractSelectionModel.java
index 7a47da0ce1..7d1c0a1926 100644
--- a/server/src/main/java/com/vaadin/ui/components/grid/AbstractSelectionModel.java
+++ b/server/src/main/java/com/vaadin/ui/components/grid/AbstractSelectionModel.java
@@ -39,6 +39,9 @@ public abstract class AbstractSelectionModel<T> extends AbstractGridExtension<T>
     @Override
     public void generateData(T item, JsonObject jsonObject) {
         if (isSelected(item)) {
+            // Pre-emptive update in case used a stale element in selection.
+            refreshData(item);
+
             jsonObject.put(DataCommunicatorConstants.SELECTED, true);
         }
     }
diff --git a/server/src/main/java/com/vaadin/ui/components/grid/SingleSelectionModelImpl.java b/server/src/main/java/com/vaadin/ui/components/grid/SingleSelectionModelImpl.java
index aae0da8e14..008165734c 100644
--- a/server/src/main/java/com/vaadin/ui/components/grid/SingleSelectionModelImpl.java
+++ b/server/src/main/java/com/vaadin/ui/components/grid/SingleSelectionModelImpl.java
@@ -107,28 +107,17 @@ public class SingleSelectionModelImpl<T> extends AbstractSelectionModel<T>
      *         item (or {@code null} if no selection), {@code false} otherwise.
      */
     protected boolean isKeySelected(String key) {
-        return Objects.equals(key, getSelectedKey());
+        return isSelected(getData(key));
     }
 
     /**
-     * Returns the communication key of the selected item or {@code null} if no
-     * item is selected.
+     * Sets the selected item. If the item is {@code null}, clears the current
+     * selection if any.
      *
-     * @return the key of the selected item if any, {@code null} otherwise.
-     */
-    protected String getSelectedKey() {
-        return itemToKey(selectedItem);
-    }
-
-    /**
-     * Sets the selected item based on the given communication key. If the key
-     * is {@code null}, clears the current selection if any.
-     *
-     * @param key
-     *            the key of the selected item or {@code null} to clear
-     *            selection
+     * @param item
+     *            the selected item or {@code null} to clear selection
      */
-    protected void doSetSelectedKey(String key) {
+    protected void doSetSelected(T item) {
         if (getParent() == null) {
             throw new IllegalStateException(
                     "Trying to update selection for grid selection model that has been detached from the grid.");
@@ -137,7 +126,7 @@ public class SingleSelectionModelImpl<T> extends AbstractSelectionModel<T>
         if (selectedItem != null) {
             getGrid().getDataCommunicator().refresh(selectedItem);
         }
-        selectedItem = getData(key);
+        selectedItem = item;
         if (selectedItem != null) {
             getGrid().getDataCommunicator().refresh(selectedItem);
         }
@@ -158,12 +147,13 @@ public class SingleSelectionModelImpl<T> extends AbstractSelectionModel<T>
             throw new IllegalStateException("Client tried to update selection"
                     + " although user selection is disallowed");
         }
-        if (isKeySelected(key)) {
+        T item = getData(key);
+        if (isSelected(item)) {
             return;
         }
 
-        T oldSelection = this.getSelectedItem().orElse(null);
-        doSetSelectedKey(key);
+        T oldSelection = selectedItem;
+        doSetSelected(item);
         fireEvent(new SingleSelectionEvent<>(getGrid(), asSingleSelect(),
                 oldSelection, true));
     }
@@ -177,36 +167,18 @@ public class SingleSelectionModelImpl<T> extends AbstractSelectionModel<T>
      *            the item to select or {@code null} to clear selection
      */
     protected void setSelectedFromServer(T item) {
-        // TODO creates a key if item not in data provider
-        String key = itemToKey(item);
-
-        if (isSelected(item) || isKeySelected(key)) {
+        if (isSelected(item)) {
+            // Avoid generating an extra key when item matches a stale one.
             return;
         }
 
         T oldSelection = this.getSelectedItem()
                 .orElse(asSingleSelect().getEmptyValue());
-        doSetSelectedKey(key);
+        doSetSelected(item);
         fireEvent(new SingleSelectionEvent<>(getGrid(), asSingleSelect(),
                 oldSelection, false));
     }
 
-    /**
-     * Returns the communication key assigned to the given item.
-     *
-     * @param item
-     *            the item whose key to return
-     * @return the assigned key
-     */
-    protected String itemToKey(T item) {
-        if (item == null) {
-            return null;
-        } else {
-            // TODO creates a key if item not in data provider
-            return getGrid().getDataCommunicator().getKeyMapper().key(item);
-        }
-    }
-
     @Override
     public Set<T> getSelectedItems() {
         if (selectedItem != null) {
@@ -289,6 +261,12 @@ public class SingleSelectionModelImpl<T> extends AbstractSelectionModel<T>
 
     @Override
     public boolean isSelected(T item) {
+        // Quick comparison of objects directly
+        if (Objects.equals(item, selectedItem)) {
+            return true;
+        }
+
+        // Id based check
         return item != null && selectedItem != null
                 && getGrid().getDataProvider().getId(selectedItem)
                         .equals(getGrid().getDataProvider().getId(item));
diff --git a/server/src/test/java/com/vaadin/data/GridAsSingleSelectInBinderTest.java b/server/src/test/java/com/vaadin/data/GridAsSingleSelectInBinderTest.java
index 077ca185a2..e38b86f6db 100644
--- a/server/src/test/java/com/vaadin/data/GridAsSingleSelectInBinderTest.java
+++ b/server/src/test/java/com/vaadin/data/GridAsSingleSelectInBinderTest.java
@@ -34,7 +34,8 @@ public class GridAsSingleSelectInBinderTest
             extends SingleSelectionModelImpl<Sex> {
 
         public void setSelectedFromClient(Sex item) {
-            setSelectedFromClient(itemToKey(item));
+            setSelectedFromClient(
+                    getGrid().getDataCommunicator().getKeyMapper().key(item));
         }
     }
 
diff --git a/server/src/test/java/com/vaadin/data/provider/ReplaceListDataProviderTest.java b/server/src/test/java/com/vaadin/data/provider/ReplaceListDataProviderTest.java
index 7552691413..5c7448d17e 100644
--- a/server/src/test/java/com/vaadin/data/provider/ReplaceListDataProviderTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/ReplaceListDataProviderTest.java
@@ -45,5 +45,4 @@ public class ReplaceListDataProviderTest {
         Assert.assertTrue("Old test object should be stale",
                 dataProvider.isStale(TEST_OBJECT));
     }
-
 }
diff --git a/uitest/src/main/java/com/vaadin/tests/data/DataProviderRefresh.java b/uitest/src/main/java/com/vaadin/tests/data/DataProviderRefresh.java
new file mode 100644
index 0000000000..4476073b6d
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/data/DataProviderRefresh.java
@@ -0,0 +1,118 @@
+package com.vaadin.tests.data;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.Random;
+import java.util.stream.Stream;
+
+import com.vaadin.data.provider.AbstractDataProvider;
+import com.vaadin.data.provider.Query;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.Grid;
+
+public class DataProviderRefresh extends AbstractTestUI {
+
+    public static class Bean implements Serializable {
+
+        private String value;
+        private final int id;
+
+        public Bean(String value, int id) {
+            this.value = value;
+            this.id = id;
+        }
+
+        public String getValue() {
+            return value;
+        }
+
+        public int getId() {
+            return id;
+        }
+
+        public void setValue(String value) {
+            this.value = value;
+        }
+
+        @Override
+        public String toString() {
+            return "{ " + value + ", " + id + " }";
+        }
+    }
+
+    /**
+     * A dummy data provider for testing item replacement and stale elements.
+     */
+    public class ReplaceListDataProvider
+            extends AbstractDataProvider<Bean, Void> {
+
+        private final List<Bean> backend;
+
+        public ReplaceListDataProvider(List<Bean> items) {
+            backend = items;
+        }
+
+        @Override
+        public void refreshItem(Bean item) {
+            if (replaceItem(item)) {
+                super.refreshItem(item);
+            }
+        }
+
+        private boolean replaceItem(Bean item) {
+            for (int i = 0; i < backend.size(); ++i) {
+                if (getId(backend.get(i)).equals(getId(item))) {
+                    if (backend.get(i).equals(item)) {
+                        return false;
+                    }
+                    backend.set(i, item);
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        @Override
+        public boolean isInMemory() {
+            return true;
+        }
+
+        @Override
+        public int size(Query<Bean, Void> t) {
+            return backend.size();
+        }
+
+        @Override
+        public Stream<Bean> fetch(Query<Bean, Void> query) {
+            return backend.stream().skip(query.getOffset())
+                    .limit(query.getLimit());
+        }
+
+        @Override
+        public Object getId(Bean item) {
+            return item.getId();
+        }
+    }
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        Grid<Bean> grid = new Grid<>();
+        ArrayList<Bean> arrayList = new ArrayList<>();
+        Bean foo = new Bean("Foo", 10);
+        arrayList.add(foo);
+        arrayList.add(new Bean("Baz", 11));
+        ReplaceListDataProvider dataProvider = new ReplaceListDataProvider(
+                arrayList);
+        grid.setDataProvider(dataProvider);
+        grid.addColumn(Object::toString).setCaption("toString");
+        addComponent(grid);
+        addComponent(new Button("Replace item",
+                e -> dataProvider.refreshItem(new Bean("Bar", 10))));
+        addComponent(new Button("Select old", e -> grid.select(foo)));
+    }
+
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/data/DataProviderRefreshTest.java b/uitest/src/test/java/com/vaadin/tests/data/DataProviderRefreshTest.java
new file mode 100644
index 0000000000..f32ec543a8
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/data/DataProviderRefreshTest.java
@@ -0,0 +1,66 @@
+package com.vaadin.tests.data;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.GridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class DataProviderRefreshTest extends SingleBrowserTest {
+
+    @Test
+    public void select_and_replace() {
+        openTestURL();
+        GridElement grid = $(GridElement.class).first();
+        Assert.assertFalse("Row should not be initially selected",
+                grid.getRow(0).isSelected());
+        // Select item before replace
+        $(ButtonElement.class).caption("Select old").first().click();
+        Assert.assertTrue("Row should be selected",
+                grid.getRow(0).isSelected());
+
+        $(ButtonElement.class).caption("Replace item").first().click();
+        Assert.assertTrue("Row should still be selected after item replace",
+                grid.getRow(0).isSelected());
+        Assert.assertEquals("Grid content was not updated.", "{ Bar, 10 }",
+                grid.getCell(0, 0).getText());
+
+        // Deselect row
+        grid.getCell(0, 0).click();
+        Assert.assertFalse("Row should be deselected after click",
+                grid.getRow(0).isSelected());
+
+        Assert.assertEquals("Second row was affected", "{ Baz, 11 }",
+                grid.getCell(1, 0).getText());
+    }
+
+    @Test
+    public void replace_and_select() {
+        openTestURL();
+        GridElement grid = $(GridElement.class).first();
+        Assert.assertFalse("Row should not be initially selected",
+                grid.getRow(0).isSelected());
+
+        // Replace item before select
+        $(ButtonElement.class).caption("Replace item").first().click();
+        Assert.assertFalse("Row should not be selected after item replace",
+                grid.getRow(0).isSelected());
+        Assert.assertEquals("Grid content was not updated.", "{ Bar, 10 }",
+                grid.getCell(0, 0).getText());
+
+        $(ButtonElement.class).caption("Select old").first().click();
+        Assert.assertTrue("Row should be selected",
+                grid.getRow(0).isSelected());
+        Assert.assertEquals("Grid content should not update.", "{ Bar, 10 }",
+                grid.getCell(0, 0).getText());
+
+        // Deselect row
+        grid.getCell(0, 0).click();
+        Assert.assertFalse("Row should be deselected after click",
+                grid.getRow(0).isSelected());
+
+        Assert.assertEquals("Second row was affected", "{ Baz, 11 }",
+                grid.getCell(1, 0).getText());
+    }
+}