diff --git a/plugins/ui/src/deephaven/ui/_internal/utils.py b/plugins/ui/src/deephaven/ui/_internal/utils.py index 066b6bbcb..534deca51 100644 --- a/plugins/ui/src/deephaven/ui/_internal/utils.py +++ b/plugins/ui/src/deephaven/ui/_internal/utils.py @@ -127,7 +127,7 @@ def _wrapped_callable( func: Callable, *args: Any, **kwargs: Any, -) -> None: +) -> Any: """ Filter the args and kwargs and call the specified function with the filtered args and kwargs. @@ -139,6 +139,9 @@ def _wrapped_callable( func: The function to call *args: args, used by the dispatcher **kwargs: kwargs, used by the dispatcher + + Returns: + The result of the function call. """ args = args if max_args is None else args[:max_args] kwargs = ( @@ -146,7 +149,7 @@ def _wrapped_callable( if kwargs_set is None else {k: v for k, v in kwargs.items() if k in kwargs_set} ) - func(*args, **kwargs) + return func(*args, **kwargs) def wrap_callable(func: Callable) -> Callable: diff --git a/plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py b/plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py index b0dfc8ca4..b622466f7 100644 --- a/plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py +++ b/plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py @@ -19,6 +19,7 @@ from .._internal import wrap_callable from ..elements import Element from ..renderer import NodeEncoder, Renderer, RenderedNode +from ..renderer.NodeEncoder import CALLABLE_KEY from .._internal import RenderContext, StateUpdateCallable, ExportedRenderState from .ErrorCode import ErrorCode @@ -98,6 +99,24 @@ class ElementMessageStream(MessageStream): Callables and render functions to be called on the next render loop. """ + _callable_dict: dict[str, Callable] + """ + Dict of callable IDs to callables. + This is intended to be used by the renderer and can be replaced on each render. + """ + + _temp_callable_dict: dict[str, Callable] + """ + Dict of callable IDs to callables returned from other callables. + These are not generated by the renderer and can be removed by the client. + This should not be cleaned out on each render like _callable_dict. + """ + + _next_temp_callable_id: int + """ + The next ID to use for temporary callables. + """ + _render_lock: threading.Lock """ Lock to ensure only one thread is rendering at a time. @@ -147,6 +166,9 @@ def __init__(self, element: Element, connection: MessageStream): self._renderer = Renderer(self._context) self._update_queue = Queue() self._callable_queue = Queue() + self._callable_dict = {} + self._temp_callable_dict = {} + self._next_temp_callable_id = 0 self._render_lock = threading.Lock() self._is_dirty = False self._render_state = _RenderState.IDLE @@ -327,6 +349,8 @@ def _make_request(self, method: str, *params: Any) -> dict[str, Any]: def _make_dispatcher(self) -> Dispatcher: dispatcher = Dispatcher() dispatcher["setState"] = self._set_state + dispatcher["callCallable"] = self._call_callable + dispatcher["closeCallable"] = self._close_callable return dispatcher def _set_state(self, state: ExportedRenderState) -> None: @@ -340,6 +364,58 @@ def _set_state(self, state: ExportedRenderState) -> None: self._context.import_state(state) self._mark_dirty() + def _call_callable(self, callable_id: str, args: Any) -> Any: + """ + Call a callable by its ID. + If the result is a callable, it is registered as a temporary callable. + + Args: + callable_id: The ID of the callable to call + args: The array of arguments to pass to the callable. These will be spread as positional args to the callable. + """ + logger.debug("Calling callable %s with %s", callable_id, args) + fn = self._callable_dict.get(callable_id) or self._temp_callable_dict.get( + callable_id + ) + if fn is None: + logger.error("Callable not found: %s", callable_id) + return + result = fn(*args) + + def serialize_callables(node: Any) -> Any: + if callable(node): + new_id = f"tempCb{self._next_temp_callable_id}" + self._next_temp_callable_id += 1 + self._temp_callable_dict[new_id] = node + return { + CALLABLE_KEY: new_id, + } + raise TypeError( + f"A Deephaven UI callback returned a non-serializable value. Object of type {type(node).__name__} is not JSON serializable" + ) + + try: + return json.dumps(result, default=serialize_callables) + except Exception as e: + # This is shown to the user in the Python console + # The stack trace from logger.exception is useless to the user + # Stack trace only includes the internals of the serialization process + logger.error(e) + return { + "serialization_error": f"Cannot serialize callable {callable_id} result" + } + + def _close_callable(self, callable_id: str) -> None: + """ + Close a callable by its ID. + + Args: + callable_id: The ID of the callable to close + """ + logger.debug("Closing callable %s", callable_id) + self._callable_dict.pop(callable_id, None) + self._temp_callable_dict.pop(callable_id, None) + def _send_document_update( self, root: RenderedNode, state: ExportedRenderState ) -> None: @@ -366,11 +442,11 @@ def _send_document_update( payload = json.dumps(request) logger.debug(f"Sending payload: {payload}") - dispatcher = self._make_dispatcher() + callable_dict = {} for callable, callable_id in callable_id_dict.items(): logger.debug("Registering callable %s", callable_id) - dispatcher[callable_id] = wrap_callable(callable) - self._dispatcher = dispatcher + callable_dict[callable_id] = wrap_callable(callable) + self._callable_dict = callable_dict if self._is_closed: # The connection is closed, so this component will not update anymore # delete the context so the objects in the collected scope are released diff --git a/plugins/ui/src/js/src/widget/WidgetHandler.tsx b/plugins/ui/src/js/src/widget/WidgetHandler.tsx index 8b2c3b02b..43f665322 100644 --- a/plugins/ui/src/js/src/widget/WidgetHandler.tsx +++ b/plugins/ui/src/js/src/widget/WidgetHandler.tsx @@ -18,7 +18,7 @@ import { WidgetDescriptor } from '@deephaven/dashboard'; import { useWidget } from '@deephaven/jsapi-bootstrap'; import type { dh } from '@deephaven/jsapi-types'; import Log from '@deephaven/log'; -import { EMPTY_FUNCTION } from '@deephaven/utils'; +import { EMPTY_FUNCTION, assertNotNull } from '@deephaven/utils'; import { CALLABLE_KEY, OBJECT_KEY, @@ -35,7 +35,7 @@ import { METHOD_DOCUMENT_UPDATED, } from './WidgetTypes'; import DocumentHandler from './DocumentHandler'; -import { getComponentForElement } from './WidgetUtils'; +import { getComponentForElement, wrapCallable } from './WidgetUtils'; import WidgetErrorView from './WidgetErrorView'; import ReactPanelContentOverlayContext from '../layout/ReactPanelContentOverlayContext'; @@ -116,6 +116,15 @@ function WidgetHandler({ [jsonClient] ); + const callableFinalizationRegistry = useMemo( + () => + new FinalizationRegistry(callableId => { + log.debug2('Closing callable', callableId); + jsonClient?.request('closeCallable', [callableId]); + }), + [jsonClient] + ); + const parseDocument = useCallback( /** * Parse the data from the server, replacing some of the nodes on the way. @@ -127,6 +136,7 @@ function WidgetHandler({ * @returns The parsed data */ (data: string) => { + assertNotNull(jsonClient); // Keep track of exported objects that are no longer in use after this render. // We close those objects that are no longer referenced, as they will never be referenced again. const deadObjectMap = new Map(exportedObjectMap.current); @@ -136,10 +146,11 @@ function WidgetHandler({ if (isCallableNode(value)) { const callableId = value[CALLABLE_KEY]; log.debug2('Registering callableId', callableId); - return async (...args: unknown[]) => { - log.debug('Callable called', callableId, ...args); - return jsonClient?.request(callableId, args); - }; + return wrapCallable( + jsonClient, + callableId, + callableFinalizationRegistry + ); } if (isObjectNode(value)) { // Replace this node with the exported object @@ -183,7 +194,7 @@ function WidgetHandler({ ); return parsedData; }, - [jsonClient] + [jsonClient, callableFinalizationRegistry] ); const updateExportedObjects = useCallback( diff --git a/plugins/ui/src/js/src/widget/WidgetUtils.test.tsx b/plugins/ui/src/js/src/widget/WidgetUtils.test.tsx index a9e8fba29..c3d5e7894 100644 --- a/plugins/ui/src/js/src/widget/WidgetUtils.test.tsx +++ b/plugins/ui/src/js/src/widget/WidgetUtils.test.tsx @@ -1,10 +1,16 @@ import React from 'react'; +import type { JSONRPCServerAndClient } from 'json-rpc-2.0'; import { Text } from '@deephaven/components'; +import { TestUtils } from '@deephaven/utils'; import { ELEMENT_NAME, ELEMENT_PREFIX, } from '../elements/model/ElementConstants'; -import { ElementNode, ELEMENT_KEY } from '../elements/utils/ElementUtils'; +import { + ElementNode, + ELEMENT_KEY, + CALLABLE_KEY, +} from '../elements/utils/ElementUtils'; import HTMLElementView from '../elements/HTMLElementView'; import IconElementView from '../elements/IconElementView'; import { @@ -12,8 +18,23 @@ import { getComponentForElement, getComponentTypeForElement, getPreservedData, + wrapCallable, } from './WidgetUtils'; +const mockJsonRequest = jest.fn(() => + Promise.resolve(JSON.stringify({ result: 'mock' })) +); + +const mockJsonClient = TestUtils.createMockProxy({ + request: mockJsonRequest, +}); + +const mockFinilizationRegistry = TestUtils.createMockProxy< + FinalizationRegistry +>({ + register: jest.fn(), +}); + describe('getComponentTypeForElement', () => { it.each( Object.keys(elementComponentMap) as (keyof typeof elementComponentMap)[] @@ -91,3 +112,102 @@ describe('getPreservedData', () => { expect(actual).toEqual({ panelIds: widgetData.panelIds }); }); }); + +describe('wrapCallable', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('should return a function that sends a request to the client', () => { + wrapCallable(mockJsonClient, 'testMethod', mockFinilizationRegistry)(); + expect(mockJsonClient.request).toHaveBeenCalledWith('callCallable', [ + 'testMethod', + [], + ]); + }); + + it('should return a function that sends a request to the client with args', () => { + wrapCallable( + mockJsonClient, + 'testMethod', + mockFinilizationRegistry + )('a', { b: 'b' }); + expect(mockJsonClient.request).toHaveBeenCalledWith('callCallable', [ + 'testMethod', + ['a', { b: 'b' }], + ]); + }); + + it('should register the function in the finalization registry', () => { + const wrapped = wrapCallable( + mockJsonClient, + 'testMethod', + mockFinilizationRegistry + ); + + expect(mockFinilizationRegistry.register).toHaveBeenCalledWith( + wrapped, + 'testMethod', + wrapped + ); + }); + + it('should wrap returned callables', async () => { + mockJsonRequest.mockResolvedValueOnce( + JSON.stringify({ + [CALLABLE_KEY]: 'nestedCb', + }) + ); + + const wrappedResult = await wrapCallable( + mockJsonClient, + 'testMethod', + mockFinilizationRegistry + )(); + expect(wrappedResult).toBeInstanceOf(Function); + + expect(mockFinilizationRegistry.register).toHaveBeenCalledTimes(2); + expect(mockFinilizationRegistry.register).toHaveBeenLastCalledWith( + wrappedResult, + 'nestedCb', + wrappedResult + ); + }); + + it('should wrap nested returned callables', async () => { + mockJsonRequest.mockResolvedValueOnce( + JSON.stringify({ + nestedCallable: { + [CALLABLE_KEY]: 'nestedCb', + }, + someOtherProp: 'mock', + }) + ); + + const wrappedResult = (await wrapCallable( + mockJsonClient, + 'testMethod', + mockFinilizationRegistry + )()) as { nestedCallable: () => void; someOtherProp: string }; + + expect(wrappedResult).toMatchObject({ + nestedCallable: expect.any(Function), + someOtherProp: 'mock', + }); + + expect(mockFinilizationRegistry.register).toHaveBeenCalledTimes(2); + expect(mockFinilizationRegistry.register).toHaveBeenLastCalledWith( + wrappedResult.nestedCallable, + 'nestedCb', + wrappedResult.nestedCallable + ); + }); + + it('should reject if the result is not parseable', () => { + mockJsonRequest.mockResolvedValueOnce('not a json string'); + + expect( + wrapCallable(mockJsonClient, 'testMethod', mockFinilizationRegistry)() + ).rejects.toBeInstanceOf(Error); + }); +}); diff --git a/plugins/ui/src/js/src/widget/WidgetUtils.tsx b/plugins/ui/src/js/src/widget/WidgetUtils.tsx index 367752436..ebec26892 100644 --- a/plugins/ui/src/js/src/widget/WidgetUtils.tsx +++ b/plugins/ui/src/js/src/widget/WidgetUtils.tsx @@ -1,6 +1,7 @@ /* eslint-disable react/jsx-props-no-spreading */ /* eslint-disable import/prefer-default-export */ import React, { ComponentType } from 'react'; +import type { JSONRPCServerAndClient } from 'json-rpc-2.0'; // Importing `Item` and `Section` compnents directly since they should not be // wrapped due to how Spectrum collection components consume them. import { @@ -26,12 +27,15 @@ import { View, } from '@deephaven/components'; import { ValueOf } from '@deephaven/utils'; +import Log from '@deephaven/log'; import { ReadonlyWidgetData, WidgetAction, isWidgetError } from './WidgetTypes'; import { ElementNode, ELEMENT_KEY, isElementNode, wrapElementChildren, + isCallableNode, + CALLABLE_KEY, wrapTextChildren, } from '../elements/utils/ElementUtils'; import HTMLElementView from '../elements/HTMLElementView'; @@ -70,6 +74,8 @@ const shouldWrapTextChildren = new Set([ ELEMENT_NAME.view, ]); +const log = Log.module('@deephaven/js-plugin-ui/WidgetUtils'); + /* * Map element node names to their corresponding React components */ @@ -177,6 +183,56 @@ export function getPreservedData( ); } +/** + * Wraps a callable returned by the server so any returned callables are also wrapped. + * The callable will also be added to the finalization registry so it can be cleaned up + * when there are no more strong references to the callable. + * @param jsonClient The JSON client to send callable requests to + * @param callableId The callableId to return a wrapped callable for + * @param registry The finalization registry to register the callable with. + * @returns A wrapped callable that will automatically wrap any nested callables returned by the server + */ +export function wrapCallable( + jsonClient: JSONRPCServerAndClient, + callableId: string, + registry: FinalizationRegistry +): (...args: unknown[]) => Promise { + const callable = async (...args: unknown[]) => { + log.debug2(`Callable ${callableId} called`, args); + const resultString = await jsonClient.request('callCallable', [ + callableId, + args, + ]); + + log.debug2(`Callable ${callableId} result string`, resultString); + + try { + // Do NOT add anything that logs result + // It creates a strong ref to the result object in the console logs + // As a result, any returned callables will never be GC'd and the finalization registry will never clean them up + const result = JSON.parse(resultString, (key, value) => { + if (isCallableNode(value)) { + const nestedCallable = wrapCallable( + jsonClient, + value[CALLABLE_KEY], + registry + ); + return nestedCallable; + } + return value; + }); + + return result; + } catch { + throw new Error(`Error parsing callable result: ${resultString}`); + } + }; + + registry.register(callable, callableId, callable); + + return callable; +} + /** * Get the name of an error type * @param error Name of an error