[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

@@ -28,7 +28,7 @@ import {
   useMutationObserver,
 } from '../../observer/mutation_observer';
 import { useResizeObserver } from '../../observer/resize_observer';
-import { DEFAULT_ROW_HEIGHT } from '../row_height_utils';
+import { DEFAULT_ROW_HEIGHT, RowHeightUtils } from '../row_height_utils';
 import { EuiDataGridCell } from './data_grid_cell';
 import {
   DataGridSortingContext,
@@ -44,10 +44,12 @@ import {
 import {
   EuiDataGridBodyProps,
   EuiDataGridInMemoryValues,
+  EuiDataGridRowHeightsOptions,
   EuiDataGridRowManager,
   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';
 
@@ -253,6 +255,67 @@ export function getParentCellContent(_element: Node | HTMLElement) {
   return element;
 }
 
+// computes the unconstrained (total possible) height of a grid
+const useUnconstrainedHeight = ({
+  rowHeightUtils,
+  startRow,
+  endRow,
+  getCorrectRowIndex,
+  rowHeightsOptions,
+  defaultHeight,
+  headerRowHeight,
+  footerRowHeight,
+}: {
+  rowHeightUtils: RowHeightUtils;
+  startRow: number;
+  endRow: number;
+  getCorrectRowIndex: (rowIndex: number) => number;
+  rowHeightsOptions?: EuiDataGridRowHeightsOptions;
+  defaultHeight: number;
+  headerRowHeight: number;
+  footerRowHeight: number;
+}) => {
+  // when a row height is updated, force a re-render of the grid body to update the unconstrained height
+  const forceRender = useForceRender();
+  useEffect(() => {
+    rowHeightUtils.setRerenderGridBody(forceRender);
+  }, [rowHeightUtils, forceRender]);
+
+  let knownHeight = 0; // tracks the pixel height of rows we know the size of
+  let knownRowCount = 0; // how many rows we know the size of
+  for (let i = startRow; i < endRow; i++) {
+    const correctRowIndex = getCorrectRowIndex(i); // map visible row to logical row
+
+    // lookup the height configuration of this row
+    const rowHeightOption = rowHeightUtils.getRowHeightOption(
+      correctRowIndex,
+      rowHeightsOptions
+    );
+
+    if (rowHeightOption) {
+      // this row's height is known
+      knownRowCount++;
+      knownHeight += rowHeightUtils.getCalculatedHeight(
+        rowHeightOption,
+        defaultHeight,
+        correctRowIndex,
+        rowHeightUtils.isRowHeightOverride(correctRowIndex, rowHeightsOptions)
+      );
+    }
+  }
+
+  // how many rows to provide space for on the screen
+  const rowCountToAffordFor = endRow - startRow;
+
+  const unconstrainedHeight =
+    defaultHeight * (rowCountToAffordFor - knownRowCount) + // guess how much space is required for unknown rows
+    knownHeight + // computed pixel height of the known rows
+    headerRowHeight + // account for header
+    footerRowHeight; // account for footer
+
+  return unconstrainedHeight;
+};
+
 export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
   props
 ) => {
@@ -573,11 +636,16 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
     }
   }, [getRowHeight]);
 
-  const rowCountToAffordFor = pagination
-    ? pagination.pageSize
-    : visibleRowIndices.length;
-  const unconstrainedHeight =
-    defaultHeight * rowCountToAffordFor + headerRowHeight + footerRowHeight;
+  const unconstrainedHeight = useUnconstrainedHeight({
+    rowHeightUtils,
+    startRow,
+    endRow,
+    getCorrectRowIndex,
+    rowHeightsOptions,
+    defaultHeight,
+    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);
+  }, []);
+};