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.

client/src/main/java/com/vaadin/client/connectors/data/DataCommunicatorConnector.java

@@ -25,6 +25,7 @@ import com.vaadin.client.data.AbstractRemoteDataSource;
 import com.vaadin.client.data.DataSource;
 import com.vaadin.client.extensions.AbstractExtensionConnector;
 import com.vaadin.data.provider.DataCommunicator;
+import com.vaadin.shared.Range;
 import com.vaadin.shared.data.DataCommunicatorClientRpc;
 import com.vaadin.shared.data.DataCommunicatorConstants;
 import com.vaadin.shared.data.DataRequestRpc;
@@ -97,15 +98,7 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
             getRpcProxy(DataRequestRpc.class).requestRows(firstRowIndex,
                     numberOfRows, getCachedRange().getStart(),
                     getCachedRange().length());
-
-            JsonArray dropped = Json.createArray();
-            int i = 0;
-            for (String key : droppedKeys) {
-                dropped.set(i++, key);
-            }
-            droppedKeys.clear();
-
-            getRpcProxy(DataRequestRpc.class).dropRows(dropped);
+            sendDroppedRows();
         }
 
         @Override
@@ -113,6 +106,12 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
             return row.getString(DataCommunicatorConstants.KEY);
         }
 
+        @Override
+        protected void dropFromCache(Range range) {
+            super.dropFromCache(range);
+            sendDroppedRows();
+        }
+
         @Override
         protected void onDropFromCache(int rowIndex, JsonObject removed) {
             droppedKeys.add(getRowKey(removed));
@@ -135,6 +134,22 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
                 setRowData(index, Collections.singletonList(rowData));
             }
         }
+
+        /**
+         * Inform the server of any dropped rows.
+         */
+        private void sendDroppedRows() {
+            if (!droppedKeys.isEmpty()) {
+                JsonArray dropped = Json.createArray();
+                int i = 0;
+                for (String key : droppedKeys) {
+                    dropped.set(i++, key);
+                }
+                droppedKeys.clear();
+
+                getRpcProxy(DataRequestRpc.class).dropRows(dropped);
+            }
+        }
     }
 
     private DataSource<JsonObject> ds = new VaadinDataSource();

client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java

@@ -17,6 +17,7 @@ package com.vaadin.client.connectors.treegrid;
 
 import java.util.Collection;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 import java.util.logging.Logger;
 
@@ -56,6 +57,10 @@ import elemental.json.JsonObject;
 @Connect(com.vaadin.ui.TreeGrid.class)
 public class TreeGridConnector extends GridConnector {
 
+    private static enum AwaitingRowsState {
+        NONE, COLLAPSE, EXPAND
+    }
+
     public TreeGridConnector() {
         registerRpc(FocusRpc.class, (rowIndex, cellIndex) -> {
             getWidget().focusCell(rowIndex, cellIndex);
@@ -68,6 +73,8 @@ public class TreeGridConnector extends GridConnector {
 
     private Set<String> rowKeysPendingExpand = new HashSet<>();
 
+    private AwaitingRowsState awaitingRowsState = AwaitingRowsState.NONE;
+
     @Override
     public TreeGrid getWidget() {
         return (TreeGrid) super.getWidget();
@@ -158,16 +165,14 @@ public class TreeGridConnector extends GridConnector {
         registerRpc(TreeGridClientRpc.class, new TreeGridClientRpc() {
 
             @Override
-            public void setExpanded(String key) {
-                rowKeysPendingExpand.add(key);
-                Range cache = ((AbstractRemoteDataSource) getDataSource())
-                        .getCachedRange();
-                checkExpand(cache.getStart(), cache.length());
+            public void setExpanded(List<String> keys) {
+                rowKeysPendingExpand.addAll(keys);
+                checkExpand();
             }
 
             @Override
-            public void setCollapsed(String key) {
-                rowKeysPendingExpand.remove(key);
+            public void setCollapsed(List<String> keys) {
+                rowKeysPendingExpand.removeAll(keys);
             }
 
             @Override
@@ -189,12 +194,18 @@ public class TreeGridConnector extends GridConnector {
 
             @Override
             public void dataRemoved(int firstRowIndex, int numberOfRows) {
-                // NO-OP
+                if (awaitingRowsState == AwaitingRowsState.COLLAPSE) {
+                    awaitingRowsState = AwaitingRowsState.NONE;
+                }
+                checkExpand();
             }
 
             @Override
             public void dataAdded(int firstRowIndex, int numberOfRows) {
-                // NO-OP
+                if (awaitingRowsState == AwaitingRowsState.EXPAND) {
+                    awaitingRowsState = AwaitingRowsState.NONE;
+                }
+                checkExpand();
             }
 
             @Override
@@ -204,7 +215,7 @@ public class TreeGridConnector extends GridConnector {
 
             @Override
             public void resetDataAndSize(int estimatedNewDataSize) {
-                // NO-OP
+                awaitingRowsState = AwaitingRowsState.NONE;
             }
         });
     }
@@ -232,16 +243,43 @@ public class TreeGridConnector extends GridConnector {
         return cell.getColumn().getRenderer() instanceof HierarchyRenderer;
     }
 
+    /**
+     * Delegates to {@link #setCollapsed(int, boolean, boolean)}, with
+     * {@code userOriginated} as {@code true}.
+     *
+     * @see #setCollapsed(int, boolean, boolean)
+     */
     private void setCollapsed(int rowIndex, boolean collapsed) {
-        String rowKey = getRowKey(getDataSource().getRow(rowIndex));
-        getRpcProxy(NodeCollapseRpc.class).setNodeCollapsed(rowKey, rowIndex,
-                collapsed, true);
+        setCollapsed(rowIndex, collapsed, true);
     }
 
-    private void setCollapsedServerInitiated(int rowIndex, boolean collapsed) {
+    /**
+     * Set the collapse state for the row in the given index.
+     * <p>
+     * Calling this method will have no effect if a response has not yet been
+     * received for a previous call to this method.
+     *
+     * @param rowIndex
+     *            index of the row to set the state for
+     * @param collapsed
+     *            {@code true} to collapse the row, {@code false} to expand the
+     *            row
+     * @param userOriginated
+     *            whether this method was originated from a user interaction
+     */
+    private void setCollapsed(int rowIndex, boolean collapsed,
+            boolean userOriginated) {
+        if (isAwaitingRowChange()) {
+            return;
+        }
+        if (collapsed) {
+            awaitingRowsState = AwaitingRowsState.COLLAPSE;
+        } else {
+            awaitingRowsState = AwaitingRowsState.EXPAND;
+        }
         String rowKey = getRowKey(getDataSource().getRow(rowIndex));
         getRpcProxy(NodeCollapseRpc.class).setNodeCollapsed(rowKey, rowIndex,
-                collapsed, false);
+                collapsed, userOriginated);
     }
 
     /**
@@ -347,8 +385,20 @@ public class TreeGridConnector extends GridConnector {
         }
     }
 
+    private boolean isAwaitingRowChange() {
+        return awaitingRowsState != AwaitingRowsState.NONE;
+    }
+
+    private void checkExpand() {
+        Range cache = ((AbstractRemoteDataSource) getDataSource())
+                .getCachedRange();
+        checkExpand(cache.getStart(), cache.length());
+    }
+
     private void checkExpand(int firstRowIndex, int numberOfRows) {
-        if (rowKeysPendingExpand.isEmpty()) {
+        if (rowKeysPendingExpand.isEmpty() || isAwaitingRowChange()) {
+            // will not perform the check if an expand or collapse action is
+            // already pending or there are no rows pending expand
             return;
         }
         for (int rowIndex = firstRowIndex; rowIndex < firstRowIndex
@@ -356,7 +406,7 @@ public class TreeGridConnector extends GridConnector {
             String rowKey = getDataSource().getRow(rowIndex)
                     .getString(DataCommunicatorConstants.KEY);
             if (rowKeysPendingExpand.remove(rowKey)) {
-                setCollapsedServerInitiated(rowIndex, false);
+                setCollapsed(rowIndex, false, false);
                 return;
             }
         }

client/src/main/java/com/vaadin/client/data/AbstractRemoteDataSource.java

@@ -356,7 +356,13 @@ public abstract class AbstractRemoteDataSource<T> implements DataSource<T> {
         dropFromCache(cacheParition[2]);
     }
 
-    private void dropFromCache(Range range) {
+    /**
+     * Drop the given range of rows from this data source's cache.
+     *
+     * @param range
+     *            the range of rows to drop
+     */
+    protected void dropFromCache(Range range) {
         for (int i = range.getStart(); i < range.getEnd(); i++) {
             // Called after dropping from cache. Dropped row is passed as a
             // parameter, but is no longer present in the DataSource

server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java

@@ -283,10 +283,11 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     @Override
     protected void onDropRows(JsonArray keys) {
         for (int i = 0; i < keys.length(); i++) {
-            // cannot drop expanded rows since the parent item is needed always
-            // when fetching more rows
+            // cannot drop keys of expanded rows, parents of expanded rows or
+            // rows that are pending expand
             String itemKey = keys.getString(i);
-            if (mapper.isCollapsed(itemKey) && !rowKeysPendingExpand.contains(itemKey)) {
+            if (!mapper.isKeyStored(itemKey)
+                    && !rowKeysPendingExpand.contains(itemKey)) {
                 getActiveDataHandler().dropActiveData(itemKey);
             }
         }
@@ -390,9 +391,15 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
      *            the key of the row, not {@code null}
      * @param expandedRowIndex
      *            the index of the row to expand
+     * @param userOriginated
+     *            whether this expand was originated from the server or client
      * @return {@code true} if the row was expanded, {@code false} otherwise
      */
-    public boolean doExpand(String expandedRowKey, final int expandedRowIndex) {
+    public boolean doExpand(String expandedRowKey, final int expandedRowIndex,
+            boolean userOriginated) {
+        if (!userOriginated && !rowKeysPendingExpand.contains(expandedRowKey)) {
+            return false;
+        }
         if (expandedRowIndex < 0 | expandedRowIndex >= mapper.getTreeSize()) {
             throw new IllegalArgumentException("Invalid row index "
                     + expandedRowIndex + " when tree grid size of "
@@ -403,7 +410,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         Objects.requireNonNull(expandedItem,
                 "Cannot find item for given key " + expandedRowKey);
 
-        final int expandedNodeSize = doSizeQuery(expandedItem);
+        int expandedNodeSize = doSizeQuery(expandedItem);
         if (expandedNodeSize == 0) {
             // TODO handle 0 size -> not expandable
             throw new IllegalStateException("Row with index " + expandedRowIndex
@@ -413,7 +420,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         if (!mapper.isCollapsed(expandedRowKey)) {
             return false;
         }
-        mapper.expand(expandedRowKey, expandedRowIndex, expandedNodeSize);
+        expandedNodeSize = mapper.expand(expandedRowKey, expandedRowIndex,
+                expandedNodeSize);
         rowKeysPendingExpand.remove(expandedRowKey);
 
         getClientRpc().insertRows(expandedRowIndex + 1, expandedNodeSize);

server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java

@@ -17,9 +17,13 @@ package com.vaadin.data.provider;
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.TreeSet;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
@@ -139,6 +143,13 @@ class HierarchyMapper implements Serializable {
     /** The expanded nodes in the tree. */
     private final TreeSet<TreeNode> nodes = new TreeSet<>();
 
+    /**
+     * Map of collapsed subtrees. The keys of this map are the collapsed
+     * subtrees parent keys and values are the {@code TreeSet}s of the subtree's
+     * expanded {@code TreeNode}s.
+     */
+    private final Map<String, TreeSet<TreeNode>> collapsedNodes = new HashMap<>();
+
     /**
      * Resets the tree, sets given the root level size.
      *
@@ -146,6 +157,7 @@ class HierarchyMapper implements Serializable {
      *            the number of items in the root level
      */
     public void reset(int rootLevelSize) {
+        collapsedNodes.clear();
         nodes.clear();
         nodes.add(new TreeNode(null, 0, rootLevelSize));
     }
@@ -172,6 +184,33 @@ class HierarchyMapper implements Serializable {
         return !getNodeForKey(itemKey).isPresent();
     }
 
+    /**
+     * Return whether the given item key is still being used in this mapper.
+     *
+     * @param itemKey
+     *            the item key to look for
+     * @return {@code true} if the item key is still used, {@code false}
+     *         otherwise
+     */
+    public boolean isKeyStored(String itemKey) {
+        if (getNodeForKey(itemKey).isPresent()) {
+            return true;
+        }
+        // Is the key used in a collapsed subtree?
+        for (Entry<String, TreeSet<TreeNode>> entry : collapsedNodes.entrySet()) {
+            if (entry.getKey() != null && entry.getKey().equals(itemKey)) {
+                return true;
+            }
+            for (TreeNode subTreeNode : entry.getValue()) {
+                if (subTreeNode.getParentKey() != null
+                        && subTreeNode.getParentKey().equals(itemKey)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     /**
      * Return the depth of expanded node's subtree.
      * <p>
@@ -222,24 +261,37 @@ class HierarchyMapper implements Serializable {
     /**
      * Expands the node in the given index and with the given key.
      *
-     * @param expanedRowKey
+     * @param expandedRowKey
      *            the key of the expanded item
      * @param expandedRowIndex
      *            the index of the expanded item
      * @param expandedNodeSize
-     *            the size of the subtree of the expanded node
+     *            the size of the subtree of the expanded node, used if
+     *            previously unknown
      * @throws IllegalStateException
      *             if the node was expanded already
+     * @return the actual size of the expand
      */
-    protected void expand(String expanedRowKey, int expandedRowIndex,
+    protected int expand(String expandedRowKey, int expandedRowIndex,
             int expandedNodeSize) {
         if (expandedNodeSize < 1) {
             throw new IllegalArgumentException(
                     "The expanded node's size cannot be less than 1, was "
                             + expandedNodeSize);
         }
-        TreeNode newNode = new TreeNode(expanedRowKey, expandedRowIndex + 1,
-                expandedNodeSize);
+        TreeNode newNode;
+        TreeSet<TreeNode> subTree = null;
+        if (collapsedNodes.containsKey(expandedRowKey)) {
+            subTree = collapsedNodes.remove(expandedRowKey);
+            newNode = subTree.first();
+            int offset = expandedRowIndex - newNode.getStartIndex() + 1;
+            subTree.forEach(node -> node.push(offset));
+            expandedNodeSize = newNode.getEndIndex() - newNode.getStartIndex()
+                    + 1;
+        } else {
+            newNode = new TreeNode(expandedRowKey, expandedRowIndex + 1,
+                    expandedNodeSize);
+        }
 
         boolean added = nodes.add(newNode);
         if (!added) {
@@ -248,19 +300,26 @@ class HierarchyMapper implements Serializable {
         }
 
         // push end indexes for parent nodes
+        final int expandSize = expandedNodeSize;
         List<TreeNode> updated = nodes.headSet(newNode, false).stream()
                 .filter(node -> node.getEndIndex() >= expandedRowIndex)
                 .collect(Collectors.toList());
         nodes.removeAll(updated);
-        updated.stream().forEach(node -> node.pushEnd(expandedNodeSize));
+        updated.stream().forEach(node -> node.pushEnd(expandSize));
         nodes.addAll(updated);
 
         // push start and end indexes for later nodes
         updated = nodes.tailSet(newNode, false).stream()
                 .collect(Collectors.toList());
         nodes.removeAll(updated);
-        updated.stream().forEach(node -> node.push(expandedNodeSize));
+        updated.stream().forEach(node -> node.push(expandSize));
         nodes.addAll(updated);
+
+        if (subTree != null) {
+            nodes.addAll(subTree);
+        }
+
+        return expandSize;
     }
 
     /**
@@ -291,6 +350,12 @@ class HierarchyMapper implements Serializable {
                     + " is different for the collapsed node " + collapsedNode);
         }
 
+        Set<TreeNode> subTreeNodes = nodes
+                .tailSet(collapsedNode).stream().filter(node -> collapsedNode
+                        .getEndIndex() >= node.getEndIndex())
+                .collect(Collectors.toSet());
+        collapsedNodes.put(collapsedNode.getParentKey(), new TreeSet<>(subTreeNodes));
+
         // remove complete subtree
         AtomicInteger removedSubTreeSize = new AtomicInteger(
                 collapsedNode.getEndIndex() - collapsedNode.startIndex + 1);

server/src/main/java/com/vaadin/ui/TreeGrid.java

@@ -18,6 +18,7 @@ package com.vaadin.ui;
 import java.io.Serializable;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
@@ -227,8 +228,8 @@ public class TreeGrid<T> extends Grid<T> {
                                 .get(rowKey), true);
                     }
                 } else {
-                    if (getDataCommunicator().doExpand(rowKey, rowIndex)
-                            && userOriginated) {
+                    if (getDataCommunicator().doExpand(rowKey, rowIndex,
+                            userOriginated) && userOriginated) {
                         fireExpandEvent(getDataCommunicator().getKeyMapper()
                                 .get(rowKey), true);
                     }
@@ -441,34 +442,69 @@ public class TreeGrid<T> extends Grid<T> {
     }
 
     /**
-     * Expands the given item.
+     * Expands the given items.
      * <p>
-     * If the item is currently expanded, does nothing. If the item does not
-     * have any children, does nothing.
+     * If an item is currently expanded, does nothing. If an item does not have
+     * any children, does nothing.
      *
-     * @param item
-     *            the item to expand
+     * @param items
+     *            the items to expand
      */
-    public void expand(T item) {
-        getDataCommunicator().setPendingExpand(item).ifPresent(key -> {
-            getRpcProxy(TreeGridClientRpc.class).setExpanded(key);
-            fireExpandEvent(item, false);
-        });
+    public void expand(T... items) {
+        expand(Arrays.asList(items));
     }
 
     /**
-     * Collapses the given item.
+     * Expands the given items.
      * <p>
-     * If the item is already collapsed, does nothing.
+     * If an item is currently expanded, does nothing. If an item does not have
+     * any children, does nothing.
      *
-     * @param item
-     *            the item to collapse
+     * @param items
+     *            the items to expand
      */
-    public void collapse(T item) {
-        getDataCommunicator().collapseItem(item).ifPresent(key -> {
-            getRpcProxy(TreeGridClientRpc.class).setCollapsed(key);
-            fireCollapseEvent(item, false);
-        });
+    public void expand(Collection<T> items) {
+        List<String> expandedKeys = new ArrayList<>();
+        List<T> expandedItems = new ArrayList<>();
+        items.forEach(item -> getDataCommunicator().setPendingExpand(item)
+                .ifPresent(key -> {
+                    expandedKeys.add(key);
+                    expandedItems.add(item);
+                }));
+        getRpcProxy(TreeGridClientRpc.class).setExpanded(expandedKeys);
+        expandedItems.forEach(item -> fireExpandEvent(item, false));
+    }
+
+    /**
+     * Collapse the given items.
+     * <p>
+     * For items that are already collapsed, does nothing.
+     *
+     * @param items
+     *            the collection of items to collapse
+     */
+    public void collapse(T... items) {
+        collapse(Arrays.asList(items));
+    }
+
+    /**
+     * Collapse the given items.
+     * <p>
+     * For items that are already collapsed, does nothing.
+     *
+     * @param items
+     *            the collection of items to collapse
+     */
+    public void collapse(Collection<T> items) {
+        List<String> collapsedKeys = new ArrayList<>();
+        List<T> collapsedItems = new ArrayList<>();
+        items.forEach(item -> getDataCommunicator().collapseItem(item)
+                .ifPresent(key -> {
+                    collapsedKeys.add(key);
+                    collapsedItems.add(item);
+                }));
+        getRpcProxy(TreeGridClientRpc.class).setCollapsed(collapsedKeys);
+        collapsedItems.forEach(item -> fireCollapseEvent(item, false));
     }
 
     @Override

shared/src/main/java/com/vaadin/shared/ui/treegrid/TreeGridClientRpc.java

@@ -15,6 +15,8 @@
  */
 package com.vaadin.shared.ui.treegrid;
 
+import java.util.List;
+
 import com.vaadin.shared.communication.ClientRpc;
 
 /**
@@ -26,22 +28,22 @@ import com.vaadin.shared.communication.ClientRpc;
 public interface TreeGridClientRpc extends ClientRpc {
 
     /**
-     * Inform the client that an item with the given key has been expanded by
-     * the server.
+     * Inform the client that a list of items with the given keys have been
+     * expanded by the server.
      *
-     * @param key
-     *            the communication key of the expanded item
+     * @param keys
+     *            the communication keys of the expanded items
      */
-    public void setExpanded(String key);
+    public void setExpanded(List<String> keys);
 
     /**
-     * Inform the client that an item with the given key has been collapsed by
-     * the server.
+     * Inform the client that a list of items with the given keys have been
+     * collapsed by the server.
      *
-     * @param key
-     *            the communication key of the collapsed item
+     * @param keys
+     *            the communication keys of the collapsed items
      */
-    public void setCollapsed(String key);
+    public void setCollapsed(List<String> keys);
 
     /**
      * Clear all pending expands from the client.

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

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,

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;

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;
+            }
+        });
+    }
+}

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;

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")));