Reset HierarchicalDataCommunicator on changes (#9275)

Reset HDC when encountering unexpected changes in the data.

Additionally this patch fixes an issue with client and server caches
getting out of sync during resets.

8 files changed, 327 insertions, 55 deletions

diff --git a/client/src/main/java/com/vaadin/client/data/AbstractRemoteDataSource.java b/client/src/main/java/com/vaadin/client/data/AbstractRemoteDataSource.java
index 6ec549a16c..c03d3b2936 100644
--- a/client/src/main/java/com/vaadin/client/data/AbstractRemoteDataSource.java
+++ b/client/src/main/java/com/vaadin/client/data/AbstractRemoteDataSource.java
@@ -904,7 +904,8 @@ public abstract class AbstractRemoteDataSource<T> implements DataSource<T> {
      */
     protected void resetDataAndSize(int newSize) {
         size = newSize;
-        dropFromCache(getCachedRange());
+        indexToRowMap.clear();
+        keyToIndexMap.clear();
         cached = Range.withLength(0, 0);
 
         getHandlers().forEach(dch -> dch.resetDataAndSize(newSize));
diff --git a/server/src/main/java/com/vaadin/data/HierarchyData.java b/server/src/main/java/com/vaadin/data/HierarchyData.java
index df5fbbbbf5..0cf680db2f 100644
--- a/server/src/main/java/com/vaadin/data/HierarchyData.java
+++ b/server/src/main/java/com/vaadin/data/HierarchyData.java
@@ -27,7 +27,7 @@ import java.util.stream.Stream;
 
 /**
  * Class for representing hierarchical data.
- * 
+ *
  * @author Vaadin Ltd
  * @since 8.1
  *
@@ -95,13 +95,13 @@ public class HierarchyData<T> implements Serializable {
      * Adds a data item as a child of {@code parent}. Call with {@code null} as
      * parent to add a root level item. The given parent item must already exist
      * in this structure, and an item can only be added to this structure once.
-     * 
+     *
      * @param parent
      *            the parent item for which the items are added as children
      * @param item
      *            the item to add
      * @return this
-     * 
+     *
      * @throws IllegalArgumentException
      *             if parent is not null and not already added to this structure
      * @throws IllegalArgumentException
@@ -129,13 +129,13 @@ public class HierarchyData<T> implements Serializable {
      * {@code null} as parent to add root level items. The given parent item
      * must already exist in this structure, and an item can only be added to
      * this structure once.
-     * 
+     *
      * @param parent
      *            the parent item for which the items are added as children
      * @param items
      *            the list of items to add
      * @return this
-     * 
+     *
      * @throws IllegalArgumentException
      *             if parent is not null and not already added to this structure
      * @throws IllegalArgumentException
@@ -155,13 +155,13 @@ public class HierarchyData<T> implements Serializable {
      * {@code null} as parent to add root level items. The given parent item
      * must already exist in this structure, and an item can only be added to
      * this structure once.
-     * 
+     *
      * @param parent
      *            the parent item for which the items are added as children
      * @param items
      *            the collection of items to add
      * @return this
-     * 
+     *
      * @throws IllegalArgumentException
      *             if parent is not null and not already added to this structure
      * @throws IllegalArgumentException
@@ -180,13 +180,13 @@ public class HierarchyData<T> implements Serializable {
      * with {@code null} as parent to add root level items. The given parent
      * item must already exist in this structure, and an item can only be added
      * to this structure once.
-     * 
+     *
      * @param parent
      *            the parent item for which the items are added as children
      * @param items
      *            stream of items to add
      * @return this
-     * 
+     *
      * @throws IllegalArgumentException
      *             if parent is not null and not already added to this structure
      * @throws IllegalArgumentException
@@ -203,11 +203,11 @@ public class HierarchyData<T> implements Serializable {
     /**
      * Remove a given item from this structure. Additionally, this will
      * recursively remove any descendants of the item.
-     * 
+     *
      * @param item
      *            the item to remove, or null to clear all data
      * @return this
-     * 
+     *
      * @throws IllegalArgumentException
      *             if the item does not exist in this structure
      */
@@ -219,25 +219,32 @@ public class HierarchyData<T> implements Serializable {
         new ArrayList<>(getChildren(item)).forEach(child -> removeItem(child));
         itemToWrapperMap.get(itemToWrapperMap.get(item).getParent())
                 .removeChild(item);
+        if (item != null) {
+            // remove non root item from backing map
+            itemToWrapperMap.remove(item);
+        }
         return this;
     }
 
     /**
      * Clear all items from this structure. Shorthand for calling
      * {@link #removeItem(Object)} with null.
+     *
+     * @return this
      */
-    public void clear() {
+    public HierarchyData<T> clear() {
         removeItem(null);
+        return this;
     }
 
     /**
      * Get the immediate child items for the given item.
-     * 
+     *
      * @param item
      *            the item for which to retrieve child items for, null to
      *            retrieve all root items
      * @return a list of child items for the given item
-     * 
+     *
      * @throws IllegalArgumentException
      *             if the item does not exist in this structure
      */
@@ -249,7 +256,15 @@ public class HierarchyData<T> implements Serializable {
         return itemToWrapperMap.get(item).getChildren();
     }
 
-    private boolean contains(T item) {
+    /**
+     * Check whether the given item is in this hierarchy.
+     *
+     * @param item
+     *            the item to check
+     * @return {@code true} if the item is in this hierarchy, {@code false} if
+     *         not
+     */
+    public boolean contains(T item) {
         return itemToWrapperMap.containsKey(item);
     }
 
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 820e3d8f2c..5d3bd81a05 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -22,7 +22,6 @@ import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -153,13 +152,27 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     private void loadRequestedRows() {
         final Range requestedRows = getPushRows();
         if (!requestedRows.isEmpty()) {
-            doPushRows(requestedRows);
+            doPushRows(requestedRows, 0);
         }
 
         setPushRows(Range.withLength(0, 0));
     }
 
-    private void doPushRows(final Range requestedRows) {
+    /**
+     * Attempts to push the requested range of rows to the client. Will trigger
+     * a reset if the data provider is unable to provide the requested number of
+     * items.
+     *
+     * @param requestedRows
+     *            the range of rows to push
+     * @param insertRowsCount
+     *            number of rows to insert, beginning at the start index of
+     *            {@code requestedRows}, 0 to not insert any rows
+     * @return {@code true} if the range was successfully pushed to the client,
+     *         {@code false} if the push was unsuccessful and a reset was
+     *         triggered
+     */
+    private boolean doPushRows(final Range requestedRows, int insertRowsCount) {
         Stream<TreeLevelQuery> levelQueries = mapper.splitRangeToLevelQueries(
                 requestedRows.getStart(), requestedRows.getEnd() - 1);
 
@@ -173,7 +186,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             List<T> results = doFetchQuery(query.startIndex, query.size,
                     getKeyMapper().get(query.node.getParentKey()))
                             .collect(Collectors.toList());
-            // TODO if the size differers from expected, all goes to hell
+
             fetchedItems.addAll(results);
             List<JsonObject> rowData = results.stream()
                     .map(item -> createDataObject(item, query.depth))
@@ -181,33 +194,31 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             mapper.reorderLevelQueryResultsToFlatOrdering(rowDataMapper, query,
                     rowData);
         });
-        verifyNoNullItems(dataObjects, requestedRows);
+
+        if (hasNullItems(dataObjects, requestedRows)) {
+            reset();
+            return false;
+        }
+
+        if (insertRowsCount > 0) {
+            getClientRpc().insertRows(requestedRows.getStart(),
+                    insertRowsCount);
+        }
 
         sendData(requestedRows.getStart(), Arrays.asList(dataObjects));
         getActiveDataHandler().addActiveData(fetchedItems.stream());
         getActiveDataHandler().cleanUp(fetchedItems.stream());
+        return true;
     }
 
-    /*
-     * Verify that there are no null objects in the array, to fail eagerly and
-     * not just on the client side.
-     */
-    private void verifyNoNullItems(JsonObject[] dataObjects,
+    private boolean hasNullItems(JsonObject[] dataObjects,
             Range requestedRange) {
-        List<Integer> nullItems = new ArrayList<>(0);
-        AtomicInteger indexCounter = new AtomicInteger();
-        Stream.of(dataObjects).forEach(object -> {
-            int index = indexCounter.getAndIncrement();
+        for (JsonObject object : dataObjects) {
             if (object == null) {
-                nullItems.add(index);
+                return true;
             }
-        });
-        if (!nullItems.isEmpty()) {
-            throw new IllegalStateException("For requested rows "
-                    + requestedRange + ", there was null items for indexes "
-                    + nullItems.stream().map(Object::toString)
-                            .collect(Collectors.joining(", ")));
         }
+        return false;
     }
 
     private JsonObject createDataObject(T item, int depth) {
@@ -288,9 +299,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             String itemKey = keys.getString(i);
             if (!mapper.isKeyStored(itemKey)
                     && !rowKeysPendingExpand.contains(itemKey)) {
-                // FIXME: cache invalidated incorrectly, active keys being
-                // dropped prematurely
-                // getActiveDataHandler().dropActiveData(itemKey);
+                getActiveDataHandler().dropActiveData(itemKey);
             }
         }
     }
@@ -379,6 +388,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
                 collapsedRowIndex);
         getClientRpc().removeRows(collapsedRowIndex + 1,
                 collapsedSubTreeSize);
+
         // FIXME seems like a slight overkill to do this just for refreshing
         // expanded status
         refresh(collapsedItem);
@@ -414,9 +424,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
         int expandedNodeSize = doSizeQuery(expandedItem);
         if (expandedNodeSize == 0) {
-            // TODO handle 0 size -> not expandable
-            throw new IllegalStateException("Row with index " + expandedRowIndex
-                    + " returned no child nodes.");
+            reset();
+            return false;
         }
 
         if (!mapper.isCollapsed(expandedRowKey)) {
@@ -426,17 +435,16 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
                 expandedNodeSize);
         rowKeysPendingExpand.remove(expandedRowKey);
 
-        getClientRpc().insertRows(expandedRowIndex + 1, expandedNodeSize);
-        // TODO optimize by sending "just enough" of the expanded items
-        // directly
-        doPushRows(
-                Range.withLength(expandedRowIndex + 1, expandedNodeSize));
+        boolean success = doPushRows(Range.withLength(expandedRowIndex + 1,
+                Math.min(expandedNodeSize, latestCacheSize)), expandedNodeSize);
 
-        // expanded node needs to be updated to be marked as expanded
-        // FIXME seems like a slight overkill to do this just for refreshing
-        // expanded status
-        refresh(expandedItem);
-        return true;
+        if (success) {
+            // FIXME seems like a slight overkill to do this just for refreshing
+            // expanded status
+            refresh(expandedItem);
+            return true;
+        }
+        return false;
     }
 
     /**
diff --git a/server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java b/server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java
index b86c3186c6..8f70dd8e8c 100644
--- a/server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java
+++ b/server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java
@@ -53,7 +53,7 @@ public class InMemoryHierarchicalDataProvider<T> extends
      * <p>
      * All changes made to the given HierarchyData object will also be visible
      * through this data provider.
-     * 
+     *
      * @param hierarchyData
      *            the backing HierarchyData for this provider
      */
@@ -63,7 +63,7 @@ public class InMemoryHierarchicalDataProvider<T> extends
 
     /**
      * Return the underlying hierarchical data of this provider.
-     * 
+     *
      * @return the underlying data of this provider
      */
     public HierarchyData<T> getData() {
@@ -77,6 +77,12 @@ public class InMemoryHierarchicalDataProvider<T> extends
 
     @Override
     public boolean hasChildren(T item) {
+        if (!hierarchyData.contains(item)) {
+            throw new IllegalArgumentException("Item " + item
+                    + " could not be found in the backing HierarchyData. "
+                    + "Did you forget to refresh this data provider after item removal?");
+        }
+
         return !hierarchyData.getChildren(item).isEmpty();
     }
 
@@ -89,6 +95,13 @@ public class InMemoryHierarchicalDataProvider<T> extends
     @Override
     public Stream<T> fetchChildren(
             HierarchicalQuery<T, SerializablePredicate<T>> query) {
+        if (!hierarchyData.contains(query.getParent())) {
+            throw new IllegalArgumentException("The queried item "
+                    + query.getParent()
+                    + " could not be found in the backing HierarchyData. "
+                    + "Did you forget to refresh this data provider after item removal?");
+        }
+
         Stream<T> childStream = getFilteredStream(
                 hierarchyData.getChildren(query.getParent()).stream(),
                 query.getFilter());
diff --git a/server/src/test/java/com/vaadin/data/provider/InMemoryHierarchicalDataProviderTest.java b/server/src/test/java/com/vaadin/data/provider/InMemoryHierarchicalDataProviderTest.java
index 6e4d66bb34..99b7a04b19 100644
--- a/server/src/test/java/com/vaadin/data/provider/InMemoryHierarchicalDataProviderTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/InMemoryHierarchicalDataProviderTest.java
@@ -73,6 +73,13 @@ public class InMemoryHierarchicalDataProviderTest extends
     }
 
     @Test
+    public void hierarchyData_re_add_removed_item() {
+        StrBean item = rootData.get(0);
+        data.removeItem(item).addItem(null, item);
+        Assert.assertTrue(data.getChildren(null).contains(item));
+    }
+
+    @Test
     public void setFilter() {
         getDataProvider().setFilter(item -> item.getValue().equals("Xyz")
                 || item.getValue().equals("Baz"));
diff --git a/testbench-api/src/main/java/com/vaadin/testbench/elements/TreeGridElement.java b/testbench-api/src/main/java/com/vaadin/testbench/elements/TreeGridElement.java
index 13dfbbf565..3eb8b30b5e 100644
--- a/testbench-api/src/main/java/com/vaadin/testbench/elements/TreeGridElement.java
+++ b/testbench-api/src/main/java/com/vaadin/testbench/elements/TreeGridElement.java
@@ -119,6 +119,24 @@ public class TreeGridElement extends GridElement {
     }
 
     /**
+     * Check whether the given indices correspond to a cell that contains a
+     * visible hierarchy toggle element.
+     *
+     * @param rowIndex
+     *            0-based row index
+     * @param hierarchyColumnIndex
+     *            0-based index of the hierarchy column
+     * @return {@code true} if this cell has the expand toggle visible
+     */
+    public boolean hasExpandToggle(int rowIndex, int hierarchyColumnIndex) {
+        WebElement expandElement = getExpandElement(rowIndex,
+                hierarchyColumnIndex);
+        List<String> classes = Arrays
+                .asList(expandElement.getAttribute("class").split(" "));
+        return classes.contains("expanded") || classes.contains("collapsed");
+    }
+
+    /**
      * Gets the expand/collapse element for the given row.
      *
      * @param rowIndex
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridChangingHierarchy.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridChangingHierarchy.java
new file mode 100644
index 0000000000..c7943524e4
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridChangingHierarchy.java
@@ -0,0 +1,89 @@
+package com.vaadin.tests.components.treegrid;
+
+import java.util.stream.Stream;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.ValueProvider;
+import com.vaadin.data.provider.HierarchicalQuery;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.server.SerializablePredicate;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.TreeGrid;
+
+public class TreeGridChangingHierarchy extends AbstractTestUI {
+
+    private static class TestDataProvider
+            extends InMemoryHierarchicalDataProvider<String> {
+
+        private HierarchyData<String> hierarchyData;
+
+        public TestDataProvider(HierarchyData<String> hierarchyData) {
+            super(hierarchyData);
+            this.hierarchyData = hierarchyData;
+        }
+
+        @Override
+        public boolean hasChildren(String item) {
+            if (!hierarchyData.contains(item)) {
+                return false;
+            }
+            return super.hasChildren(item);
+        }
+
+        @Override
+        public Stream<String> fetchChildren(
+                HierarchicalQuery<String, SerializablePredicate<String>> query) {
+            if (!hierarchyData.contains(query.getParent())) {
+                return Stream.empty();
+            }
+            return super.fetchChildren(query);
+        }
+    }
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        HierarchyData<String> data = new HierarchyData<>();
+        data.addItems(null, "a", "b", "c").addItem("b", "b/a");
+
+        TreeGrid<String> grid = new TreeGrid<>();
+        grid.setDataProvider(new TestDataProvider(data));
+        grid.addColumn(ValueProvider.identity());
+
+        Button btn = new Button("add items to a and refresh");
+        btn.addClickListener(event -> {
+            data.addItems("a", "a/a", "a/b");
+            grid.getDataProvider().refreshItem("a");
+        });
+        Button btn2 = new Button("add items to a/a and refresh");
+        btn2.addClickListener(event -> {
+            data.addItems("a/a", "a/a/a", "a/a/c").addItem("a/a/a", "a/a/a/a");
+            grid.getDataProvider().refreshItem("a/a");
+        });
+        Button btn3 = new Button("remove a/a");
+        btn3.addClickListener(event -> {
+            data.removeItem("a/a");
+        });
+        Button btn4 = new Button("remove children of a/a");
+        btn4.addClickListener(event -> {
+            data.removeItem("a/a/a");
+            data.removeItem("a/a/c");
+        });
+        Button btn5 = new Button("remove a");
+        btn5.addClickListener(event -> {
+            data.removeItem("a");
+        });
+        Button btn6 = new Button("remove children of a");
+        btn6.addClickListener(event -> {
+            data.removeItem("a/a");
+            data.removeItem("a/b");
+        });
+        Button btn7 = new Button("remove children of a/a/a");
+        btn7.addClickListener(event -> {
+            data.removeItem("a/a/a/a");
+        });
+
+        addComponents(grid, btn, btn2, btn3, btn4, btn5, btn6, btn7);
+    }
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridChangingHierarchyTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridChangingHierarchyTest.java
new file mode 100644
index 0000000000..89273af661
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridChangingHierarchyTest.java
@@ -0,0 +1,121 @@
+package com.vaadin.tests.components.treegrid;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.openqa.selenium.By;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridChangingHierarchyTest extends SingleBrowserTest {
+
+    private TreeGridElement grid;
+    private ButtonElement addItemsToABtn;
+    private ButtonElement addItemsToAABtn;
+    private ButtonElement removeAABtn;
+    private ButtonElement removeChildrenOfAABtn;
+    private ButtonElement removeABtn;
+    private ButtonElement removeChildrenOfABtn;
+    private ButtonElement removeChildrenOfAAABtn;
+
+    @Before
+    public void before() {
+        setDebug(true);
+        openTestURL();
+        grid = $(TreeGridElement.class).first();
+        addItemsToABtn = $(ButtonElement.class).get(0);
+        addItemsToAABtn = $(ButtonElement.class).get(1);
+        removeAABtn = $(ButtonElement.class).get(2);
+        removeChildrenOfAABtn = $(ButtonElement.class).get(3);
+        removeABtn = $(ButtonElement.class).get(4);
+        removeChildrenOfABtn = $(ButtonElement.class).get(5);
+        removeChildrenOfAAABtn = $(ButtonElement.class).get(6);
+    }
+
+    @After
+    public void after() {
+        assertNoErrorNotifications();
+        assertFalse(isElementPresent(By.className("v-errorindicator")));
+    }
+
+    @Test
+    public void removing_items_from_hierarchy() {
+        addItemsToABtn.click();
+        addItemsToAABtn.click();
+        grid.expandWithClick(0);
+        grid.expandWithClick(1);
+        grid.collapseWithClick(0);
+        removeAABtn.click();
+        // Item removed from hierarchy. when encountering less children than
+        // expected, should reset:
+        grid.expandWithClick(0);
+        // expand "a" after the reset:
+        grid.expandWithClick(0);
+        // "a/a" should be removed from a's children:
+        assertEquals("a/b", grid.getCell(1, 0).getText());
+    }
+
+    @Test
+    public void removing_all_children_from_item() {
+        addItemsToABtn.click();
+        assertTrue(grid.isRowCollapsed(0, 0));
+        // drop added children from backing data source
+        removeChildrenOfABtn.click();
+        // changes are not refreshed, thus the row should still appear as
+        // collapsed
+        assertTrue(grid.isRowCollapsed(0, 0));
+        // when encountering 0 children, will reset
+        grid.expandWithClick(0);
+        assertEquals(3, grid.getRowCount());
+        assertFalse(grid.hasExpandToggle(0, 0));
+        // verify other items still expand/collapse correctly:
+        grid.expandWithClick(1);
+        assertEquals("b/a", grid.getCell(2, 0).getText());
+        assertEquals(4, grid.getRowCount());
+        grid.collapseWithClick(1);
+        assertEquals("c", grid.getCell(2, 0).getText());
+        assertEquals(3, grid.getRowCount());
+    }
+
+    @Test
+    public void removal_of_deeply_nested_items() {
+        addItemsToABtn.click();
+        addItemsToAABtn.click();
+        grid.expandWithClick(0);
+        grid.expandWithClick(1);
+        grid.expandWithClick(2);
+        removeChildrenOfAAABtn.click();
+        grid.collapseWithClick(1);
+        // reset should get triggered here
+        grid.expandWithClick(1);
+        grid.expandWithClick(0);
+        grid.expandWithClick(1);
+        assertEquals("a/a/a", grid.getCell(2, 0).getText());
+        assertFalse(grid.hasExpandToggle(2, 0));
+    }
+
+    @Test
+    public void changing_selection_from_selected_removed_item() {
+        addItemsToABtn.click();
+        grid.expandWithClick(0);
+        grid.getCell(1, 0).click();
+        removeChildrenOfABtn.click();
+        grid.collapseWithClick(0);
+        grid.getCell(1, 0).click();
+        assertTrue(grid.getRow(1).isSelected());
+    }
+
+    @Test
+    public void remove_item_from_root() {
+        addItemsToABtn.click();
+        removeABtn.click();
+        grid.expandWithClick(0);
+        assertEquals("b", grid.getCell(0, 0).getText());
+    }
+}