Addressed Sagar's comment on allowing user to pick which dimension co…

…mbinations to track

Signed-off-by: Peter Alfonsi <[email protected]>

plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java

@@ -158,7 +158,7 @@ private EhcacheDiskCache(Builder<K, V> builder) {
         );
         this.cache = buildCache(Duration.ofMillis(expireAfterAccess.getMillis()), builder);
         List<String> dimensionNames = Objects.requireNonNull(builder.dimensionNames, "Dimension names can't be null");
-        this.statsHolder = new StatsHolder(dimensionNames, builder.getSettings());
+        this.statsHolder = new StatsHolder(dimensionNames, builder.getSettings(), StatsHolder.TrackingMode.ALL_COMBINATIONS);
     }
 
     @SuppressWarnings({ "rawtypes" })

server/src/main/java/org/opensearch/common/cache/stats/MultiDimensionCacheStats.java

@@ -13,7 +13,9 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 
 /**
@@ -23,11 +25,6 @@
  */
 public class MultiDimensionCacheStats implements CacheStats {
 
-    /**
-     * For memory purposes, don't track stats for more than this many distinct combinations of dimension values.
-     */
-    public final static int DEFAULT_MAX_DIMENSION_VALUES = 20_000;
-
     // The value of the tier dimension for entries in this Stats object. This is handled separately for efficiency,
     // as it always has the same value for every entry in the stats object.
     // Package-private for testing.
@@ -62,28 +59,31 @@ public CacheStatsResponse getTotalStats() {
 
     /**
      * Get the stats response aggregated by dimensions. If there are no values for the specified dimensions,
-     * returns an all-zero response.
+     * returns an all-zero response. If the specified dimensions don't form a valid key, as determined by the statsHolder's
+     * tracking mode, throws an IllegalArgumentException.
      */
     @Override
     public CacheStatsResponse getStatsByDimensions(List<CacheStatsDimension> dimensions) {
-        if (!checkDimensionNames(dimensions)) {
-            throw new IllegalArgumentException("Can't get stats for unrecognized dimensions");
+        List<CacheStatsDimension> modifiedDimensions = new ArrayList<>(dimensions);
+        CacheStatsDimension tierDim = getTierDimension(dimensions);
+        if (tierDim != null) {
+            modifiedDimensions.remove(tierDim);
+        }
+
+        if (!checkDimensions(modifiedDimensions)) {
+            throw new IllegalArgumentException("Can't retrieve stats for this combination of dimensions");
         }
 
-        CacheStatsDimension tierDim = getTierDimension(dimensions);
         if (tierDim == null || tierDim.dimensionValue.equals(tierDimensionValue)) {
             // If there is no tier dimension, or if the tier dimension value matches the one for this stats object, return an aggregated
             // response over the non-tier dimensions
-            List<CacheStatsDimension> modifiedDimensions = new ArrayList<>(dimensions);
-            if (tierDim != null) {
-                modifiedDimensions.remove(tierDim);
-            }
-
             ConcurrentMap<StatsHolder.Key, CacheStatsResponse> map = statsHolder.getStatsMap();
-
             CacheStatsResponse response = new CacheStatsResponse();
 
-            if (modifiedDimensions.size() == statsHolder.getDimensionNames().size()) {
+            // In the SEPARATE_DIMENSIONS_ONLY and SPECIFIC_COMBINATIONS cases, we don't do any adding; just return directly from the map.
+            // Also do this if mode is ALL_COMBINATIONS and our dimensions have a value for every dimension name.
+            if (statsHolder.mode != StatsHolder.TrackingMode.ALL_COMBINATIONS
+                || modifiedDimensions.size() == statsHolder.getDimensionNames().size()) {
                 CacheStatsResponse resultFromMap = map.getOrDefault(new StatsHolder.Key(modifiedDimensions), new CacheStatsResponse());
                 response.add(resultFromMap); // Again return a copy
                 return response;
@@ -111,17 +111,38 @@ private CacheStatsDimension getTierDimension(List<CacheStatsDimension> dimension
         return null;
     }
 
-    private boolean checkDimensionNames(List<CacheStatsDimension> dimensions) {
-        for (CacheStatsDimension dim : dimensions) {
-            if (!(statsHolder.getDimensionNames().contains(dim.dimensionName)
-                || dim.dimensionName.equals(CacheStatsDimension.TIER_DIMENSION_NAME))) {
-                // Reject dimension names that aren't in the list and aren't the tier dimension
-                return false;
-            }
+    // Check the dimensions passed in are a valid request, according to the stats holder's tracking mode
+    private boolean checkDimensions(List<CacheStatsDimension> dimensions) {
+        switch (statsHolder.mode) {
+            case SEPARATE_DIMENSIONS_ONLY:
+                if (!(dimensions.size() == 1 && statsHolder.getDimensionNames().contains(dimensions.get(0).dimensionName))) {
+                    return false;
+                }
+                break;
+            case ALL_COMBINATIONS:
+                for (CacheStatsDimension dim : dimensions) {
+                    if (!statsHolder.getDimensionNames().contains(dim.dimensionName)) {
+                        return false;
+                    }
+                }
+                break;
+            case SPECIFIC_COMBINATIONS:
+                if (!statsHolder.getSpecificCombinations().contains(getDimensionNamesSet(dimensions))) {
+                    return false;
+                }
+                break;
         }
         return true;
     }
 
+    private Set<String> getDimensionNamesSet(List<CacheStatsDimension> dimensions) {
+        Set<String> dimSet = new HashSet<>();
+        for (CacheStatsDimension dim : dimensions) {
+            dimSet.add(dim.dimensionName);
+        }
+        return dimSet;
+    }
+
     @Override
     public long getTotalHits() {
         return statsHolder.getTotalStats().getHits();

server/src/main/java/org/opensearch/common/cache/stats/StatsHolder.java

@@ -17,6 +17,7 @@
 import org.opensearch.core.common.io.stream.Writeable;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
@@ -40,6 +41,33 @@ public class StatsHolder implements Writeable {
     // The list of permitted dimensions.
     private final List<String> dimensionNames;
 
+    /**
+     * Determines which combinations of dimension values are tracked separately by this StatsHolder. In every case,
+     * incoming keys still must have all dimension values populated.
+     */
+    public enum TrackingMode {
+        /**
+         * Tracks stats for each dimension separately. Does not support retrieving stats by combinations of dimension values,
+         * only by a single dimension value.
+         */
+        SEPARATE_DIMENSIONS_ONLY,
+        /**
+         * Tracks stats for every combination of dimension values. Can retrieve stats for any combination of dimensions,
+         * by adding together the combinations.
+         */
+        ALL_COMBINATIONS,
+        /**
+         * Tracks stats for a specified subset of combinations. Each combination is kept aggregated in memory. Only stats for
+         * the pre-specified combinations can be retrieved.
+         */
+        SPECIFIC_COMBINATIONS
+    }
+
+    // The mode for this instance.
+    public final TrackingMode mode;
+    // The specific combinations of dimension names to track, if mode is SPECIFIC_COMBINATIONS.
+    private final Set<Set<String>> specificCombinations;
+
     // A map from a set of cache stats dimensions -> stats for that combination of dimensions.
     private final ConcurrentMap<Key, CacheStatsResponse> statsMap;
 
@@ -48,11 +76,30 @@ public class StatsHolder implements Writeable {
 
     private final Logger logger = LogManager.getLogger(StatsHolder.class);
 
-    public StatsHolder(List<String> dimensionNames, Settings settings) {
+    public StatsHolder(List<String> dimensionNames, Settings settings, TrackingMode mode) {
+        assert (!mode.equals(TrackingMode.SPECIFIC_COMBINATIONS))
+            : "Must use constructor specifying specificCombinations when tracking mode is set to SPECIFIC_COMBINATIONS";
         this.dimensionNames = dimensionNames;
         this.statsMap = new ConcurrentHashMap<>();
         this.totalStats = new CacheStatsResponse();
         this.maxDimensionValues = MAX_DIMENSION_VALUES_SETTING.get(settings);
+        this.mode = mode;
+        this.specificCombinations = new HashSet<>();
+    }
+
+    public StatsHolder(List<String> dimensionNames, Settings settings, TrackingMode mode, Set<Set<String>> specificCombinations) {
+        if (!mode.equals(TrackingMode.SPECIFIC_COMBINATIONS)) {
+            logger.warn("Ignoring specific combinations; tracking mode is not set to SPECIFIC_COMBINATIONS");
+        }
+        this.dimensionNames = dimensionNames;
+        this.statsMap = new ConcurrentHashMap<>();
+        this.totalStats = new CacheStatsResponse();
+        this.maxDimensionValues = MAX_DIMENSION_VALUES_SETTING.get(settings);
+        this.mode = mode;
+        for (Set<String> combination : specificCombinations) {
+            assert combination.size() > 0 : "Must have at least one dimension name in the combination to record";
+        }
+        this.specificCombinations = specificCombinations;
     }
 
     public StatsHolder(StreamInput in) throws IOException {
@@ -64,6 +111,13 @@ public StatsHolder(StreamInput in) throws IOException {
         this.statsMap = new ConcurrentHashMap<Key, CacheStatsResponse>(readMap);
         this.totalStats = new CacheStatsResponse(in);
         this.maxDimensionValues = in.readVInt();
+        this.mode = in.readEnum(TrackingMode.class);
+        this.specificCombinations = new HashSet<>();
+        int numCombinations = in.readVInt();
+        for (int i = 0; i < numCombinations; i++) {
+            String[] names = in.readStringArray();
+            specificCombinations.add(new HashSet<>(List.of(names)));
+        }
     }
 
     public List<String> getDimensionNames() {
@@ -78,6 +132,7 @@ public CacheStatsResponse getTotalStats() {
         return totalStats;
     }
 
+    // For all these increment functions, the dimensions list comes from the key, and contains all dimensions present in dimensionNames.
     public void incrementHitsByDimensions(List<CacheStatsDimension> dimensions) {
         internalIncrement(dimensions, (response, amount) -> response.hits.inc(amount), 1);
     }
@@ -121,13 +176,51 @@ public long count() {
     }
 
     private void internalIncrement(List<CacheStatsDimension> dimensions, BiConsumer<CacheStatsResponse, Long> incrementer, long amount) {
-        CacheStatsResponse stats = internalGetStats(dimensions);
-        incrementer.accept(stats, amount);
-        incrementer.accept(totalStats, amount);
+        for (CacheStatsResponse stats : getStatsToIncrement(dimensions)) {
+            incrementer.accept(stats, amount);
+            incrementer.accept(totalStats, amount);
+        }
+    }
+
+    private List<CacheStatsResponse> getStatsToIncrement(List<CacheStatsDimension> keyDimensions) {
+        List<CacheStatsResponse> result = new ArrayList<>();
+        switch (mode) {
+            case SEPARATE_DIMENSIONS_ONLY:
+                for (CacheStatsDimension dim : keyDimensions) {
+                    result.add(internalGetStats(List.of(dim)));
+                }
+                break;
+            case ALL_COMBINATIONS:
+                assert keyDimensions.size() == dimensionNames.size();
+                result.add(internalGetStats(keyDimensions));
+                break;
+            case SPECIFIC_COMBINATIONS:
+                for (Set<String> combination : specificCombinations) {
+                    result.add(internalGetStats(filterDimensionsMatchingCombination(combination, keyDimensions)));
+                }
+                break;
+        }
+        return result;
+    }
+
+    private List<CacheStatsDimension> filterDimensionsMatchingCombination(
+        Set<String> dimCombination,
+        List<CacheStatsDimension> dimensions
+    ) {
+        List<CacheStatsDimension> result = new ArrayList<>();
+        for (CacheStatsDimension dim : dimensions) {
+            if (dimCombination.contains(dim.dimensionName)) {
+                result.add(dim);
+            }
+        }
+        return result;
+    }
+
+    Set<Set<String>> getSpecificCombinations() {
+        return specificCombinations;
     }
 
     private CacheStatsResponse internalGetStats(List<CacheStatsDimension> dimensions) {
-        assert dimensions.size() == dimensionNames.size();
         CacheStatsResponse response = statsMap.get(new Key(dimensions));
         if (response == null) {
             response = new CacheStatsResponse();
@@ -151,6 +244,13 @@ public void writeTo(StreamOutput out) throws IOException {
         );
         totalStats.writeTo(out);
         out.writeVInt(maxDimensionValues);
+        out.writeEnum(mode);
+        // Write Set<Set<String>> as repeated String[]
+        out.writeVInt(specificCombinations.size());
+        for (Set<String> combination : specificCombinations) {
+            out.writeStringArray(combination.toArray(new String[0]));
+        }
+
     }
 
     /**

server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java

@@ -58,7 +58,7 @@ public OpenSearchOnHeapCache(Builder<K, V> builder) {
         }
         cache = cacheBuilder.build();
         this.dimensionNames = Objects.requireNonNull(builder.dimensionNames, "Dimension names can't be null");
-        this.statsHolder = new StatsHolder(dimensionNames, builder.getSettings());
+        this.statsHolder = new StatsHolder(dimensionNames, builder.getSettings(), StatsHolder.TrackingMode.ALL_COMBINATIONS);
         this.removalListener = builder.getRemovalListener();
     }