ui: refine hook to schedule data refreshes

This commit refines the logic in the `useFetchDataWithPolling`
hook. Namely, it has been renamed to `useScheduleFunction` to
be more generic. It also fixes a bug in the insights pages
where we would not update stale data  if `shouldPoll` was false
(which is the case for custom time intervals). We fix this by
always fetching immediately if the data is invalid.

Tests were also added to verify the hook behaviour.

Epic: none

Release note: None

pkg/ui/workspaces/cluster-ui/src/insights/workloadInsights/statementInsights/statementInsightsView.tsx

@@ -56,7 +56,7 @@ import moment from "moment";
 import styles from "src/statementsPage/statementsPage.module.scss";
 import sortableTableStyles from "src/sortedtable/sortedtable.module.scss";
 import { commonStyles } from "../../../common";
-import { useFetchDataWithPolling } from "src/util/hooks";
+import { useScheduleFunction } from "src/util/hooks";
 import { InlineAlert } from "@cockroachlabs/ui-components";
 import { insights } from "src/util";
 import { Anchor } from "src/anchor";
@@ -129,14 +129,17 @@ export const StatementInsightsView: React.FC<StatementInsightsViewProps> = ({
   }, [refreshStatementInsights, timeScale]);
 
   const shouldPoll = timeScale.key !== "Custom";
-  const clearPolling = useFetchDataWithPolling(
+  const [refetch, clearPolling] = useScheduleFunction(
     refresh,
-    isDataValid,
-    lastUpdated,
-    shouldPoll,
+    shouldPoll, // Don't reschedule refresh if we have a custom time interval.
     10 * 1000, // 10s polling interval
+    lastUpdated,
   );
 
+  useEffect(() => {
+    if (!isDataValid) refetch();
+  }, [isDataValid, refetch]);
+
   useEffect(() => {
     // We use this effect to sync settings defined on the URL (sort, filters),
     // with the redux store. The only time we do this is when the user navigates

pkg/ui/workspaces/cluster-ui/src/insights/workloadInsights/transactionInsights/transactionInsightsView.tsx

@@ -51,7 +51,7 @@ import { TxnInsightsRequest } from "src/api";
 import styles from "src/statementsPage/statementsPage.module.scss";
 import sortableTableStyles from "src/sortedtable/sortedtable.module.scss";
 import { commonStyles } from "../../../common";
-import { useFetchDataWithPolling } from "src/util/hooks";
+import { useScheduleFunction } from "src/util/hooks";
 import { InlineAlert } from "@cockroachlabs/ui-components";
 import { insights } from "src/util";
 import { Anchor } from "src/anchor";
@@ -122,14 +122,17 @@ export const TransactionInsightsView: React.FC<TransactionInsightsViewProps> = (
   }, [refreshTransactionInsights, timeScale]);
 
   const shouldPoll = timeScale.key !== "Custom";
-  const clearPolling = useFetchDataWithPolling(
+  const [refetch, clearPolling] = useScheduleFunction(
     refresh,
-    isDataValid,
-    lastUpdated,
-    shouldPoll,
+    shouldPoll, // Don't reschedule refresh if we have a custom time interval.
     10 * 1000, // 10s polling interval
+    lastUpdated,
   );
 
+  useEffect(() => {
+    if (!isDataValid) refetch();
+  }, [isDataValid, refetch]);
+
   useEffect(() => {
     // We use this effect to sync settings defined on the URL (sort, filters),
     // with the redux store. The only time we do this is when the user navigates

pkg/ui/workspaces/cluster-ui/src/util/hooks.spec.tsx

@@ -0,0 +1,142 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+import React from "react";
+import { useScheduleFunction } from "./hooks";
+import {
+  render,
+  cleanup,
+  screen,
+  fireEvent,
+  waitFor,
+} from "@testing-library/react";
+import moment from "moment";
+
+describe("useScheduleFunction", () => {
+  let mockFn: jest.Mock;
+
+  beforeEach(() => {
+    mockFn = jest.fn();
+  });
+
+  afterEach(() => {
+    cleanup();
+  });
+
+  it("should schedule the function according to the lastCompleted time and interval provided", async () => {
+    const timeoutMs = 1500;
+    const TestScheduleHook = () => {
+      const [lastUpdated, setLastUpdated] = React.useState<moment.Moment>(null);
+
+      const setLastUpdatedToNow = () => {
+        mockFn();
+        setLastUpdated(moment.utc());
+      };
+
+      useScheduleFunction(setLastUpdatedToNow, true, timeoutMs, lastUpdated);
+      return <div />;
+    };
+
+    const { unmount } = render(<TestScheduleHook />);
+    await waitFor(
+      () => new Promise(res => setTimeout(res, timeoutMs * 3 + 1000)), // Add 0.5s of buffer.
+      { timeout: timeoutMs * 3 + 1500 },
+    );
+    // 3 intervals and the initial call on mount, since last completed was initially null.
+    expect(mockFn).toBeCalledTimes(4);
+
+    unmount();
+    // Verify scheduling is stopped on unmount.
+    await waitFor(
+      () => new Promise(res => setTimeout(res, timeoutMs + 500)), // Add 0.5s of buffer.
+      { timeout: timeoutMs + 550 },
+    );
+    expect(mockFn).toBeCalledTimes(4);
+    // });
+  }, 20000);
+
+  it("should schedule the function immediately if returned scheduleNow is used", async () => {
+    const TestScheduleHook = () => {
+      const [lastUpdated, setLastUpdated] = React.useState<moment.Moment>(
+        moment.utc(),
+      );
+
+      const setLastUpdatedToNow = () => {
+        mockFn();
+        setLastUpdated(moment.utc());
+      };
+
+      const [scheduleNow] = useScheduleFunction(
+        setLastUpdatedToNow,
+        true,
+        3000, // Longer timeout.
+        lastUpdated,
+      );
+
+      const onClick = () => {
+        scheduleNow();
+      };
+
+      return (
+        <div>
+          <button onClick={onClick}></button>
+        </div>
+      );
+    };
+
+    render(<TestScheduleHook />);
+    const button = await screen.getByRole("button");
+
+    fireEvent.click(button);
+    await waitFor(() => expect(mockFn).toBeCalledTimes(1), { timeout: 1500 });
+
+    fireEvent.click(button);
+    await waitFor(() => expect(mockFn).toBeCalledTimes(2), { timeout: 1500 });
+
+    // Verify that the schedule is set up correctly after by waiting the interval length.
+    await waitFor(() => new Promise(res => setTimeout(res, 3100)), {
+      timeout: 4000,
+    });
+    expect(mockFn).toBeCalledTimes(3);
+  }, 10000);
+
+  it("should clear the scheduled func immediately if clear is used", async () => {
+    const TestScheduleHook = () => {
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      const [_, clear] = useScheduleFunction(mockFn, true, 1000, moment.utc());
+
+      React.useEffect(() => {
+        clear();
+      });
+
+      return <div />;
+    };
+
+    render(<TestScheduleHook />);
+    await waitFor(() => new Promise(res => setTimeout(res, 5000)), {
+      timeout: 5500,
+    });
+    expect(mockFn).toBeCalledTimes(0);
+  }, 10000);
+
+  it("should not reschedule the func if shouldReschedule=false", async () => {
+    const TestScheduleHook = () => {
+      // Since we pass in a last completed time here, we should expect no calls.
+      useScheduleFunction(mockFn, false, 100, moment.utc());
+      return <div />;
+    };
+
+    render(<TestScheduleHook />);
+    await waitFor(() => new Promise(res => setTimeout(res, 3000)), {
+      timeout: 3500,
+    });
+    expect(mockFn).toBeCalledTimes(0);
+  });
+});

pkg/ui/workspaces/cluster-ui/src/util/hooks.ts

@@ -15,70 +15,87 @@ export const usePrevious = <T>(value: T): T | undefined => {
   const ref = useRef<T>();
   useEffect(() => {
     ref.current = value;
-  });
+  }, [value]);
   return ref.current;
 };
 
-const MIN_REQUEST_DELAY_MS = 1000;
+const MIN_REQUEST_DELAY_MS = 500;
+
+type VoidFunction = () => void;
 
 /**
- * useFetchDataWithPolling allows the setup of making data requests with optional polling.
- * Data will automatically be fetched when the data is not valid or was never updated.
- * Requests will be made at most once per polling interval.
+ * useScheduleFunction allows the scheduling of a function call at an interval.
+ * It will be scheduled as follows:
+ * 1. Call immediately if
+ *        - `scheduleNow` callback is used
+ *        - last completed time is not set
+ * 2. Otherwise, reschedule the function if shouldReschedule is true  based on the
+ * last completed time and the scheduleInterval provided.
+ *
  * @param callbackFn The call back function to be called at the provided interval.
- * @param pollingIntervalMs interval in ms to fetch data
- * @param isDataValid whether the current dasta is valid, if the data is not valid we fetch immediately
- * @param shouldPoll whether we should setup polling
- * @param lastUpdated the time the data was last updated
- * @returns a function to stop polling
+ * @param scheduleIntervalMs scheduling interval in millis
+ * @param shouldReschedule whether we should continue to reschedule the function after completion
+ * @param lastCompleted the time the function was last completed
+ * @returns a tuple containing a function to schedule the function immediately (clearing the prev schedule)
+ * and a function to clear the schedule
  */
-export const useFetchDataWithPolling = (
+export const useScheduleFunction = (
   callbackFn: () => void,
-  isDataValid: boolean,
-  lastUpdated: moment.Moment,
-  shouldPoll: boolean,
-  pollingIntervalMs: number | null,
-): (() => void) => {
+  shouldReschedule: boolean,
+  scheduleIntervalMs: number | null,
+  lastCompleted: moment.Moment | null,
+): [VoidFunction, VoidFunction] => {
   const lastReqMade = useRef<moment.Moment>(null);
   const refreshDataTimeout = useRef<NodeJS.Timeout>(null);
 
-  const clearRefreshDataTimeout = useCallback(() => {
+  // useRef so we don't have to include this in our dep array.
+  const clearSchedule = useCallback(() => {
     if (refreshDataTimeout.current != null) {
       clearTimeout(refreshDataTimeout.current);
     }
   }, []);
 
-  useEffect(() => {
-    const now = moment();
-    let nextRefresh = null;
-
-    if (!isDataValid || !lastUpdated) {
-      // At most we will make all reqs managed by this hook once every MIN_REQUEST_DELAY_MS.
-      nextRefresh = lastReqMade.current
-        ? lastReqMade.current.clone().add(MIN_REQUEST_DELAY_MS, "milliseconds")
-        : now;
-    } else if (lastUpdated && pollingIntervalMs) {
-      nextRefresh = lastUpdated.clone().add(pollingIntervalMs, "milliseconds");
-    } else {
-      return;
-    }
+  const schedule = useCallback(
+    (scheduleNow = false) => {
+      const now = moment.utc();
+      let nextRefresh: moment.Moment;
+      if (scheduleNow) {
+        nextRefresh =
+          lastReqMade.current
+            ?.clone()
+            .add(MIN_REQUEST_DELAY_MS, "milliseconds") ?? now;
+      } else if (shouldReschedule && scheduleIntervalMs) {
+        nextRefresh = (lastCompleted ?? now)
+          .clone()
+          .add(scheduleIntervalMs, "milliseconds");
+      } else {
+        // Either we don't need to schedule the function again or we have
+        // invalid params to the hook.
+        return;
+      }
 
-    if (shouldPoll) {
+      const timeoutMs = Math.max(0, nextRefresh.diff(now, "millisecond"));
       refreshDataTimeout.current = setTimeout(() => {
-        lastReqMade.current = now;
+        lastReqMade.current = moment.utc();
+        // TODO (xinhaoz) if we can swap to using the fetch API more directly here
+        // we can abort the api call on refreshes.
         callbackFn();
-      }, Math.max(0, nextRefresh.diff(now, "millisecond")));
-    }
+      }, timeoutMs);
+    },
+    [shouldReschedule, scheduleIntervalMs, lastCompleted, callbackFn],
+  );
+
+  useEffect(() => {
+    if (!lastCompleted) schedule(true);
+    else schedule();
+
+    return clearSchedule;
+  }, [lastCompleted, schedule, clearSchedule]);
 
-    return clearRefreshDataTimeout;
-  }, [
-    callbackFn,
-    clearRefreshDataTimeout,
-    isDataValid,
-    lastUpdated,
-    shouldPoll,
-    pollingIntervalMs,
-  ]);
+  const scheduleNow = useCallback(() => {
+    clearSchedule();
+    schedule(true);
+  }, [schedule, clearSchedule]);
 
-  return clearRefreshDataTimeout;
+  return [scheduleNow, clearSchedule];
 };