diff --git a/src/DataGrid.tsx b/src/DataGrid.tsx
index dc6f6aeb58..3a40ee1103 100644
--- a/src/DataGrid.tsx
+++ b/src/DataGrid.tsx
@@ -300,7 +300,6 @@ function DataGrid<R, SR, K extends Key>(
   const [isDragging, setDragging] = useState(false);
   const [draggedOverRowIdx, setOverRowIdx] = useState<number | undefined>(undefined);
   const [scrollToPosition, setScrollToPosition] = useState<PartialPosition | null>(null);
-  const [shouldFocusCell, setShouldFocusCell] = useState(false);
   const [previousRowIdx, setPreviousRowIdx] = useState(-1);
 
   const getColumnWidth = useCallback(
@@ -470,16 +469,6 @@ function DataGrid<R, SR, K extends Key>(
     latestDraggedOverRowIdx.current = rowIdx;
   }, []);
 
-  const focusCellOrCellContent = useCallback(() => {
-    const cell = getCellToScroll(gridRef.current!);
-    if (cell === null) return;
-
-    scrollIntoView(cell);
-    // Focus cell content when available instead of the cell itself
-    const elementToFocus = cell.querySelector<Element & HTMLOrSVGElement>('[tabindex="0"]') ?? cell;
-    elementToFocus.focus({ preventScroll: true });
-  }, [gridRef]);
-
   /**
    * effects
    */
@@ -494,13 +483,6 @@ function DataGrid<R, SR, K extends Key>(
     }
   }, [selectedCellIsWithinSelectionBounds, selectedPosition]);
 
-  useLayoutEffect(() => {
-    if (shouldFocusCell) {
-      setShouldFocusCell(false);
-      focusCellOrCellContent();
-    }
-  }, [shouldFocusCell, focusCellOrCellContent]);
-
   useImperativeHandle(ref, () => ({
     element: gridRef.current,
     scrollToCell({ idx, rowIdx }) {
@@ -753,6 +735,16 @@ function DataGrid<R, SR, K extends Key>(
     );
   }
 
+  function focusCellOrCellContent() {
+    const cell = getCellToScroll(gridRef.current!);
+    if (cell === null) return;
+
+    scrollIntoView(cell);
+    // Focus cell content when available instead of the cell itself
+    const elementToFocus = cell.querySelector<Element & HTMLOrSVGElement>('[tabindex="0"]') ?? cell;
+    elementToFocus.focus({ preventScroll: true });
+  }
+
   function selectCell(position: Position, enableEditor?: Maybe<boolean>): void {
     if (!isCellWithinSelectionBounds(position)) return;
     commitEditorChanges();
@@ -766,8 +758,10 @@ function DataGrid<R, SR, K extends Key>(
       // Avoid re-renders if the selected cell state is the same
       scrollIntoView(getCellToScroll(gridRef.current!));
     } else {
-      setShouldFocusCell(true);
-      setSelectedPosition({ ...position, mode: 'SELECT' });
+      flushSync(() => {
+        setSelectedPosition({ ...position, mode: 'SELECT' });
+      });
+      focusCellOrCellContent();
     }
 
     if (onSelectedCellChange && !samePosition) {
@@ -915,6 +909,10 @@ function DataGrid<R, SR, K extends Key>(
     );
   }
 
+  function cancelEditMode() {
+    setSelectedPosition(({ idx, rowIdx }) => ({ idx, rowIdx, mode: 'SELECT' }));
+  }
+
   function getCellEditor(rowIdx: number) {
     if (selectedPosition.rowIdx !== rowIdx || selectedPosition.mode === 'SELECT') return;
 
@@ -923,8 +921,12 @@ function DataGrid<R, SR, K extends Key>(
     const colSpan = getColSpan(column, lastFrozenColumnIndex, { type: 'ROW', row });
 
     const closeEditor = (shouldFocusCell: boolean) => {
-      setShouldFocusCell(shouldFocusCell);
-      setSelectedPosition(({ idx, rowIdx }) => ({ idx, rowIdx, mode: 'SELECT' }));
+      if (shouldFocusCell) {
+        flushSync(cancelEditMode);
+        focusCellOrCellContent();
+      } else {
+        cancelEditMode();
+      }
     };
 
     const onRowChange = (row: R, commitChanges: boolean, shouldFocusCell: boolean) => {
@@ -935,8 +937,11 @@ function DataGrid<R, SR, K extends Key>(
         // SELECT and this results in onRowChange getting called twice.
         flushSync(() => {
           updateRow(column, selectedPosition.rowIdx, row);
-          closeEditor(shouldFocusCell);
+          cancelEditMode();
         });
+        if (shouldFocusCell) {
+          focusCellOrCellContent();
+        }
       } else {
         setSelectedPosition((position) => ({ ...position, row }));
       }
@@ -944,7 +949,7 @@ function DataGrid<R, SR, K extends Key>(
 
     if (rows[selectedPosition.rowIdx] !== selectedPosition.originalRow) {
       // Discard changes if rows are updated from outside
-      closeEditor(false);
+      cancelEditMode();
     }
 
     return (
diff --git a/src/cellRenderers/renderCheckbox.tsx b/src/cellRenderers/renderCheckbox.tsx
index a24ac0ebe4..be2c798f62 100644
--- a/src/cellRenderers/renderCheckbox.tsx
+++ b/src/cellRenderers/renderCheckbox.tsx
@@ -22,9 +22,16 @@ const checkbox = css`
 
 const checkboxClassname = `rdg-checkbox-input ${checkbox}`;
 
-export function renderCheckbox({ onChange, indeterminate, ...props }: RenderCheckboxProps) {
+export function renderCheckbox({
+  onChange,
+  indeterminate,
+  checked,
+  ...props
+}: RenderCheckboxProps) {
   function handleChange(e: React.ChangeEvent<HTMLInputElement>) {
-    onChange(e.target.checked, (e.nativeEvent as MouseEvent).shiftKey);
+    // https://github.com/facebook/react/issues/31358
+    // onChange(e.target.checked, (e.nativeEvent as MouseEvent).shiftKey);
+    onChange(!checked, (e.nativeEvent as MouseEvent).shiftKey);
   }
 
   return (
@@ -36,6 +43,7 @@ export function renderCheckbox({ onChange, indeterminate, ...props }: RenderChec
       }}
       type="checkbox"
       className={checkboxClassname}
+      checked={checked}
       onChange={handleChange}
       {...props}
     />