Fix Grid handling state changes and RPC calls deferred (#16188)

Some column changes in Grid were not correctly handled by the
RpcDataProviderExtension. These cases are now correctly handled.

Change-Id: I966b1c71d26e77e30e7dd84f26ab9704bd4f1f0f

3 files changed, 120 insertions, 138 deletions

diff --git a/client/src/com/vaadin/client/connectors/GridConnector.java b/client/src/com/vaadin/client/connectors/GridConnector.java
index 450df31b36..593a8c8574 100644
--- a/client/src/com/vaadin/client/connectors/GridConnector.java
+++ b/client/src/com/vaadin/client/connectors/GridConnector.java
@@ -27,8 +27,6 @@ import java.util.Map;
 import java.util.Set;
 import java.util.logging.Logger;
 
-import com.google.gwt.core.client.Scheduler;
-import com.google.gwt.core.client.Scheduler.ScheduledCommand;
 import com.google.gwt.user.client.ui.Widget;
 import com.vaadin.client.ComponentConnector;
 import com.vaadin.client.ConnectorHierarchyChangeEvent;
@@ -160,12 +158,13 @@ public class GridConnector extends AbstractHasComponentsConnector implements
         public Object getValue(final JsonObject obj) {
             final JsonObject rowData = obj.getObject(GridState.JSONKEY_DATA);
 
-            assert rowData.hasKey(id) : "Could not find data for column with id "
-                    + id;
+            if (rowData.hasKey(id)) {
+                final JsonValue columnValue = rowData.get(id);
 
-            final JsonValue columnValue = rowData.get(id);
+                return rendererConnector.decode(columnValue);
+            }
 
-            return rendererConnector.decode(columnValue);
+            return null;
         }
 
         /*
@@ -202,19 +201,8 @@ public class GridConnector extends AbstractHasComponentsConnector implements
 
                 @Override
                 public void bind(final int rowIndex) {
-                    /*
-                     * Because most shared state handling is deferred, we must
-                     * defer this too to ensure the editorConnector references
-                     * in shared state are up to date before opening the editor.
-                     * Yes, this is a hack on top of a hack.
-                     */
-                    Scheduler.get().scheduleDeferred(new ScheduledCommand() {
-                        @Override
-                        public void execute() {
-                            serverInitiated = true;
-                            GridConnector.this.getWidget().editRow(rowIndex);
-                        }
-                    });
+                    serverInitiated = true;
+                    GridConnector.this.getWidget().editRow(rowIndex);
                 }
 
                 @Override
@@ -225,15 +213,8 @@ public class GridConnector extends AbstractHasComponentsConnector implements
 
                 @Override
                 public void confirmBind() {
-                    /*
-                     * See comment in bind()
-                     */
-                    Scheduler.get().scheduleDeferred(new ScheduledCommand() {
-                        @Override
-                        public void execute() {
-                            endRequest();
-                        }
-                    });
+                    endRequest();
+
                 }
 
                 @Override
@@ -447,62 +428,42 @@ public class GridConnector extends AbstractHasComponentsConnector implements
             updateSelectionFromState();
         }
 
-        /*
-         * The operations in here have been made deferred.
-         * 
-         * The row data needed to react to column changes comes in the RPC
-         * calls. Since state is always updated before RPCs are called, we need
-         * to be sure that RPC is called before Grid reacts to state changes.
-         * 
-         * Note that there are still some methods annotated with @OnStateChange
-         * that aren't deferred. That's okay, though.
-         */
-
-        Scheduler.get().scheduleDeferred(new ScheduledCommand() {
-            @Override
-            public void execute() {
-                // Column updates
-                if (stateChangeEvent.hasPropertyChanged("columns")) {
-
-                    // Remove old columns
-                    purgeRemovedColumns();
-
-                    // Add new columns
-                    for (GridColumnState state : getState().columns) {
-                        if (!columnIdToColumn.containsKey(state.id)) {
-                            addColumnFromStateChangeEvent(state);
-                        }
-                        updateColumnFromState(columnIdToColumn.get(state.id),
-                                state);
-                    }
-                }
+        // Column updates
+        if (stateChangeEvent.hasPropertyChanged("columns")) {
 
-                if (stateChangeEvent.hasPropertyChanged("columnOrder")) {
-                    if (orderNeedsUpdate(getState().columnOrder)) {
-                        updateColumnOrderFromState(getState().columnOrder);
-                    }
-                }
+            // Remove old columns
+            purgeRemovedColumns();
 
-                if (stateChangeEvent.hasPropertyChanged("header")) {
-                    updateHeaderFromState(getState().header);
+            // Add new columns
+            for (GridColumnState state : getState().columns) {
+                if (!columnIdToColumn.containsKey(state.id)) {
+                    addColumnFromStateChangeEvent(state);
                 }
+                updateColumnFromState(columnIdToColumn.get(state.id), state);
+            }
+        }
 
-                if (stateChangeEvent.hasPropertyChanged("footer")) {
-                    updateFooterFromState(getState().footer);
-                }
+        if (stateChangeEvent.hasPropertyChanged("columnOrder")) {
+            if (orderNeedsUpdate(getState().columnOrder)) {
+                updateColumnOrderFromState(getState().columnOrder);
+            }
+        }
 
-                if (stateChangeEvent.hasPropertyChanged("editorEnabled")) {
-                    getWidget().setEditorEnabled(getState().editorEnabled);
-                }
+        if (stateChangeEvent.hasPropertyChanged("header")) {
+            updateHeaderFromState(getState().header);
+        }
 
-                if (stateChangeEvent.hasPropertyChanged("frozenColumnCount")) {
-                    getWidget().setFrozenColumnCount(
-                            getState().frozenColumnCount);
-                }
+        if (stateChangeEvent.hasPropertyChanged("footer")) {
+            updateFooterFromState(getState().footer);
+        }
 
-            }
-        });
+        if (stateChangeEvent.hasPropertyChanged("editorEnabled")) {
+            getWidget().setEditorEnabled(getState().editorEnabled);
+        }
 
+        if (stateChangeEvent.hasPropertyChanged("frozenColumnCount")) {
+            getWidget().setFrozenColumnCount(getState().frozenColumnCount);
+        }
     }
 
     private void updateColumnOrderFromState(List<String> stateColumnOrder) {
diff --git a/server/src/com/vaadin/data/RpcDataProviderExtension.java b/server/src/com/vaadin/data/RpcDataProviderExtension.java
index 10857f8d6a..2e10ccfef1 100644
--- a/server/src/com/vaadin/data/RpcDataProviderExtension.java
+++ b/server/src/com/vaadin/data/RpcDataProviderExtension.java
@@ -404,9 +404,9 @@ public class RpcDataProviderExtension extends AbstractExtension {
                         itemId);
                 valueChangeListeners.put(itemId, listener);
 
-                for (final Object propertyId : item.getItemPropertyIds()) {
-                    final Property<?> property = item
-                            .getItemProperty(propertyId);
+                for (final Column column : getGrid().getColumns()) {
+                    final Property<?> property = item.getItemProperty(column
+                            .getPropertyId());
                     if (property instanceof ValueChangeNotifier) {
                         ((ValueChangeNotifier) property)
                                 .addValueChangeListener(listener);
@@ -423,9 +423,9 @@ public class RpcDataProviderExtension extends AbstractExtension {
                         .remove(itemId);
 
                 if (listener != null) {
-                    for (final Object propertyId : item.getItemPropertyIds()) {
+                    for (final Column column : getGrid().getColumns()) {
                         final Property<?> property = item
-                                .getItemProperty(propertyId);
+                                .getItemProperty(column.getPropertyId());
                         if (property instanceof ValueChangeNotifier) {
                             ((ValueChangeNotifier) property)
                                     .removeValueChangeListener(listener);
@@ -436,30 +436,46 @@ public class RpcDataProviderExtension extends AbstractExtension {
         }
 
         /**
-         * Manages removed properties in active rows.
+         * Manages removed columns in active rows.
+         * <p>
+         * This method does <em>not</em> send data again to the client.
          * 
-         * @param removedPropertyIds
-         *            the property ids that have been removed from the container
+         * @param removedColumns
+         *            the columns that have been removed from the grid
          */
-        public void propertiesRemoved(
-                @SuppressWarnings("unused") Collection<Object> removedPropertyIds) {
-            /*
-             * no-op, for now.
-             * 
-             * The Container should be responsible for cleaning out any
-             * ValueChangeListeners from removed Properties. Components will
-             * benefit from this, however.
-             */
+        public void columnsRemoved(Collection<Column> removedColumns) {
+            if (removedColumns.isEmpty()) {
+                return;
+            }
+
+            for (int i = activeRange.getStart(); i < activeRange.getEnd(); i++) {
+                final Object itemId = container.getIdByIndex(i);
+                final Item item = container.getItem(itemId);
+                final GridValueChangeListener listener = valueChangeListeners
+                        .get(itemId);
+                assert (listener != null) : "a listener should've been pre-made by addValueChangeListeners";
+
+                for (final Column column : removedColumns) {
+                    final Property<?> property = item.getItemProperty(column
+                            .getPropertyId());
+                    if (property instanceof ValueChangeNotifier) {
+                        ((ValueChangeNotifier) property)
+                                .removeValueChangeListener(listener);
+                    }
+                }
+            }
         }
 
         /**
-         * Manages added properties in active rows.
+         * Manages added columns in active rows.
+         * <p>
+         * This method sends the data for the changed rows to client side.
          * 
-         * @param addedPropertyIds
-         *            the property ids that have been added to the container
+         * @param addedColumns
+         *            the columns that have been added to the grid
          */
-        public void propertiesAdded(Collection<Object> addedPropertyIds) {
-            if (addedPropertyIds.isEmpty()) {
+        public void columnsAdded(Collection<Column> addedColumns) {
+            if (addedColumns.isEmpty()) {
                 return;
             }
 
@@ -470,9 +486,9 @@ public class RpcDataProviderExtension extends AbstractExtension {
                         .get(itemId);
                 assert (listener != null) : "a listener should've been pre-made by addValueChangeListeners";
 
-                for (final Object propertyId : addedPropertyIds) {
-                    final Property<?> property = item
-                            .getItemProperty(propertyId);
+                for (final Column column : addedColumns) {
+                    final Property<?> property = item.getItemProperty(column
+                            .getPropertyId());
                     if (property instanceof ValueChangeNotifier) {
                         ((ValueChangeNotifier) property)
                                 .addValueChangeListener(listener);
@@ -924,37 +940,24 @@ public class RpcDataProviderExtension extends AbstractExtension {
     }
 
     /**
-     * Informs this data provider that some of the properties have been removed
-     * from the container.
-     * <p>
-     * Please note that we could add our own
-     * {@link com.vaadin.data.Container.PropertySetChangeListener
-     * PropertySetChangeListener} to the container, but then we'd need to
-     * implement the same bookeeping for finding what's added and removed that
-     * Grid already does in its own listener.
+     * Informs this data provider that given columns have been removed from
+     * grid.
      * 
      * @param removedColumns
-     *            a list of property ids for the removed columns
+     *            a list of removed columns
      */
-    public void propertiesRemoved(List<Object> removedColumns) {
-        activeRowHandler.propertiesRemoved(removedColumns);
+    public void columnsRemoved(List<Column> removedColumns) {
+        activeRowHandler.columnsRemoved(removedColumns);
     }
 
     /**
-     * Informs this data provider that some of the properties have been added to
-     * the container.
-     * <p>
-     * Please note that we could add our own
-     * {@link com.vaadin.data.Container.PropertySetChangeListener
-     * PropertySetChangeListener} to the container, but then we'd need to
-     * implement the same bookeeping for finding what's added and removed that
-     * Grid already does in its own listener.
+     * Informs this data provider that given columns have been added to grid.
      * 
-     * @param addedPropertyIds
-     *            a list of property ids for the added columns
+     * @param addedColumns
+     *            a list of added columns
      */
-    public void propertiesAdded(HashSet<Object> addedPropertyIds) {
-        activeRowHandler.propertiesAdded(addedPropertyIds);
+    public void columnsAdded(List<Column> addedColumns) {
+        activeRowHandler.columnsAdded(addedColumns);
     }
 
     public DataProviderKeyMapper getKeyMapper() {
diff --git a/server/src/com/vaadin/ui/Grid.java b/server/src/com/vaadin/ui/Grid.java
index 316f7682de..a4a78f4ff3 100644
--- a/server/src/com/vaadin/ui/Grid.java
+++ b/server/src/com/vaadin/ui/Grid.java
@@ -2447,28 +2447,30 @@ public class Grid extends AbstractComponent implements SelectionNotifier,
             Collection<?> properties = new HashSet<Object>(event.getContainer()
                     .getContainerPropertyIds());
 
-            // Cleanup columns that are no longer in grid
-            List<Object> removedColumns = new LinkedList<Object>();
-            for (Object columnId : columns.keySet()) {
-                if (!properties.contains(columnId)) {
-                    removedColumns.add(columnId);
+            // Find columns that need to be removed.
+            List<Column> removedColumns = new LinkedList<Column>();
+            for (Object propertyId : columns.keySet()) {
+                if (!properties.contains(propertyId)) {
+                    removedColumns.add(getColumn(propertyId));
                 }
             }
-            for (Object columnId : removedColumns) {
-                removeColumn(columnId);
-                columnKeys.remove(columnId);
+
+            // Actually remove columns.
+            for (Column column : removedColumns) {
+                Object propertyId = column.getPropertyId();
+                internalRemoveColumn(propertyId);
+                columnKeys.remove(propertyId);
             }
-            datasourceExtension.propertiesRemoved(removedColumns);
+            datasourceExtension.columnsRemoved(removedColumns);
 
             // Add new columns
-            HashSet<Object> addedPropertyIds = new HashSet<Object>();
+            List<Column> addedColumns = new LinkedList<Column>();
             for (Object propertyId : properties) {
                 if (!columns.containsKey(propertyId)) {
-                    appendColumn(propertyId);
-                    addedPropertyIds.add(propertyId);
+                    addedColumns.add(appendColumn(propertyId));
                 }
             }
-            datasourceExtension.propertiesAdded(addedPropertyIds);
+            datasourceExtension.columnsAdded(addedColumns);
 
             if (getFrozenColumnCount() > columns.size()) {
                 setFrozenColumnCount(columns.size());
@@ -2950,7 +2952,14 @@ public class Grid extends AbstractComponent implements SelectionNotifier,
         } else {
             addColumnProperty(propertyId, String.class, "");
         }
-        return getColumn(propertyId);
+
+        // Inform the data provider of this new column.
+        Column column = getColumn(propertyId);
+        List<Column> addedColumns = new ArrayList<Column>();
+        addedColumns.add(column);
+        datasourceExtension.columnsAdded(addedColumns);
+
+        return column;
     }
 
     /**
@@ -3012,10 +3021,12 @@ public class Grid extends AbstractComponent implements SelectionNotifier,
      * Removes all columns from this Grid.
      */
     public void removeAllColumns() {
+        List<Column> removed = new ArrayList<Column>(columns.values());
         Set<Object> properties = new HashSet<Object>(columns.keySet());
         for (Object propertyId : properties) {
             removeColumn(propertyId);
         }
+        datasourceExtension.columnsRemoved(removed);
     }
 
     /**
@@ -3093,6 +3104,13 @@ public class Grid extends AbstractComponent implements SelectionNotifier,
      *            The property id of column to be removed
      */
     public void removeColumn(Object propertyId) {
+        List<Column> removed = new ArrayList<Column>();
+        removed.add(getColumn(propertyId));
+        internalRemoveColumn(propertyId);
+        datasourceExtension.columnsRemoved(removed);
+    }
+
+    private void internalRemoveColumn(Object propertyId) {
         setEditorField(propertyId, null);
         header.removeColumn(propertyId);
         footer.removeColumn(propertyId);