From 944dac63adbaef3014c698c9c3181e5a7cf8b874 Mon Sep 17 00:00:00 2001
From: Larry Booker <lbooker42@gmail.com>
Date: Fri, 13 Sep 2024 09:50:46 -0700
Subject: [PATCH] Working through problems.

---
 .../deephaven/engine/table/impl/SourceTable.java |  1 +
 .../impl/locations/TrackedTableLocationKey.java  | 13 +++++++++++++
 .../impl/AbstractTableLocationProvider.java      | 16 ++++++++++------
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/SourceTable.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/SourceTable.java
index a6c652a8510..b41ccaeba49 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/SourceTable.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/SourceTable.java
@@ -139,6 +139,7 @@ private void initializeAvailableLocations() {
                 if (isRefreshing()) {
                     final TableLocationSubscriptionBuffer locationBuffer =
                             new TableLocationSubscriptionBuffer(locationProvider);
+                    manage(locationBuffer);
                     try (final TableLocationSubscriptionBuffer.LocationUpdate locationUpdate =
                             locationBuffer.processPending()) {
                         maybeRemoveLocations(locationUpdate.getPendingRemovedLocationKeys());
diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/TrackedTableLocationKey.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/TrackedTableLocationKey.java
index 81b23b06dbd..7515e8f43c5 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/TrackedTableLocationKey.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/TrackedTableLocationKey.java
@@ -36,6 +36,19 @@ public ImmutableTableLocationKey getKey() {
         return locationKey;
     }
 
+    @Override
+    public int hashCode() {
+        return locationKey.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj instanceof ImmutableTableLocationKey) {
+            return locationKey.equals(obj);
+        }
+        return (this == obj);
+    }
+
     /**
      * This {@link TrackedTableLocationKey} should manage the given {@link TableLocation} and store a reference to it.
      * 
diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/impl/AbstractTableLocationProvider.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/impl/AbstractTableLocationProvider.java
index 233ce1e0bf7..45b337e136c 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/impl/AbstractTableLocationProvider.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/locations/impl/AbstractTableLocationProvider.java
@@ -4,6 +4,7 @@
 package io.deephaven.engine.table.impl.locations.impl;
 
 import io.deephaven.base.verify.Assert;
+import io.deephaven.engine.liveness.LivenessScopeStack;
 import io.deephaven.engine.liveness.ReferenceCountedLivenessNode;
 import io.deephaven.engine.table.Table;
 import io.deephaven.engine.table.impl.locations.*;
@@ -124,6 +125,7 @@ protected void destroy() {
                 AbstractTableLocationProvider.this.destroy();
             }
         };
+        LivenessScopeStack.peek().manage(livenessNode);
     }
 
     /**
@@ -431,15 +433,14 @@ public TableLocation getTableLocationIfPresent(@NotNull final TableLocationKey t
         // The intent is to create a TableLocation exactly once to replace the TableLocationKey placeholder that was
         // added in handleTableLocationKey.
         if (!(current instanceof TableLocation)) {
-            final TableLocationKey immutableKey = (TableLocationKey) current;
+            final TrackedTableLocationKey trackedKey = (TrackedTableLocationKey) current;
             // noinspection SynchronizationOnLocalVariableOrMethodParameter
-            synchronized (immutableKey) {
-                current = tableLocations.get(immutableKey);
-                if (immutableKey == current) {
+            synchronized (trackedKey) {
+                current = tableLocations.get(trackedKey);
+                if (trackedKey == current) {
                     // Make a new location, have the tracked key manage it, then replace the key with the
                     // new location in the map and return it. Note, this may contend for the lock on tableLocations
-                    final TrackedTableLocationKey trackedKey = (TrackedTableLocationKey) current;
-                    final TableLocation newLocation = makeTableLocation(immutableKey);
+                    final TableLocation newLocation = makeTableLocation(trackedKey.getKey());
                     trackedKey.manageTableLocation(newLocation);
                     tableLocations.add(current = newLocation);
                 }
@@ -522,6 +523,9 @@ private static TableLocationKey toKey(@NotNull final Object keyOrLocation) {
         if (keyOrLocation instanceof TableLocation) {
             return ((TableLocation) keyOrLocation).getKey();
         }
+        if (keyOrLocation instanceof TrackedTableLocationKey) {
+            return (((TrackedTableLocationKey) keyOrLocation).getKey());
+        }
         if (keyOrLocation instanceof TableLocationKey) {
             return ((TableLocationKey) keyOrLocation);
         }