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.

3 files changed, 98 insertions, 27 deletions

diff --git a/client/src/main/java/com/vaadin/client/connectors/data/DataCommunicatorConnector.java b/client/src/main/java/com/vaadin/client/connectors/data/DataCommunicatorConnector.java
index 1b55d61a4e..f7669c53ca 100644
--- a/client/src/main/java/com/vaadin/client/connectors/data/DataCommunicatorConnector.java
+++ b/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
@@ -114,6 +107,12 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
         }
 
         @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();
diff --git a/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java b/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
index 2075b6cff1..11848f69e0 100644
--- a/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
+++ b/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;
             }
         }
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 d1ff848401..6ec549a16c 100644
--- a/client/src/main/java/com/vaadin/client/data/AbstractRemoteDataSource.java
+++ b/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