Improve expand and collapse of items in TreeGrid (#9159)

Fixes a race condition when expanding multiple items.
Only one expand or collapse request should be sent from
the client before waiting for a response, otherwise the
indexing in subsequent requests will be incorrect.

Adds API to collapse and expand multiple items from the
server, reducing the amount of round trips with multiple
item expands and collapses.

HierarchyMapper now correctly keeps expanded nodes expanded
if their parent is collapsed.

6 files changed, 172 insertions, 25 deletions

diff --git a/uitest/src/main/java/com/vaadin/tests/performance/TreeGridMemory.java b/uitest/src/main/java/com/vaadin/tests/performance/TreeGridMemory.java
index 105a05fe49..531cd8e283 100644
--- a/uitest/src/main/java/com/vaadin/tests/performance/TreeGridMemory.java
+++ b/uitest/src/main/java/com/vaadin/tests/performance/TreeGridMemory.java
@@ -1,5 +1,6 @@
 package com.vaadin.tests.performance;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
@@ -8,6 +9,7 @@ import javax.servlet.annotation.WebServlet;
 import com.vaadin.annotations.VaadinServletConfiguration;
 import com.vaadin.data.HierarchyData;
 import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinServlet;
 import com.vaadin.tests.data.bean.Address;
 import com.vaadin.tests.data.bean.Person;
@@ -23,6 +25,16 @@ public class TreeGridMemory extends AbstractBeansMemoryTest<TreeGrid<Person>> {
     public static class Servlet extends VaadinServlet {
     }
 
+    private boolean initiallyExpanded = false;
+
+    @Override
+    protected void init(VaadinRequest request) {
+        if (request.getParameter("initiallyExpanded") != null) {
+            initiallyExpanded = true;
+        }
+        super.init(request);
+    }
+
     @Override
     protected TreeGrid<Person> createComponent() {
         TreeGrid<Person> treeGrid = new TreeGrid<>();
@@ -43,20 +55,25 @@ public class TreeGridMemory extends AbstractBeansMemoryTest<TreeGrid<Person>> {
     protected void setInMemoryContainer(TreeGrid<Person> treeGrid,
             List<Person> data) {
         HierarchyData<Person> hierarchyData = new HierarchyData<>();
-        if (data.size() % 2 == 0) {
+        treeGrid.setDataProvider(
+                new InMemoryHierarchicalDataProvider<>(hierarchyData));
+        List<Person> toExpand = new ArrayList<>();
+        if (data.size() != 0 && data.size() % 2 == 0) {
             // treat list as if it were a balanced binary tree
             hierarchyData.addItem(null, data.get(0));
             int n = 0;
             while (2 * n + 2 < data.size()) {
                 hierarchyData.addItems(data.get(n),
                         data.subList(2 * n + 1, 2 * n + 3));
+                toExpand.add(data.get(n));
                 n++;
             }
         } else {
             hierarchyData.addItems(null, data);
         }
-        treeGrid.setDataProvider(
-                new InMemoryHierarchicalDataProvider<>(hierarchyData));
+        if (initiallyExpanded) {
+            treeGrid.expand(toExpand);
+        }
     }
 
     @Override
diff --git a/uitest/src/main/java/com/vaadin/tests/performance/TreeTableMemory.java b/uitest/src/main/java/com/vaadin/tests/performance/TreeTableMemory.java
index 203c74fd4d..d62bcd046f 100644
--- a/uitest/src/main/java/com/vaadin/tests/performance/TreeTableMemory.java
+++ b/uitest/src/main/java/com/vaadin/tests/performance/TreeTableMemory.java
@@ -6,6 +6,7 @@ import java.util.Optional;
 import javax.servlet.annotation.WebServlet;
 
 import com.vaadin.annotations.VaadinServletConfiguration;
+import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinServlet;
 import com.vaadin.tests.data.bean.Address;
 import com.vaadin.tests.data.bean.Person;
@@ -23,6 +24,16 @@ public class TreeTableMemory extends AbstractBeansMemoryTest<TreeTable> {
     public static class Servlet extends VaadinServlet {
     }
 
+    private boolean initiallyExpanded = false;
+
+    @Override
+    protected void init(VaadinRequest request) {
+        if (request.getParameter("initiallyExpanded") != null) {
+            initiallyExpanded = true;
+        }
+        super.init(request);
+    }
+
     @Override
     protected TreeTable createComponent() {
         TreeTable treeTable = new TreeTable();
@@ -38,14 +49,19 @@ public class TreeTableMemory extends AbstractBeansMemoryTest<TreeTable> {
         container.addContainerProperty("street", String.class, null);
         container.addContainerProperty("zip", String.class, null);
         container.addContainerProperty("city", String.class, null);
+        treeTable.setContainerDataSource(container);
 
-        if (data.size() % 2 == 0) {
+        if (data.size() != 0 && data.size() % 2 == 0) {
             createItem(0, container, data);
+            treeTable.setCollapsed(0, false);
             int n = 0;
             while (2 * n + 2 < data.size()) {
                 for (int i : new Integer[] { 1, 2 }) {
                     createItem(2 * n + i, container, data);
                     container.setParent(2 * n + i, n);
+                    if (initiallyExpanded) {
+                        treeTable.setCollapsed(2 * n + i, false);
+                    }
                 }
                 n++;
             }
@@ -54,7 +70,6 @@ public class TreeTableMemory extends AbstractBeansMemoryTest<TreeTable> {
                 createItem(i, container, data);
             }
         }
-        treeTable.setContainerDataSource(container);
     }
 
     private void createItem(int index, HierarchicalContainer container,
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
index a8328bd007..b1b4da2c74 100644
--- a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
@@ -95,12 +95,6 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         assertEquals(3, grid.getRowCount());
         assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
-        // 1 | 1 should not be expanded this time
-        selectMenuPath("Component", "Features", "Server-side expand",
-                "Expand 0 | 0");
-        assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
-
         assertNoSystemNotifications();
         assertNoErrorNotifications();
     }
@@ -280,6 +274,26 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
                 "Item collapsed (user originated: true): 0 | 1"));
     }
 
+    @Test
+    public void expanded_nodes_stay_expanded_when_parent_expand_state_is_toggled() {
+        grid.expandWithClick(0);
+        grid.expandWithClick(1);
+        grid.collapseWithClick(0);
+        grid.expandWithClick(0);
+        assertCellTexts(0, 0, new String[] { "0 | 0", "1 | 0", "2 | 0", "2 | 1",
+                "2 | 2", "1 | 1", "1 | 2", "0 | 1", "0 | 2" });
+        assertEquals(9, grid.getRowCount());
+
+        grid.expandWithClick(7);
+        grid.expandWithClick(8);
+        grid.collapseWithClick(7);
+        grid.collapseWithClick(0);
+        grid.expandWithClick(1);
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "1 | 0", "2 | 0",
+                "2 | 1", "2 | 2", "1 | 1", "1 | 2", "0 | 2" });
+        assertEquals(9, grid.getRowCount());
+    }
+
     private void assertCellTexts(int startRowIndex, int cellIndex,
             String[] cellTexts) {
         int index = startRowIndex;
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridCollapseExpandTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridCollapseExpandTest.java
new file mode 100644
index 0000000000..59ff75ecd0
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridCollapseExpandTest.java
@@ -0,0 +1,62 @@
+package com.vaadin.tests.components.treegrid;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.support.ui.ExpectedCondition;
+
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTestPhantomJS2;
+
+public class TreeGridCollapseExpandTest extends SingleBrowserTestPhantomJS2 {
+
+    private TreeGridElement grid;
+
+    @Override
+    public Class<?> getUIClass() {
+        return TreeGridBasicFeatures.class;
+    }
+
+    @Test
+    public void no_race_condition_with_multiple_collapse_or_expand() {
+        openTestURL();
+        grid = $(TreeGridElement.class).first();
+        testBench().disableWaitForVaadin();
+
+        // toggle expand of two rows simultaneously
+        // only the first of the expands should occur
+        executeScript("arguments[0].click(); arguments[1].click()",
+                grid.getExpandElement(0, 0), grid.getExpandElement(1, 0));
+        waitUntilRowCountEquals(6);
+        assertCellTexts(0, 0,
+                new String[] { "0 | 0", "1 | 0", "1 | 1", "1 | 2", "0 | 1" });
+
+        // toggle collapse of the expanded first row and immediately expand the
+        // last row
+        // only the collapse should occur
+        executeScript("arguments[0].click(); arguments[1].click()",
+                grid.getExpandElement(0, 0), grid.getExpandElement(5, 0));
+        waitUntilRowCountEquals(3);
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
+    }
+
+    private void assertCellTexts(int startRowIndex, int cellIndex,
+            String[] cellTexts) {
+        int index = startRowIndex;
+        for (String cellText : cellTexts) {
+            assertEquals(cellText,
+                    grid.getRow(index).getCell(cellIndex).getText());
+            index++;
+        }
+    }
+
+    private void waitUntilRowCountEquals(int expectedCount) {
+        waitUntil(new ExpectedCondition<Boolean>() {
+            @Override
+            public Boolean apply(WebDriver input) {
+                return grid.getRowCount() == expectedCount;
+            }
+        });
+    }
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java
index 1d5d9b309d..1d99534719 100644
--- a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java
@@ -1,10 +1,12 @@
 package com.vaadin.tests.components.treegrid;
 
+import org.apache.commons.lang3.StringUtils;
 import org.junit.Assert;
 import org.junit.Test;
 
 import com.vaadin.testbench.elements.ButtonElement;
 import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.performance.TreeGridMemory;
 import com.vaadin.tests.tb3.SingleBrowserTest;
 
 public class TreeGridHugeTreeTest extends SingleBrowserTest {
@@ -48,6 +50,38 @@ public class TreeGridHugeTreeTest extends SingleBrowserTest {
                 "Granddad 2", "Dad 2/0" });
     }
 
+    @Test
+    public void collapsed_rows_invalidated_correctly() {
+        openTestURL();
+        grid = $(TreeGridElement.class).first();
+        grid.expandWithClick(2);
+        grid.expandWithClick(3);
+        grid.expandWithClick(0);
+        grid.collapseWithClick(0);
+        grid.expandWithClick(0);
+        grid.expandWithClick(1);
+        assertCellTexts(0, 0,
+                new String[] { "Granddad 0", "Dad 0/0", "Son 0/0/0" });
+    }
+
+    @Test
+    public void collapsed_subtrees_outside_of_cache_stay_expanded() {
+        getDriver().get(StringUtils.strip(getBaseURL(), "/")
+                + TreeGridMemory.PATH + "?items=200&initiallyExpanded");
+        grid = $(TreeGridElement.class).first();
+
+        String[] cellTexts = new String[100];
+        for (int i = 0; i < 100; i++) {
+            cellTexts[i] = grid.getRow(i).getCell(0).getText();
+        }
+        grid.scrollToRow(0);
+
+        grid.collapseWithClick(1);
+        grid.expandWithClick(1);
+
+        assertCellTexts(0, 0, cellTexts);
+    }
+
     private void assertCellTexts(int startRowIndex, int cellIndex,
             String[] cellTexts) {
         int index = startRowIndex;
diff --git a/uitest/src/test/java/com/vaadin/tests/performance/MemoryIT.java b/uitest/src/test/java/com/vaadin/tests/performance/MemoryIT.java
index b30dcdae4f..d3082ec336 100644
--- a/uitest/src/test/java/com/vaadin/tests/performance/MemoryIT.java
+++ b/uitest/src/test/java/com/vaadin/tests/performance/MemoryIT.java
@@ -34,19 +34,26 @@ public class MemoryIT extends SingleBrowserTest {
 
     @Test
     public void measureMemory() {
-        performTest(GridMemory.PATH, 1, "grid-v8-one-item-");
-        performTest(CompatibilityGridMemory.PATH, 1, "grid-v7-one-item-");
+        performTest(GridMemory.PATH + "?items=1", "grid-v8-one-item-");
+        performTest(CompatibilityGridMemory.PATH + "?items=1",
+                "grid-v7-one-item-");
 
-        performTest(GridMemory.PATH, 100000, "grid-v8-100thousand-items-");
-        performTest(CompatibilityGridMemory.PATH, 100000,
+        performTest(GridMemory.PATH + "?items=1", "grid-v8-100thousand-items-");
+        performTest(CompatibilityGridMemory.PATH + "?items=100000",
                 "grid-v7-100thousand-items-");
 
-        performTest(TreeGridMemory.PATH, 1, "tree-grid-one-item-");
-        performTest(TreeTableMemory.PATH, 1, "tree-table-one-item-");
+        performTest(TreeGridMemory.PATH + "?items=1", "tree-grid-one-item-");
+        performTest(TreeTableMemory.PATH + "?items=1", "tree-table-one-item-");
 
-        performTest(TreeGridMemory.PATH, 100000,
+
+        performTest(TreeGridMemory.PATH + "?items=100&initiallyExpanded",
+                "tree-grid-100-items-initially-expanded-");
+        performTest(TreeTableMemory.PATH + "?items=100&initiallyExpanded",
+                "tree-table-100-items-initially-expanded-");
+
+        performTest(TreeGridMemory.PATH + "?items=100000",
                 "tree-grid-100thousand-items-");
-        performTest(TreeTableMemory.PATH, 100000,
+        performTest(TreeTableMemory.PATH + "?items=100000",
                 "tree-table-100thousand-items-");
     }
 
@@ -54,14 +61,13 @@ public class MemoryIT extends SingleBrowserTest {
     protected void closeApplication() {
     }
 
-    private void performTest(String path, int itemsCount,
-            String teamcityStatPrefix) {
+    private void performTest(String path, String teamcityStatPrefix) {
         double lastResult = 0;
         int stableNumber = 0;
         List<Long> renderingTimes = new ArrayList<>();
         List<Long> requestTimes = new ArrayList<>();
         for (int i = 0; i < MAX_ITERATIONS; i++) {
-            openUI(path, itemsCount);
+            openUI(path);
             renderingTimes.add(testBench().totalTimeSpentRendering());
             requestTimes.add(testBench().totalTimeSpentServicingRequests());
             long currentResult = Long
@@ -104,9 +110,8 @@ public class MemoryIT extends SingleBrowserTest {
         return delta < deltaLimit;
     }
 
-    private void openUI(String path, int itemsNumber) {
-        getDriver().get(StringUtils.strip(getBaseURL(), "/") + path + "?items="
-                + itemsNumber);
+    private void openUI(String path) {
+        getDriver().get(StringUtils.strip(getBaseURL(), "/") + path);
         Assert.assertTrue(isElementPresent(By.className("v-grid"))
                 || isElementPresent(By.className("v-treegrid"))
                 || isElementPresent(By.className("v-table")));