Ensure that DataIndexes produced by a RegionedColumnSourceManager are retained by the DataIndexer

Author: rcaudy

engine/table/src/main/java/io/deephaven/engine/table/impl/dataindex/RemappedDataIndex.java

@@ -7,6 +7,7 @@
 import io.deephaven.engine.table.ColumnSource;
 import io.deephaven.engine.table.DataIndex;
 import io.deephaven.engine.table.Table;
+import io.deephaven.engine.table.impl.indexer.DataIndexer;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.*;
@@ -17,7 +18,7 @@
  * A {@link AbstractDataIndex} that remaps the key columns of another {@link AbstractDataIndex}. Used to implement
  * {@link io.deephaven.engine.table.DataIndex#remapKeyColumns(Map)}.
  */
-public class RemappedDataIndex extends AbstractDataIndex {
+public class RemappedDataIndex extends AbstractDataIndex implements DataIndexer.RetainableDataIndex {
 
     private final AbstractDataIndex sourceIndex;
     private final Map<ColumnSource<?>, ColumnSource<?>> oldToNewColumnMap;
@@ -109,4 +110,9 @@ public boolean isRefreshing() {
     public boolean isValid() {
         return sourceIndex.isValid();
     }
+
+    @Override
+    public boolean shouldRetain() {
+        return DataIndexer.RetainableDataIndex.shouldRetain(sourceIndex);
+    }
 }

engine/table/src/main/java/io/deephaven/engine/table/impl/indexer/DataIndexer.java

@@ -4,6 +4,9 @@
 package io.deephaven.engine.table.impl.indexer;
 
 import com.google.common.collect.Sets;
+import io.deephaven.base.reference.HardSimpleReference;
+import io.deephaven.base.reference.SimpleReference;
+import io.deephaven.base.reference.WeakSimpleReference;
 import io.deephaven.base.verify.Require;
 import io.deephaven.engine.liveness.LivenessScopeStack;
 import io.deephaven.engine.rowset.RowSet;
@@ -20,8 +23,6 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.lang.ref.Reference;
-import java.lang.ref.WeakReference;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.Predicate;
@@ -390,6 +391,27 @@ private static DataIndex validateAndManageCachedDataIndex(@Nullable final DataIn
         return dataIndex;
     }
 
+    /**
+     * Interface for {@link DataIndex} implementations that may opt into strong reachability within the DataIndexer's
+     * cache.
+     */
+    public interface RetainableDataIndex extends DataIndex {
+
+        /**
+         * @return Whether {@code this} should be strongly held (if {@link #addDataIndex(DataIndex) added}) to maintain
+         *         reachability
+         */
+        boolean shouldRetain();
+
+        /**
+         * @return Whether {@code dataIndex} should be strongly held (if {@link #addDataIndex(DataIndex) added}) to
+         *         maintain reachability
+         */
+        static boolean shouldRetain(@NotNull final DataIndex dataIndex) {
+            return dataIndex instanceof RetainableDataIndex && ((RetainableDataIndex) dataIndex).shouldRetain();
+        }
+    }
+
     /**
      * Node structure for our multi-level cache of indexes.
      */
@@ -399,14 +421,14 @@ private static class DataIndexCache {
         @SuppressWarnings("rawtypes")
         private static final AtomicReferenceFieldUpdater<DataIndexCache, Map> DESCENDANT_CACHES_UPDATER =
                 AtomicReferenceFieldUpdater.newUpdater(DataIndexCache.class, Map.class, "descendantCaches");
-        private static final Reference<DataIndex> MISSING_INDEX_REFERENCE = new WeakReference<>(null);
+        private static final SimpleReference<DataIndex> MISSING_INDEX_REFERENCE = new WeakSimpleReference<>(null);
 
         /** The sub-indexes below this level. */
         @SuppressWarnings("FieldMayBeFinal")
         private volatile Map<ColumnSource<?>, DataIndexCache> descendantCaches = EMPTY_DESCENDANT_CACHES;
 
         /** A reference to the index at this level. Note that there will never be an index at the "root" level. */
-        private volatile Reference<DataIndex> dataIndexReference = MISSING_INDEX_REFERENCE;
+        private volatile SimpleReference<DataIndex> dataIndexReference = MISSING_INDEX_REFERENCE;
 
         private DataIndexCache() {}
 
@@ -509,7 +531,9 @@ private boolean add(@NotNull final List<ColumnSource<?>> keyColumns, @NotNull fi
                     // noinspection SynchronizationOnLocalVariableOrMethodParameter
                     synchronized (cache) {
                         if (!isValidAndLive(cache.dataIndexReference.get())) {
-                            cache.dataIndexReference = new WeakReference<>(dataIndex);
+                            cache.dataIndexReference = RetainableDataIndex.shouldRetain(dataIndex)
+                                    ? new HardSimpleReference<>(dataIndex)
+                                    : new WeakSimpleReference<>(dataIndex);
                             return true;
                         }
                     }
@@ -544,7 +568,7 @@ private DataIndex computeIfAbsent(
                             // managed by the appropriate scope for the caller's own use. Further validation is deferred
                             // as in add.
                             dataIndex = dataIndexFactory.get();
-                            cache.dataIndexReference = new WeakReference<>(dataIndex);
+                            cache.dataIndexReference = new WeakSimpleReference<>(dataIndex);
                         }
                     }
                 }

engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/MergedDataIndex.java

@@ -19,6 +19,7 @@
 import io.deephaven.engine.table.impl.by.AggregationProcessor;
 import io.deephaven.engine.table.impl.by.AggregationRowLookup;
 import io.deephaven.engine.table.impl.dataindex.AbstractDataIndex;
+import io.deephaven.engine.table.impl.indexer.DataIndexer;
 import io.deephaven.engine.table.impl.locations.TableLocation;
 import io.deephaven.engine.table.impl.perf.QueryPerformanceRecorder;
 import io.deephaven.engine.table.impl.select.FunctionalColumn;
@@ -43,7 +44,7 @@
  *           source table".
  */
 @InternalUseOnly
-class MergedDataIndex extends AbstractDataIndex {
+class MergedDataIndex extends AbstractDataIndex implements DataIndexer.RetainableDataIndex {
 
     private static final String LOCATION_DATA_INDEX_TABLE_COLUMN_NAME = "__DataIndexTable";
 
@@ -239,4 +240,9 @@ public boolean isValid() {
         }
         return isValid = true;
     }
+
+    @Override
+    public boolean shouldRetain() {
+        return true;
+    }
 }

engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/PartitioningColumnDataIndex.java

@@ -17,6 +17,7 @@
 import io.deephaven.engine.table.impl.QueryTable;
 import io.deephaven.engine.table.impl.TableUpdateImpl;
 import io.deephaven.engine.table.impl.dataindex.AbstractDataIndex;
+import io.deephaven.engine.table.impl.indexer.DataIndexer;
 import io.deephaven.engine.table.impl.sources.ArrayBackedColumnSource;
 import io.deephaven.engine.table.impl.sources.ObjectArraySource;
 import io.deephaven.engine.table.impl.sources.ReinterpretUtils;
@@ -30,7 +31,7 @@
 /**
  * DataIndex over a partitioning column of a {@link Table} backed by a {@link RegionedColumnSourceManager}.
  */
-class PartitioningColumnDataIndex<KEY_TYPE> extends AbstractDataIndex {
+class PartitioningColumnDataIndex<KEY_TYPE> extends AbstractDataIndex implements DataIndexer.RetainableDataIndex {
 
     private static final int KEY_NOT_FOUND = (int) RowSequence.NULL_ROW_KEY;
 
@@ -318,4 +319,9 @@ public boolean isRefreshing() {
     public boolean isValid() {
         return true;
     }
+
+    @Override
+    public boolean shouldRetain() {
+        return true;
+    }
 }