Ensure recalculateColumnWidths works with refreshAll. (#11934) (#11959)

Column widths shouldn't be calculated between the clearing of cache and
re-populating it, but be delayed until the cache has some content again.
The calculations should only be triggered immediately if no rows are
expected.

Fixes #9996

4 files changed, 232 insertions, 16 deletions

diff --git a/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java b/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
index 93d2c9189a..ef240e8ba4 100644
--- a/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
+++ b/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
@@ -43,10 +43,13 @@ import com.vaadin.client.annotations.OnStateChange;
 import com.vaadin.client.communication.StateChangeEvent;
 import com.vaadin.client.connectors.AbstractListingConnector;
 import com.vaadin.client.connectors.grid.ColumnConnector.CustomColumn;
+import com.vaadin.client.data.AbstractRemoteDataSource;
 import com.vaadin.client.data.DataSource;
 import com.vaadin.client.ui.SimpleManagedLayout;
 import com.vaadin.client.widget.escalator.RowContainer;
 import com.vaadin.client.widget.grid.CellReference;
+import com.vaadin.client.widget.grid.DataAvailableEvent;
+import com.vaadin.client.widget.grid.DataAvailableHandler;
 import com.vaadin.client.widget.grid.EventCellReference;
 import com.vaadin.client.widget.grid.events.BodyClickHandler;
 import com.vaadin.client.widget.grid.events.BodyDoubleClickHandler;
@@ -100,6 +103,8 @@ public class GridConnector extends AbstractListingConnector
      */
     private class GridConnectorClientRpc implements GridClientRpc {
         private final Grid<JsonObject> grid;
+        private HandlerRegistration dataAvailableHandlerRegistration = null;
+        private boolean recalculateScheduled = false;
 
         private GridConnectorClientRpc(Grid<JsonObject> grid) {
             this.grid = grid;
@@ -138,7 +143,51 @@ public class GridConnector extends AbstractListingConnector
 
         @Override
         public void recalculateColumnWidths() {
-            grid.recalculateColumnWidths();
+            if (recalculateScheduled) {
+                return;
+            }
+
+            // Must be scheduled so that possible refreshAll has time to clear
+            // the cache.
+            recalculateScheduled = true;
+            Scheduler.get().scheduleFinally(() -> {
+                // If cache has been cleared, wait for data to become available.
+                // Don't trigger another attempt if there is already a handler
+                // waiting, that one will trigger the call when calculations are
+                // possible and clear out the registration afterwards.
+                if (((AbstractRemoteDataSource<JsonObject>) getDataSource())
+                        .getCachedRange().length() == 0
+                        && getDataSource().size() > 0) {
+                    if (dataAvailableHandlerRegistration == null) {
+                        dataAvailableHandlerRegistration = grid
+                                .addDataAvailableHandler(
+                                        new DataAvailableHandler() {
+
+                                            @Override
+                                            public void onDataAvailable(
+                                                    DataAvailableEvent event) {
+                                                if (event.getAvailableRows()
+                                                        .length() == 0
+                                                        && getDataSource()
+                                                                .size() > 0) {
+                                                    // Cache not populated yet,
+                                                    // wait for next call.
+                                                    return;
+                                                }
+                                                grid.recalculateColumnWidths();
+                                                if (dataAvailableHandlerRegistration != null) {
+                                                    dataAvailableHandlerRegistration
+                                                            .removeHandler();
+                                                    dataAvailableHandlerRegistration = null;
+                                                }
+                                            }
+                                        });
+                    }
+                } else if (dataAvailableHandlerRegistration == null) {
+                    grid.recalculateColumnWidths();
+                }
+                recalculateScheduled = false;
+            });
         }
     }
 
diff --git a/client/src/main/java/com/vaadin/client/widgets/Grid.java b/client/src/main/java/com/vaadin/client/widgets/Grid.java
index 4f13d0ae42..908446dfee 100755
--- a/client/src/main/java/com/vaadin/client/widgets/Grid.java
+++ b/client/src/main/java/com/vaadin/client/widgets/Grid.java
@@ -3393,7 +3393,8 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
                         Scheduler.get().scheduleDeferred(this);
                     }
                 } else if (currentDataAvailable.isEmpty()
-                        && dataSource.isWaitingForData()) {
+                        && (dataSource.isWaitingForData()
+                                || escalator.getBody().getRowCount() > 0)) {
                     Scheduler.get().scheduleDeferred(this);
                 } else {
                     calculate();
@@ -3406,20 +3407,16 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
 
         /**
          * Calculates and applies column widths, taking into account fixed
-         * widths and column expand rules
+         * widths and column expand rules.
          *
-         * @param immediately
-         *            <code>true</code> if the widths should be executed
-         *            immediately (ignoring lazy loading completely), or
-         *            <code>false</code> if the command should be run after a
-         *            while (duplicate non-immediately invocations are ignored).
          * @see Column#setWidth(double)
          * @see Column#setExpandRatio(int)
          * @see Column#setMinimumWidth(double)
          * @see Column#setMaximumWidth(double)
          */
         public void schedule() {
-            if (!isScheduled && isAttached()) {
+            if (!isScheduled && isAttached() && !(currentDataAvailable.isEmpty()
+                    && escalator.getBody().getRowCount() > 0)) {
                 isScheduled = true;
                 Scheduler.get().scheduleFinally(calculateCommand);
             }
@@ -3434,8 +3431,9 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
 
             // Make SelectAllCheckbox visible
             getSelectionColumn().ifPresent(col -> {
-                if (getDefaultHeaderRow() == null)
+                if (getDefaultHeaderRow() == null) {
                     return;
+                }
                 HeaderCell headerCell = getDefaultHeaderRow().getCell(col);
                 if (headerCell.getType().equals(GridStaticCellType.WIDGET)) {
                     // SelectAllCheckbox is present already
@@ -7248,6 +7246,8 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
         this.dataSource = dataSource;
         changeHandler = dataSource
                 .addDataChangeHandler(new DataChangeHandler() {
+                    private boolean recalculateColumnWidthsNeeded = false;
+
                     @Override
                     public void dataUpdated(int firstIndex, int numberOfItems) {
                         escalator.getBody().refreshRows(firstIndex,
@@ -7280,11 +7280,23 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
                             int numberOfItems) {
                         currentDataAvailable = Range.withLength(firstIndex,
                                 numberOfItems);
+                        if (recalculateColumnWidthsNeeded) {
+                            // Ensure that cache has actually been populated or
+                            // all rows removed, otherwise wait for next call.
+                            if (numberOfItems > 0
+                                    || getDataSource().size() == 0) {
+                                recalculateColumnWidths();
+                                recalculateColumnWidthsNeeded = false;
+                            }
+                        }
                         fireEvent(new DataAvailableEvent(currentDataAvailable));
                     }
 
                     @Override
                     public void resetDataAndSize(int newSize) {
+                        // It might take a while for new data to arrive,
+                        // clear the record of cached rows.
+                        currentDataAvailable = Range.emptyRange();
                         RowContainer body = escalator.getBody();
                         int oldSize = body.getRowCount();
 
@@ -7300,8 +7312,9 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
                                 // Need to recalculate column widths when the
                                 // first row is added to a non-header grid,
                                 // otherwise the checkbox will be aligned in a
-                                // wrong place.
-                                recalculateColumnWidths();
+                                // wrong place. Wait until the cache has been
+                                // populated before making the call.
+                                recalculateColumnWidthsNeeded = true;
                             }
                             body.insertRows(oldSize, newSize - oldSize);
                             cellFocusHandler.rowsAddedToBody(Range
@@ -7320,8 +7333,7 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
                                     visibleRowRange.length());
                         } else {
                             // We won't expect any data more data updates, so
-                            // just make
-                            // the bookkeeping happy
+                            // just make the bookkeeping happy.
                             dataAvailable(0, 0);
                         }
 
@@ -9307,11 +9319,11 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
             }, 50);
         }
     }
-    
+
     private void doRefreshOnResize() {
         if (escalator
                 .getInnerWidth() != autoColumnWidthsRecalculator.lastCalculatedInnerWidth) {
-           recalculateColumnWidths();
+            recalculateColumnWidths();
         }
 
         // Vertical resizing could make editor positioning invalid so it
diff --git a/uitest/src/main/java/com/vaadin/tests/components/grid/GridRecalculateColumnWidthNewItem.java b/uitest/src/main/java/com/vaadin/tests/components/grid/GridRecalculateColumnWidthNewItem.java
new file mode 100644
index 0000000000..8b4fc044c5
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/grid/GridRecalculateColumnWidthNewItem.java
@@ -0,0 +1,64 @@
+package com.vaadin.tests.components.grid;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.vaadin.data.provider.ListDataProvider;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.CheckBox;
+import com.vaadin.ui.Grid;
+
+public class GridRecalculateColumnWidthNewItem extends AbstractTestUI {
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        List<String> testItems = new ArrayList<>();
+        testItems.add("short1");
+        testItems.add("short2");
+
+        Grid<String> grid = new Grid<>();
+        grid.addColumn(String::toString).setCaption("Name");
+        grid.addColumn(item -> "col2").setCaption("Col 2");
+        grid.addColumn(item -> "col3").setCaption("Col 3");
+        grid.setDataProvider(new ListDataProvider<>(testItems));
+
+        final CheckBox recalculateCheckBox = new CheckBox(
+                "Recalculate column widths", true);
+
+        Button addButton = new Button("add row", e -> {
+            testItems.add(
+                    "Wiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiide");
+            grid.getDataProvider().refreshAll();
+            if (recalculateCheckBox.getValue()) {
+                grid.recalculateColumnWidths();
+            }
+        });
+        addButton.setId("add");
+
+        Button removeButton = new Button("remove row", e -> {
+            if (testItems.size() > 0) {
+                testItems.remove(testItems.size() - 1);
+            }
+            grid.getDataProvider().refreshAll();
+            if (recalculateCheckBox.getValue()) {
+                grid.recalculateColumnWidths();
+            }
+        });
+        removeButton.setId("remove");
+
+        addComponents(grid, addButton, removeButton, recalculateCheckBox);
+    }
+
+    @Override
+    protected String getTestDescription() {
+        return "Adding or removing a row with wider contents should update "
+                + "column widths if requested but not otherwise.";
+    }
+
+    @Override
+    protected Integer getTicketNumber() {
+        return 9996;
+    }
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/grid/GridRecalculateColumnWidthNewItemTest.java b/uitest/src/test/java/com/vaadin/tests/components/grid/GridRecalculateColumnWidthNewItemTest.java
new file mode 100644
index 0000000000..6d910208b7
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/grid/GridRecalculateColumnWidthNewItemTest.java
@@ -0,0 +1,91 @@
+package com.vaadin.tests.components.grid;
+
+import static org.hamcrest.core.IsNot.not;
+import static org.hamcrest.number.IsCloseTo.closeTo;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
+
+import java.io.IOException;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.CheckBoxElement;
+import com.vaadin.testbench.elements.GridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class GridRecalculateColumnWidthNewItemTest extends SingleBrowserTest {
+
+    GridElement grid;
+    ButtonElement addButton;
+    ButtonElement removeButton;
+
+    @Before
+    public void init() {
+        openTestURL();
+        grid = $(GridElement.class).first();
+        addButton = $(ButtonElement.class).id("add");
+        removeButton = $(ButtonElement.class).id("remove");
+    }
+
+    @Test
+    public void recalculateAfterAddingAndRemovingWorks() throws IOException {
+        assertEquals("CheckBox should be checked.", "checked",
+                $(CheckBoxElement.class).first().getValue());
+
+        int initialWidth = grid.getHeaderCell(0, 0).getSize().width;
+
+        addButton.click();
+        int newWidth = grid.getHeaderCell(0, 0).getSize().width;
+        // ensure the column width has increased significantly
+        assertThat(
+                "Unexpected column width after adding a row and calling recalculate.",
+                (double) newWidth, not(closeTo(initialWidth, 20)));
+
+        removeButton.click();
+        newWidth = grid.getHeaderCell(0, 0).getSize().width;
+        // ensure the column width has decreased significantly (even if it might
+        // not be exactly the original width)
+        assertThat(
+                "Unexpected column width after removing a row and calling recalculate.",
+                (double) newWidth, closeTo(initialWidth, 2));
+    }
+
+    @Test
+    public void addingWithoutRecalculateWorks() throws IOException {
+        CheckBoxElement checkBox = $(CheckBoxElement.class).first();
+        checkBox.click();
+        assertEquals("CheckBox should not be checked.", "unchecked",
+                checkBox.getValue());
+
+        int initialWidth = grid.getHeaderCell(0, 0).getSize().width;
+
+        addButton.click();
+        int newWidth = grid.getHeaderCell(0, 0).getSize().width;
+        // ensure the column width did not change significantly
+        assertThat(
+                "Unexpected column width after adding a row without calling recalculate.",
+                (double) newWidth, closeTo(initialWidth, 2));
+    }
+
+    @Test
+    public void removingWithoutRecalculateWorks() throws IOException {
+        // add a row before unchecking
+        addButton.click();
+
+        CheckBoxElement checkBox = $(CheckBoxElement.class).first();
+        checkBox.click();
+        assertEquals("CheckBox should not be checked.", "unchecked",
+                checkBox.getValue());
+
+        int initialWidth = grid.getHeaderCell(0, 0).getSize().width;
+
+        removeButton.click();
+        int newWidth = grid.getHeaderCell(0, 0).getSize().width;
+        // ensure the column width did not change significantly
+        assertThat(
+                "Unexpected column width after removing a row without calling recalculate.",
+                (double) newWidth, closeTo(initialWidth, 2));
+    }
+}