[EuiDataGrid] improve height calculation

src/components/datagrid/__mocks__/row_height_utils.ts

@@ -27,6 +27,7 @@ export const mockRowHeightUtils = ({
   getLineCount: jest.fn(actual.getLineCount),
   calculateHeightForLineCount: jest.fn(() => 50),
   isRowHeightOverride: jest.fn(actual.isRowHeightOverride),
+  setRerenderGridBody: jest.fn(),
 } as unknown) as ActualRowHeightUtils;
 
 export const RowHeightUtils = jest.fn(() => mockRowHeightUtils);

src/components/datagrid/body/__snapshots__/data_grid_cell.test.tsx.snap

@@ -34,6 +34,7 @@ exports[`EuiDataGridCell renders 1`] = `
       "isRowHeightOverride": [MockFunction],
       "pruneHiddenColumnHeights": [MockFunction],
       "setGrid": [MockFunction],
+      "setRerenderGridBody": [MockFunction],
       "setRowHeight": [MockFunction],
     }
   }
@@ -98,6 +99,7 @@ exports[`EuiDataGridCell renders 1`] = `
               "isRowHeightOverride": [MockFunction],
               "pruneHiddenColumnHeights": [MockFunction],
               "setGrid": [MockFunction],
+              "setRerenderGridBody": [MockFunction],
               "setRowHeight": [MockFunction],
             }
           }

src/components/datagrid/body/data_grid_body.tsx

@@ -48,6 +48,7 @@ import {
   EuiDataGridSchemaDetector,
 } from '../data_grid_types';
 import { makeRowManager } from './data_grid_row_manager';
+import { useForceRender } from '../../../services/hooks/useForceRender';
 
 export const VIRTUALIZED_CONTAINER_CLASS = 'euiDataGrid__virtualized';
 
@@ -508,6 +509,11 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
     [rowHeightUtils]
   );
 
+  const forceRender = useForceRender();
+  useEffect(() => {
+    rowHeightUtils.setRerenderGridBody(forceRender);
+  }, [rowHeightUtils, forceRender]);
+
   const [minRowHeight, setRowHeight] = useState(DEFAULT_ROW_HEIGHT);
 
   const defaultHeight = useMemo(() => {
@@ -573,11 +579,33 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
     }
   }, [getRowHeight]);
 
+  let knownHeight = 0;
+  let knownRowCount = 0;
+  for (let i = startRow; i < endRow; i++) {
+    const correctRowIndex = getCorrectRowIndex(i);
+    const rowHeightOption = rowHeightUtils.getRowHeightOption(
+      correctRowIndex,
+      rowHeightsOptions
+    );
+    if (rowHeightOption) {
+      knownRowCount++;
+      knownHeight += rowHeightUtils.getCalculatedHeight(
+        rowHeightOption,
+        defaultHeight,
+        correctRowIndex,
+        rowHeightUtils.isRowHeightOverride(correctRowIndex, rowHeightsOptions)
+      );
+    }
+  }
+
   const rowCountToAffordFor = pagination
     ? pagination.pageSize
     : visibleRowIndices.length;
   const unconstrainedHeight =
-    defaultHeight * rowCountToAffordFor + headerRowHeight + footerRowHeight;
+    defaultHeight * (rowCountToAffordFor - knownRowCount) +
+    knownHeight +
+    headerRowHeight +
+    footerRowHeight;
 
   // unable to determine this until the container's size is known anyway
   const unconstrainedWidth = 0;

src/components/datagrid/data_grid.spec.tsx

@@ -51,6 +51,35 @@ describe('EuiDataGrid', () => {
         .should('have.lengthOf', 0);
     });
   });
+
+  describe('height calculation', async () => {
+    it('computes a new unconstrained height when switching to auto height', () => {
+      const renderCellValue: EuiDataGridProps['renderCellValue'] = ({
+        rowIndex,
+        columnId,
+      }) => (
+        <>
+          row {rowIndex}
+          <br />
+          column {columnId}
+        </>
+      );
+
+      mount(<EuiDataGrid {...baseProps} renderCellValue={renderCellValue} />);
+
+      getGridData();
+      cy.get('[data-test-subj=euiDataGridBody]')
+        .invoke('outerHeight')
+        .then((firstHeight) => {
+          cy.get('[data-test-subj=dataGridDisplaySelectorPopover]').click();
+          cy.get('[data-text="Auto fit"]').click();
+
+          cy.get('[data-test-subj=euiDataGridBody]')
+            .invoke('outerHeight')
+            .should('be.greaterThan', firstHeight);
+        });
+    });
+  });
 });
 
 function getGridData() {

src/components/datagrid/row_height_utils.test.ts

@@ -282,6 +282,14 @@ describe('RowHeightUtils', () => {
 
           expect(resetRowSpy).not.toHaveBeenCalled();
         });
+
+        it('calls rerenderGridBody', () => {
+          const rerenderGridBody = jest.fn();
+          rowHeightUtils.setRerenderGridBody(rerenderGridBody);
+          expect(rerenderGridBody).toHaveBeenCalledTimes(0);
+          rowHeightUtils.setRowHeight(1, 'a', 34, 1);
+          expect(rerenderGridBody).toHaveBeenCalledTimes(1);
+        });
       });
 
       describe('getRowHeight', () => {

src/components/datagrid/row_height_utils.ts

@@ -150,6 +150,7 @@ export class RowHeightUtils {
   private timerId?: number;
   private grid?: Grid;
   private lastUpdatedRow: number = Infinity;
+  private rerenderGridBody: Function = () => {};
 
   isAutoHeight(
     rowIndex: number,
@@ -192,6 +193,7 @@ export class RowHeightUtils {
     rowHeights.set(colId, adaptedHeight);
     this.heightsCache.set(rowIndex, rowHeights);
     this.resetRow(visibleRowIndex);
+    this.rerenderGridBody();
   }
 
   pruneHiddenColumnHeights(visibleColumns: EuiDataGridColumn[]) {
@@ -229,4 +231,8 @@ export class RowHeightUtils {
   setGrid(grid: Grid) {
     this.grid = grid;
   }
+
+  setRerenderGridBody(rerenderGridBody: Function) {
+    this.rerenderGridBody = rerenderGridBody;
+  }
 }

src/services/hooks/index.ts

@@ -6,8 +6,9 @@
  * Side Public License, v 1.
  */
 
-export * from './useCombinedRefs';
-export * from './useUpdateEffect';
 export * from './useDependentState';
+export * from './useCombinedRefs';
+export * from './useForceRender';
 export * from './useIsWithinBreakpoints';
 export * from './useMouseMove';
+export * from './useUpdateEffect';

src/services/hooks/useForceRender.test.tsx

@@ -0,0 +1,47 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import React, { useImperativeHandle, createRef, forwardRef } from 'react';
+import { act } from 'react-dom/test-utils';
+import { mount } from 'enzyme';
+import { useForceRender } from './useForceRender';
+
+interface MockRefShape {
+  render: () => void;
+}
+
+describe('useForceRender', () => {
+  const renderTracker = jest.fn();
+
+  // eslint-disable-next-line local/forward-ref
+  const MockComponent = forwardRef<MockRefShape>((props, ref) => {
+    const render = useForceRender();
+
+    renderTracker();
+
+    // expose the render function on the component's ref
+    useImperativeHandle(ref, () => ({ render }), [render]);
+
+    return null;
+  });
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('causes the component to re-render', () => {
+    const ref = createRef<MockRefShape>();
+    mount(<MockComponent ref={ref} />);
+
+    expect(renderTracker).toHaveBeenCalledTimes(1);
+    act(() => {
+      ref.current!.render();
+    });
+    expect(renderTracker).toHaveBeenCalledTimes(2);
+  });
+});

src/services/hooks/useForceRender.ts

@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { useState, useCallback } from 'react';
+
+export const useForceRender = () => {
+  const [, setRenderCount] = useState(0);
+  return useCallback(() => {
+    setRenderCount((x) => x + 1);
+  }, [setRenderCount]);
+};