fix: Don't memoize retrieval of tables from JsPartitionedTable#getTable

web/client-api/src/main/java/io/deephaven/web/client/api/JsPartitionedTable.java

@@ -22,7 +22,6 @@
 import io.deephaven.web.client.fu.LazyPromise;
 import io.deephaven.web.client.state.ClientTableState;
 import io.deephaven.web.shared.data.RangeSet;
-import io.deephaven.web.shared.fu.JsConsumer;
 import jsinterop.annotations.JsIgnore;
 import jsinterop.annotations.JsMethod;
 import jsinterop.annotations.JsProperty;
@@ -56,14 +55,7 @@ public class JsPartitionedTable extends HasLifecycle implements ServerObject {
     private JsTable keys;
     private TableSubscription subscription;
 
-    /*
-     * Represents the sorta-kinda memoized results, tables that we've already locally fetched from the partitioned
-     * table, and if all references to a table are released, entries here will be replaced with unresolved instances so
-     * we don't leak server references or memory. Keys are Object[], even with one element, so that we can easily hash
-     * without an extra wrapper object. Since columns are consistent with a PartitionedTable, we will not worry about
-     * "foo" vs ["foo"] as being different entries.
-     */
-    private final Map<List<Object>, JsLazy<Promise<ClientTableState>>> tables = new HashMap<>();
+    private final Set<List<Object>> knownKeys = new HashSet<>();
 
     private Column[] keyColumns;
 
@@ -157,50 +149,13 @@ private void handleKeys(Event update) {
         added.indexIterator().forEachRemaining((long index) -> {
             // extract the key to use
             JsArray<Object> key = eventData.getColumns().map((c, p1, p2) -> eventData.getData(index, c));
-            populateLazyTable(key.asList());
-            CustomEventInit init = CustomEventInit.create();
+            knownKeys.add(key.asList());
+            CustomEventInit<JsArray<Object>> init = CustomEventInit.create();
             init.setDetail(key);
             fireEvent(EVENT_KEYADDED, init);
         });
     }
 
-    private void populateLazyTable(List<Object> key) {
-        tables.put(key, JsLazy.of(() -> {
-            // If we've entered this lambda, the JsLazy is being used, so we need to go ahead and get the tablehandle
-            final ClientTableState entry = connection.newState((c, cts, metadata) -> {
-                // TODO deephaven-core#2529 parallelize this
-                connection.newTable(
-                        descriptor.getKeyColumnNamesList().asArray(new String[0]),
-                        keyColumnTypes.toArray(new String[0]),
-                        key.stream().map(item -> new Object[] {item}).toArray(Object[][]::new),
-                        null,
-                        this)
-                        .then(table -> {
-                            GetTableRequest getTableRequest = new GetTableRequest();
-                            getTableRequest.setPartitionedTable(widget.getTicket());
-                            getTableRequest.setKeyTableTicket(table.getHandle().makeTicket());
-                            getTableRequest.setResultId(cts.getHandle().makeTicket());
-                            connection.partitionedTableServiceClient().getTable(getTableRequest, connection.metadata(),
-                                    (error, success) -> {
-                                        table.close();
-                                        c.apply(error, success);
-                                    });
-                            return null;
-                        });
-            },
-                    "partitioned table key " + key);
-
-            // later, when the CTS is released, remove this "table" from the map and replace with an unresolved JsLazy
-            entry.onRunning(
-                    JsConsumer.doNothing(),
-                    JsConsumer.doNothing(),
-                    () -> populateLazyTable(key));
-
-            // we'll make a table to return later, this func here just produces the JsLazy of the CTS
-            return entry.refetch(this, connection.metadata());
-        }));
-    }
-
     /**
      * Fetch the table with the given key.
      *
@@ -212,15 +167,36 @@ public Promise<JsTable> getTable(Object key) {
         if (!JsArray.isArray(key)) {
             key = JsArray.of(key);
         }
-        List<Object> keyList = Js.<JsArray<Object>>uncheckedCast(key).asList();
-        // Every caller gets a fresh table instance, and when all are closed, the CTS will be released.
-        // See #populateLazyTable for how that is tracked.
-        final JsLazy<Promise<ClientTableState>> entry = tables.get(keyList);
-        if (entry == null) {
+        final List<Object> keyList = Js.<JsArray<Object>>uncheckedCast(key).asList();
+        if (!knownKeys.contains(keyList)) {
             // key doesn't even exist, just hand back a null table
             return Promise.resolve((JsTable) null);
         }
-        return entry.get().then(cts -> Promise.resolve(new JsTable(cts.getConnection(), cts)));
+        final ClientTableState entry = connection.newState((c, cts, metadata) -> {
+            // TODO deephaven-core#2529 parallelize this
+            connection.newTable(
+                    descriptor.getKeyColumnNamesList().asArray(new String[0]),
+                    keyColumnTypes.toArray(new String[0]),
+                    keyList.stream().map(item -> new Object[] {item}).toArray(Object[][]::new),
+                    null,
+                    this)
+                    .then(table -> {
+                        GetTableRequest getTableRequest = new GetTableRequest();
+                        getTableRequest.setPartitionedTable(widget.getTicket());
+                        getTableRequest.setKeyTableTicket(table.getHandle().makeTicket());
+                        getTableRequest.setResultId(cts.getHandle().makeTicket());
+                        connection.partitionedTableServiceClient().getTable(getTableRequest, connection.metadata(),
+                                (error, success) -> {
+                                    table.close();
+                                    c.apply(error, success);
+                                });
+                        return null;
+                    });
+        },
+                "partitioned table key " + key);
+
+        return entry.refetch(this, connection.metadata())
+                .then(cts -> Promise.resolve(new JsTable(cts.getConnection(), cts)));
     }
 
     /**
@@ -248,9 +224,9 @@ public Promise<JsTable> getMergedTable() {
      */
     public JsSet<Object> getKeys() {
         if (subscription.getColumns().length == 1) {
-            return new JsSet<>(tables.keySet().stream().map(list -> list.get(0)).toArray());
+            return new JsSet<>(knownKeys.stream().map(list -> list.get(0)).toArray());
         }
-        return new JsSet<>(tables.keySet().stream().map(List::toArray).toArray());
+        return new JsSet<>(knownKeys.stream().map(List::toArray).toArray());
     }
 
     /**
@@ -260,7 +236,7 @@ public JsSet<Object> getKeys() {
      */
     @JsProperty(name = "size")
     public int size() {
-        return tables.size();
+        return knownKeys.size();
     }
 
     /**