fix: Heap usage request throttling

packages/console/package.json

@@ -31,6 +31,7 @@
     "@deephaven/log": "file:../log",
     "@deephaven/storage": "file:../storage",
     "@deephaven/utils": "file:../utils",
+    "@deephaven/react-hooks": "file:../react-hooks",
     "@fortawesome/react-fontawesome": "^0.2.0",
     "classnames": "^2.3.1",
     "linkifyjs": "^4.1.0",

packages/console/src/HeapUsage.tsx

@@ -1,9 +1,10 @@
-import React, { useEffect, useState, ReactElement, useRef } from 'react';
+import { useState, ReactElement, useRef, useCallback } from 'react';
 import classNames from 'classnames';
 import { Tooltip } from '@deephaven/components';
 import type { QueryConnectable } from '@deephaven/jsapi-types';
 import { Plot, ChartTheme } from '@deephaven/chart';
 import Log from '@deephaven/log';
+import { useAsyncInterval } from '@deephaven/react-hooks';
 import './HeapUsage.scss';
 
 const log = Log.module('HeapUsage');
@@ -38,55 +39,39 @@ function HeapUsage({
     usages: [],
   });
 
-  useEffect(
-    function setUsageUpdateInterval() {
-      const fetchAndUpdate = async () => {
-        try {
-          const newUsage = await connection.getWorkerHeapInfo();
-          setMemoryUsage(newUsage);
-
-          if (bgMonitoring || isOpen) {
-            const currentUsage =
-              (newUsage.totalHeapSize - newUsage.freeMemory) /
-              newUsage.maximumHeapSize;
-            const currentTime = Date.now();
-
-            const { timestamps, usages } = historyUsage.current;
-            while (
-              timestamps.length !== 0 &&
-              currentTime - timestamps[0] > monitorDuration * 1.5
-            ) {
-              timestamps.shift();
-              usages.shift();
-            }
-
-            timestamps.push(currentTime);
-            usages.push(currentUsage);
-          } else {
-            historyUsage.current = { timestamps: [], usages: [] };
-          }
-        } catch (e) {
-          log.warn('Unable to get heap usage', e);
+  const setUsageUpdateInterval = useCallback(async () => {
+    try {
+      const newUsage = await connection.getWorkerHeapInfo();
+      setMemoryUsage(newUsage);
+
+      if (bgMonitoring || isOpen) {
+        const currentUsage =
+          (newUsage.totalHeapSize - newUsage.freeMemory) /
+          newUsage.maximumHeapSize;
+        const currentTime = Date.now();
+
+        const { timestamps, usages } = historyUsage.current;
+        while (
+          timestamps.length !== 0 &&
+          currentTime - timestamps[0] > monitorDuration * 1.5
+        ) {
+          timestamps.shift();
+          usages.shift();
         }
-      };
-      fetchAndUpdate();
-
-      const updateUsage = setInterval(
-        fetchAndUpdate,
-        isOpen ? hoverUpdateInterval : defaultUpdateInterval
-      );
-      return () => {
-        clearInterval(updateUsage);
-      };
-    },
-    [
-      isOpen,
-      hoverUpdateInterval,
-      connection,
-      defaultUpdateInterval,
-      monitorDuration,
-      bgMonitoring,
-    ]
+
+        timestamps.push(currentTime);
+        usages.push(currentUsage);
+      } else {
+        historyUsage.current = { timestamps: [], usages: [] };
+      }
+    } catch (e) {
+      log.warn('Unable to get heap usage', e);
+    }
+  }, [isOpen, connection, monitorDuration, bgMonitoring]);
+
+  useAsyncInterval(
+    setUsageUpdateInterval,
+    isOpen ? hoverUpdateInterval : defaultUpdateInterval
   );
 
   const toDecimalPlace = (num: number, dec: number) =>

packages/react-hooks/src/index.ts

@@ -1,3 +1,4 @@
+export * from './useAsyncInterval';
 export { default as useContextOrThrow } from './useContextOrThrow';
 export { default as usePrevious } from './usePrevious';
 export { default as useForwardedRef } from './useForwardedRef';

packages/react-hooks/src/useAsyncInterval.test.ts

@@ -0,0 +1,135 @@
+import { renderHook, act } from '@testing-library/react-hooks';
+import { TestUtils } from '@deephaven/utils';
+import useAsyncInterval from './useAsyncInterval';
+
+beforeEach(() => {
+  jest.clearAllMocks();
+  expect.hasAssertions();
+  jest.useFakeTimers();
+  jest.spyOn(window, 'setTimeout').mockName('setTimeout');
+});
+
+afterAll(() => {
+  jest.useRealTimers();
+});
+
+describe('useAsyncInterval', () => {
+  function createCallback(ms: number) {
+    return jest.fn(
+      async (): Promise<void> =>
+        new Promise(resolve => {
+          setTimeout(resolve, ms);
+        })
+    );
+  }
+
+  const targetIntervalMs = 1000;
+
+  it('should call the callback function after the target interval', async () => {
+    const callback = createCallback(50);
+
+    renderHook(() => useAsyncInterval(callback, targetIntervalMs));
+
+    // First tick should be scheduled for target interval
+    expect(callback).not.toHaveBeenCalled();
+    expect(window.setTimeout).toHaveBeenCalledWith(
+      expect.any(Function),
+      targetIntervalMs
+    );
+
+    // Callback should be called after target interval
+    act(() => jest.advanceTimersByTime(targetIntervalMs));
+    expect(callback).toHaveBeenCalledTimes(1);
+  });
+
+  it('should adjust the target interval based on how long async call takes', async () => {
+    const callbackDelayMs = 50;
+    const callback = createCallback(callbackDelayMs);
+
+    renderHook(() => useAsyncInterval(callback, targetIntervalMs));
+
+    // Callback should be called after target interval
+    expect(callback).not.toHaveBeenCalled();
+    act(() => jest.advanceTimersByTime(targetIntervalMs));
+    expect(callback).toHaveBeenCalledTimes(1);
+
+    jest.clearAllMocks();
+
+    // Mimick the callback Promise resolving
+    act(() => jest.advanceTimersByTime(callbackDelayMs));
+    await TestUtils.flushPromises();
+
+    // Next target interval should be adjusted based on how long the callback took
+    const nextTargetIntervalMs = targetIntervalMs - callbackDelayMs;
+
+    expect(callback).not.toHaveBeenCalled();
+    expect(window.setTimeout).toHaveBeenCalledTimes(1);
+    expect(window.setTimeout).toHaveBeenCalledWith(
+      expect.any(Function),
+      nextTargetIntervalMs
+    );
+
+    act(() => jest.advanceTimersByTime(nextTargetIntervalMs));
+    expect(callback).toHaveBeenCalledTimes(1);
+  });
+
+  it('should schedule the next callback immediately if the callback takes longer than the target interval', async () => {
+    const callbackDelayMs = targetIntervalMs + 50;
+    const callback = createCallback(callbackDelayMs);
+
+    renderHook(() => useAsyncInterval(callback, targetIntervalMs));
+
+    // Callback should be called after target interval
+    expect(callback).not.toHaveBeenCalled();
+    act(() => jest.advanceTimersByTime(targetIntervalMs));
+    expect(callback).toHaveBeenCalledTimes(1);
+
+    jest.clearAllMocks();
+
+    // Mimick the callback Promise resolving
+    act(() => jest.advanceTimersByTime(callbackDelayMs));
+    await TestUtils.flushPromises();
+
+    expect(callback).not.toHaveBeenCalled();
+    expect(window.setTimeout).toHaveBeenCalledTimes(1);
+    expect(window.setTimeout).toHaveBeenCalledWith(expect.any(Function), 0);
+
+    act(() => jest.advanceTimersByTime(0));
+    expect(callback).toHaveBeenCalledTimes(1);
+  });
+
+  it('should stop calling the callback function after unmounting', async () => {
+    const callback = createCallback(50);
+
+    const { unmount } = renderHook(() =>
+      useAsyncInterval(callback, targetIntervalMs)
+    );
+
+    unmount();
+
+    act(() => jest.advanceTimersByTime(targetIntervalMs));
+
+    expect(callback).not.toHaveBeenCalled();
+  });
+
+  it('should not re-schedule callback if callback resolves after unmounting', async () => {
+    const callbackDelayMs = 50;
+    const callback = createCallback(callbackDelayMs);
+
+    const { unmount } = renderHook(() =>
+      useAsyncInterval(callback, targetIntervalMs)
+    );
+
+    act(() => jest.advanceTimersByTime(targetIntervalMs));
+    expect(callback).toHaveBeenCalledTimes(1);
+    jest.clearAllMocks();
+
+    unmount();
+
+    // Mimick the callback Promise resolving
+    act(() => jest.advanceTimersByTime(callbackDelayMs));
+    await TestUtils.flushPromises();
+
+    expect(window.setTimeout).not.toHaveBeenCalled();
+  });
+});

packages/react-hooks/src/useAsyncInterval.ts

@@ -0,0 +1,78 @@
+import { useCallback, useEffect, useRef } from 'react';
+import Log from '@deephaven/log';
+
+const log = Log.module('useAsyncInterval');
+
+/**
+ * Calls the given async callback at a target interval.
+ *
+ * If the callback takes less time than the target interval, the target interval
+ * for the next tick will be adjusted to target the remaining time in the current
+ * interval.
+ *
+ * e.g. If the target interval is 1000ms, and the callback takes 50ms to resolve,
+ * the next tick will be scheduled for 950ms from now via `setTimeout(callback, 950)`.
+ *
+ * If the callback takes longer than the target interval, the next tick will be
+ * scheduled immediately via `setTimeout(callback, 0)`. In such cases, the time
+ * between ticks may be > than the target interval, but this guarantees that
+ * a callback won't be scheduled until after the previous one has resolved.
+ * @param callback Callback to call at the target interval
+ * @param targetIntervalMs Target interval in milliseconds to call the callback
+ */
+export function useAsyncInterval(
+  callback: () => Promise<void>,
+  targetIntervalMs: number
+) {
+  const isCancelledRef = useRef(false);
+  const trackingRef = useRef({ count: 0, started: Date.now() });
+  const setTimeoutRef = useRef(0);
+
+  const tick = useCallback(async () => {
+    const now = Date.now();
+    let elapsedSinceLastTick = now - trackingRef.current.started;
+
+    trackingRef.current.count += 1;
+    trackingRef.current.started = now;
+
+    log.debug(
+      `tick #${trackingRef.current.count}.`,
+      elapsedSinceLastTick,
+      'ms elapsed since last tick.'
+    );
+
+    await callback();
+
+    if (isCancelledRef.current) {
+      return;
+    }
+
+    elapsedSinceLastTick += Date.now() - trackingRef.current.started;
+
+    // If elapsed time is > than the target interval, adjust the next tick interval
+    const nextTickInterval = Math.max(
+      0,
+      Math.min(
+        targetIntervalMs,
+        targetIntervalMs - (elapsedSinceLastTick - targetIntervalMs)
+      )
+    );
+
+    log.debug('adjusted minIntervalMs:', nextTickInterval);
+
+    setTimeoutRef.current = window.setTimeout(tick, nextTickInterval);
+  }, [callback, targetIntervalMs]);
+
+  useEffect(() => {
+    log.debug('Setting interval minIntervalMs:', targetIntervalMs);
+
+    setTimeoutRef.current = window.setTimeout(tick, targetIntervalMs);
+
+    return () => {
+      isCancelledRef.current = true;
+      window.clearTimeout(setTimeoutRef.current);
+    };
+  }, [targetIntervalMs, tick]);
+}
+
+export default useAsyncInterval;

packages/utils/src/TestUtils.ts

@@ -43,6 +43,13 @@ export type PickMethods<T> = {
 export type ConsoleMethodName = keyof PickMethods<Console>;
 
 class TestUtils {
+  /**
+   * jest.useFakeTimers mocks `process.nextTick` by default. Hold on to a
+   * reference to the real function so we can still use it.
+   */
+  static realNextTick =
+    typeof process !== 'undefined' ? process.nextTick : undefined;
+
   /**
    * Type assertion to "cast" a function to it's corresponding jest.Mock
    * function type. Note that this is a types only helper for type assertions.
@@ -158,6 +165,23 @@ class TestUtils {
     }
   }
 
+  /**
+   * Jest doesn't have a built in way to ensure native Promises have resolved
+   * when using fake timers. We can mimic this behavior by using `process.nextTick`.
+   * Since `process.nextTick` is mocked by default when using jest.useFakeTimers(),
+   * we use the "real" process.nextTick stored in `TestUtils.realNextTick`.
+   *
+   * NOTE: Jest can be configured to leave `process.nextTick` unmocked, but this
+   * requires devs to configure it on every test.
+   * e.g.
+   * jest.useFakeTimers({
+   *   doNotFake: ['nextTick'],
+   * });
+   */
+  static async flushPromises(): Promise<void> {
+    await new Promise(TestUtils.realNextTick ?? (() => undefined));
+  }
+
   static async rightClick(
     user: ReturnType<typeof userEvent.setup>,
     element: Element