Improve StatusBox rendering in Grid view (#22170)

* reduce grid task instance rerenders * fix test
apache · Mar 11, 2022 · 9eb1a1c · 9eb1a1c
1 parent 7bd8b2d
commit 9eb1a1c
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 7 deletions.
diff --git a/airflow/www/package.json b/airflow/www/package.json
@@ -87,6 +87,7 @@
     "framer-motion": "^4",
     "jquery": ">=3.5.0",
     "jshint": "^2.13.4",
+    "lodash": "^4.17.21",
     "moment-timezone": "^0.5.28",
     "nvd3": "^1.8.6",
     "react": "^17.0.2",

diff --git a/airflow/www/static/js/tree/StatusBox.jsx b/airflow/www/static/js/tree/StatusBox.jsx
@@ -20,6 +20,7 @@
 /* global stateColors */
 
 import React from 'react';
+import { isEqual } from 'lodash';
 import {
   Flex,
   Box,
@@ -30,7 +31,7 @@ import { callModal } from '../dag';
 import InstanceTooltip from './InstanceTooltip';
 
 const StatusBox = ({
-  group, instance, containerRef, extraLinks = [], ...rest
+  group, instance, containerRef, extraLinks = [],
 }) => {
   const {
     executionDate, taskId, tryNumber = 0, operator, runId,
@@ -68,7 +69,6 @@ const StatusBox = ({
         zIndex={1}
         onMouseEnter={onMouseOver}
         onMouseLeave={onMouseLeave}
-        {...rest}
       >
         <Box
           width="10px"
@@ -82,4 +82,16 @@ const StatusBox = ({
   );
 };
 
-export default StatusBox;
+// The default equality function is a shallow comparison and json objects will return false
+// This custom compare function allows us to do a deeper comparison
+const compareProps = (
+  prevProps,
+  nextProps,
+) => (
+  isEqual(prevProps.group, nextProps.group)
+  && isEqual(prevProps.instance, nextProps.instance)
+  && prevProps.extraLinks === nextProps.extraLinks
+  && prevProps.containerRef === nextProps.containerRef
+);
+
+export default React.memo(StatusBox, compareProps);
diff --git a/airflow/www/static/js/tree/dagRuns/Bar.jsx b/airflow/www/static/js/tree/dagRuns/Bar.jsx
@@ -20,6 +20,7 @@
 /* global stateColors, moment */
 
 import React from 'react';
+import { isEqual } from 'lodash';
 import {
   Flex,
   Box,
@@ -115,7 +116,7 @@ const compareProps = (
   prevProps,
   nextProps,
 ) => (
-  JSON.stringify(prevProps.run) === JSON.stringify(nextProps.run)
+  isEqual(prevProps.run, nextProps.run)
   && prevProps.max === nextProps.max
   && prevProps.index === nextProps.index
   && prevProps.totalRuns === nextProps.totalRuns

diff --git a/airflow/www/static/js/tree/renderTaskRows.jsx b/airflow/www/static/js/tree/renderTaskRows.jsx
@@ -120,7 +120,7 @@ const Row = (props) => {
           backgroundColor="white"
           borderBottom={0}
         >
-          <Collapse in={isFullyOpen}>
+          <Collapse in={isFullyOpen} unmountOnExit>
             <TaskName
               onToggle={memoizedToggle}
               isGroup={isGroup}
@@ -133,7 +133,7 @@ const Row = (props) => {
         </Td>
         <Td width={0} p={0} borderBottom={0} />
         <Td p={0} align="right" _groupHover={{ backgroundColor: 'rgba(113, 128, 150, 0.1)' }} transition="background-color 0.2s" borderBottom={0}>
-          <Collapse in={isFullyOpen}>
+          <Collapse in={isFullyOpen} unmountOnExit>
             <TaskInstances dagRunIds={dagRunIds} task={task} containerRef={containerRef} />
           </Collapse>
         </Td>

diff --git a/airflow/www/static/js/tree/renderTaskRows.test.jsx b/airflow/www/static/js/tree/renderTaskRows.test.jsx
@@ -127,13 +127,15 @@ describe('Test renderTaskRows', () => {
 
     const groupName = getByText('group_1');
 
-    expect(getAllByTestId('task-instance')).toHaveLength(2);
+    expect(getAllByTestId('task-instance')).toHaveLength(1);
     expect(groupName).toBeInTheDocument();
     expect(getByTestId('closed-group')).toBeInTheDocument();
 
     fireEvent.click(groupName);
 
     expect(getByTestId('open-group')).toBeInTheDocument();
+    // task instances are only rendered when their group is expanded
+    expect(getAllByTestId('task-instance')).toHaveLength(2);
   });
 
   test('Still renders names if there are no instances', () => {