fix: Memoize use_table_data listener (#428)

Because the listener function was not memoized in `use_table_data`, the
`use_effect` within `use_table_listener` was getting called
unnecessarily, which means in some cases table updates can be missed
during the start and stop of the listener.

Added fix and caution message in `use_table_listener` for now.

It might be worthwhile making `use_table_listener` more robust in cases
where arguments changed, but the table specifically did not. This is not
trivial though.

plugins/ui/src/deephaven/ui/hooks/use_table_data.py

@@ -140,14 +140,20 @@ def use_table_data(
     else:
         executor_name = "concurrent"
 
-    table_updated = lambda: _set_new_data(table, sentinel, set_data, set_is_sentinel)
+    # memoize table_updated (and listener) so that they don't cause a start and stop of the listener
+    table_updated = ui.use_callback(
+        lambda: _set_new_data(table, sentinel, set_data, set_is_sentinel),
+        [table, sentinel],
+    )
 
     # call table_updated in the case of new table or sentinel
     ui.use_effect(table_updated, [table, sentinel])
+    listener = ui.use_callback(
+        partial(_on_update, ctx, table_updated, executor_name),
+        [table_updated, executor_name, ctx],
+    )
 
     # call table_updated every time the table updates
-    ui.use_table_listener(
-        table, partial(_on_update, ctx, table_updated, executor_name), []
-    )
+    ui.use_table_listener(table, listener, [])
 
     return transform(data, is_sentinel)

plugins/ui/src/deephaven/ui/hooks/use_table_listener.py

@@ -74,6 +74,8 @@ def use_table_listener(
 ) -> None:
     """
     Listen to a table and call a listener when the table updates.
+    If any dependencies change, the listener will be recreated.
+    In this case, updates may be missed if the table updates while the listener is being recreated.
 
     Args:
         table: The table to listen to.

plugins/ui/test/deephaven/ui/test_hooks.py

@@ -2,25 +2,30 @@
 import unittest
 from operator import itemgetter
 from queue import Queue
+from time import sleep
 from typing import Callable
 from unittest.mock import Mock
 from .BaseTest import BaseTestCase
 
 LISTENER_TIMEOUT = 2.0
+QUEUE_TIMEOUT = 1.0
 
 
-def render_hook(fn: Callable):
+def render_hook(fn: Callable, queue=None):
     """
     Render a hook function and return the context, result, and a rerender function for updating it
 
     Args:
       fn: Callable:
         The function to render. Pass in a function with a hook call within it.
         Re-render will call the same function but with the new args passed in.
+      queue: Queue:
+        The queue to put items on. If not provided, a new queue will be created.
     """
     from deephaven.ui._internal.RenderContext import RenderContext
 
-    queue = Queue()
+    if queue is None:
+        queue = Queue()
 
     context = RenderContext(lambda x: queue.put(x), lambda x: queue.put(x))
 
@@ -42,6 +47,47 @@ def _rerender(*args, **kwargs):
     return return_dict
 
 
+class NotifyQueue(Queue):
+    """
+    A queue that notifies a function when an item is put on it
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._notify_fn = None
+
+    def put(self, item: object, block: bool = True, timeout: float = None) -> None:
+        """
+        Put an item on the queue and notify the function
+
+        Args:
+            item: The item to put on the queue
+            block: True if the call should block until the item is put on the queue
+            timeout: The time to wait for the item to be put on the queue
+
+        Returns:
+            None
+        """
+        super().put(item)
+        if self._notify_fn:
+            self._notify_fn(self)
+
+    def call_after_put(self, fn: Callable[["NotifyQueue"], None]) -> None:
+        """
+        Register a function to be called after an item is put on the queue
+
+        Args:
+            fn: The function to call after an item is put on the queue
+        """
+        self._notify_fn = fn
+
+    def unregister_notify(self) -> None:
+        """
+        Unregister the notify function
+        """
+        self._notify_fn = None
+
+
 class HooksTest(BaseTestCase):
     def test_state(self):
         from deephaven.ui.hooks import use_state
@@ -274,6 +320,34 @@ def _test_table_data(t=table):
 
         self.assertEqual(result, expected)
 
+    def verify_queue_has_size(self, queue: NotifyQueue, size: int):
+        """
+        Verify that the queue has the expected size in a multi-threaded context
+
+        Args:
+            queue: The queue to check
+            size: The expected size of the quexue
+
+        Returns:
+            None
+        """
+        event = threading.Event()
+
+        def check_size(q):
+            if q.qsize() == size:
+                event.set()
+
+        # call after each put in case the queue is not at the correct size yet
+        queue.call_after_put(check_size)
+
+        # call now in case the queue is (or was) already at the correct size
+        check_size(queue)
+
+        if not event.wait(timeout=QUEUE_TIMEOUT):
+            assert False, f"queue did not reach size {size}"
+
+        queue.unregister_notify()
+
     def test_swapping_table_data(self):
         from deephaven.ui.hooks import use_table_data
         from deephaven import new_table
@@ -292,7 +366,9 @@ def _test_table_data(t=table):
             result = use_table_data(t, sentinel="sentinel")
             return result
 
-        render_result = render_hook(_test_table_data)
+        queue = NotifyQueue()
+
+        render_result = render_hook(_test_table_data, queue)
 
         result, rerender = itemgetter("result", "rerender")(render_result)
 
@@ -313,10 +389,13 @@ def _test_table_data(t=table):
         self.verify_table_updated(table_writer, dynamic_table, (1, "Testing"))
 
         rerender(dynamic_table)
+        # the queue should have two items eventually, one for each set_state in _set_new_data in use_table_data
+        # this check is needed because the set_state calls come from the listener, which is called in a separate thread
+        # so the queue might not have the correct size immediately
+        self.verify_queue_has_size(queue, 2)
         result = rerender(dynamic_table)
 
         expected = {"Numbers": [1], "Words": ["Testing"]}
-
         self.assertEqual(result, expected)
 
     def test_column_data(self):