Fix race condition between rendering TreeGrid and setting up headers (#9105)

* Fix race condition between rendering Grid and setting up headers
* Remove Grid default headers when initializing the widget
* Fix header clean up, add null check to select all checkbox

2 files changed, 21 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 0181980ef8..5ede13ca68 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
@@ -202,6 +202,11 @@ public class GridConnector extends AbstractListingConnector
     protected void init() {
         super.init();
 
+        // Remove default headers when initializing Grid widget
+        while (getWidget().getHeaderRowCount() > 0) {
+            getWidget().removeHeaderRow(0);
+        }
+
         registerRpc(GridClientRpc.class, new GridClientRpc() {
 
             @Override
@@ -308,24 +313,22 @@ public class GridConnector extends AbstractListingConnector
      */
     @OnStateChange("header")
     void updateHeader() {
-        Scheduler.get().scheduleFinally(() -> {
-            final Grid<JsonObject> grid = getWidget();
-            final SectionState state = getState().header;
-
-            while (grid.getHeaderRowCount() > 0) {
-                grid.removeHeaderRow(0);
-            }
+        final Grid<JsonObject> grid = getWidget();
+        final SectionState state = getState().header;
 
-            for (RowState rowState : state.rows) {
-                HeaderRow row = grid.appendHeaderRow();
+        while (grid.getHeaderRowCount() > 0) {
+            grid.removeHeaderRow(0);
+        }
 
-                if (rowState.defaultHeader) {
-                    grid.setDefaultHeaderRow(row);
-                }
+        for (RowState rowState : state.rows) {
+            HeaderRow row = grid.appendHeaderRow();
 
-                updateStaticRow(rowState, row);
+            if (rowState.defaultHeader) {
+                grid.setDefaultHeaderRow(row);
             }
-        });
+
+            updateStaticRow(rowState, row);
+        }
     }
 
     @OnStateChange("rowHeight")
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 cb50dd5977..27db38c4a1 100755
--- a/client/src/main/java/com/vaadin/client/widgets/Grid.java
+++ b/client/src/main/java/com/vaadin/client/widgets/Grid.java
@@ -3045,8 +3045,10 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
          * Sets the select all checkbox visible or hidden.
          */
         protected void doSetSelectAllCheckBoxVisible() {
-            assert selectAllCheckBox != null : "Select All Checkbox has not been created for selection column.";
-            assert selectionCell != null : "Default header cell for selection column not been set.";
+            if (selectAllCheckBox == null || selectionCell == null) {
+                // There is no default header row to display select all checkbox
+                return;
+            }
 
             if (selectAllCheckBoxVisible) {
                 selectionCell.setWidget(selectAllCheckBox);