Modified factories to take new arguments

Signed-off-by: Peter Alfonsi <petealft@amazon.com>

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

@@ -482,6 +482,7 @@ public void onEvent(CacheEvent<? extends ICacheKey<K>, ? extends byte[]> event)
                     this.removalListener.onRemoval(new RemovalNotification<>(event.getKey(), valueSerializer.deserialize(event.getOldValue()), RemovalReason.EVICTED));
                     stats.decrementEntriesByDimensions(event.getKey().dimensions);
                     stats.incrementMemorySizeByDimensions(event.getKey().dimensions, -getOldValuePairSize(event));
+                    stats.incrementEvictionsByDimensions(event.getKey().dimensions);
                     assert event.getNewValue() == null;
                     break;
                 case REMOVED:
@@ -556,11 +557,30 @@ public EhcacheDiskCacheFactory() {}
         public <K, V> ICache<K, V> create(CacheConfig<K, V> config, CacheType cacheType, Map<String, Factory> cacheFactories) {
             Map<String, Setting<?>> settingList = EhcacheDiskCacheSettings.getSettingListForCacheType(cacheType);
             Settings settings = config.getSettings();
+
+            Serializer<K, byte[]> keySerializer = null;
+            try {
+                keySerializer = (Serializer<K, byte[]>) config.getKeySerializer();
+            } catch (ClassCastException e) {
+                throw new IllegalArgumentException("EhcacheDiskCache requires a key serializer of type Serializer<K, byte[]>");
+            }
+
+            Serializer<V, byte[]> valueSerializer = null;
+            try {
+                valueSerializer = (Serializer<V, byte[]>) config.getValueSerializer();
+            } catch (ClassCastException e) {
+                throw new IllegalArgumentException("EhcacheDiskCache requires a value serializer of type Serializer<V, byte[]>");
+            }
+
             return new Builder<K, V>().setStoragePath((String) settingList.get(DISK_STORAGE_PATH_KEY).get(settings))
                 .setDiskCacheAlias((String) settingList.get(DISK_CACHE_ALIAS_KEY).get(settings))
                 .setCacheType(cacheType)
                 .setKeyType((config.getKeyType()))
                 .setValueType(config.getValueType())
+                .setKeySerializer(keySerializer)
+                .setValueSerializer(valueSerializer)
+                .setShardIdDimensionName(config.getDimensionNames().get(0)) // TODO: Rework this to pass in whole list, once stats is changed
+                .setWeigher(config.getWeigher())
                 .setRemovalListener(config.getRemovalListener())
                 .setExpireAfterAccess((TimeValue) settingList.get(DISK_CACHE_EXPIRE_AFTER_ACCESS_KEY).get(settings))
                 .setMaximumWeightInBytes((Long) settingList.get(DISK_MAX_SIZE_IN_BYTES_KEY).get(settings))

plugins/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java

@@ -110,6 +110,10 @@ public void testBasicGetAndPutUsingFactory() throws IOException {
                 new CacheConfig.Builder<String, String>().setValueType(String.class)
                     .setKeyType(String.class)
                     .setRemovalListener(removalListener)
+                    .setKeySerializer(new StringSerializer())
+                    .setValueSerializer(new StringSerializer())
+                    .setDimensionNames(List.of(dimensionName))
+                    .setWeigher(getWeigher())
                     .setSettings(
                         Settings.builder()
                             .put(
@@ -316,19 +320,11 @@ public void testEvictions() throws Exception {
             String value = generateRandomString(100);
 
             // Trying to generate more than 100kb to cause evictions.
-            long sizeOfAttemptedAdds = 0;
-            long sizeOfAttemptedAddsValue = 0;
             for (int i = 0; i < 1000; i++) {
                 String key = "Key" + i;
                 ICacheKey<String> iCacheKey = getICacheKey((key));
-                sizeOfAttemptedAdds += weigher.applyAsLong(iCacheKey, value);
                 ehcacheTest.put(iCacheKey, value);
-
             }
-            /*System.out.println("Total size of attempted adds = " + sizeOfAttemptedAdds);
-            System.out.println("Total size of attempted adds (value only) = " + sizeOfAttemptedAddsValue);
-            System.out.println("Total memory size = " + ehcacheTest.stats().getTotalMemorySize());*/
-            // TODO: Figure out why ehcache is evicting at ~30-40% of its max size rather than 100% (see commented out prints above)
             assertTrue(mockRemovalListener.onRemovalCount.get() > 0);
             assertEquals(660, ehcacheTest.stats().getTotalEvictions());
             ehcacheTest.close();
@@ -539,8 +535,9 @@ public String load(ICacheKey<String> key) throws Exception {
     }
 
     public void testMemoryTracking() throws Exception {
-        // This test leaks threads because of an issue in Ehcache:
+        // TODO: This test leaks threads because of an issue in Ehcache:
         // https://github.com/ehcache/ehcache3/issues/3204
+
         // Test all cases for EhCacheEventListener.onEvent and check stats memory usage is updated correctly
         Settings settings = Settings.builder().build();
         ToLongBiFunction<ICacheKey<String>, String> weigher = getWeigher();
@@ -685,7 +682,7 @@ public void onRemoval(RemovalNotification<ICacheKey<K>, V> notification) {
         }
     }
 
-    private static class StringSerializer implements Serializer<String, byte[]> {
+    static class StringSerializer implements Serializer<String, byte[]> {
         private final Charset charset = StandardCharsets.UTF_8;
         @Override
         public byte[] serialize(String object) {

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

@@ -142,9 +142,12 @@ public static class OpenSearchOnHeapCacheFactory implements Factory {
         public <K, V> ICache<K, V> create(CacheConfig<K, V> config, CacheType cacheType, Map<String, Factory> cacheFactories) {
             Map<String, Setting<?>> settingList = OpenSearchOnHeapCacheSettings.getSettingListForCacheType(cacheType);
             Settings settings = config.getSettings();
-            return new Builder<K, V>().setMaximumWeightInBytes(
-                ((ByteSizeValue) settingList.get(MAXIMUM_SIZE_IN_BYTES_KEY).get(settings)).getBytes()
-            ).setWeigher(config.getWeigher()).setRemovalListener(config.getRemovalListener()).build();
+            return new Builder<K, V>()
+                .setShardIdDimensionName(config.getDimensionNames().get(0)) //TODO: Make it accept >1 dimension names
+                .setMaximumWeightInBytes(((ByteSizeValue) settingList.get(MAXIMUM_SIZE_IN_BYTES_KEY).get(settings)).getBytes())
+                .setWeigher(config.getWeigher())
+                .setRemovalListener(config.getRemovalListener())
+                .build();
         }
 
         @Override

server/src/main/java/org/opensearch/common/cache/store/config/CacheConfig.java

@@ -11,8 +11,10 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.cache.RemovalListener;
 import org.opensearch.common.cache.ICacheKey;
+import org.opensearch.common.cache.serializer.Serializer;
 import org.opensearch.common.settings.Settings;
 
+import java.util.List;
 import java.util.function.ToLongBiFunction;
 
 /**
@@ -42,12 +44,20 @@ public class CacheConfig<K, V> {
 
     private final RemovalListener<ICacheKey<K>, V> removalListener;
 
+    private final Serializer<K, ?> keySerializer;
+    private final Serializer<V, ?> valueSerializer;
+
+    private final List<String> dimensionNames;
+
     private CacheConfig(Builder<K, V> builder) {
         this.keyType = builder.keyType;
         this.valueType = builder.valueType;
         this.settings = builder.settings;
         this.removalListener = builder.removalListener;
         this.weigher = builder.weigher;
+        this.keySerializer = builder.keySerializer;
+        this.valueSerializer = builder.valueSerializer;
+        this.dimensionNames = builder.dimensionNames;
     }
 
     public RemovalListener<ICacheKey<K>, V> getRemovalListener() {
@@ -70,6 +80,18 @@ public ToLongBiFunction<ICacheKey<K>, V> getWeigher() {
         return weigher;
     }
 
+    public Serializer<K, ?> getKeySerializer() {
+        return keySerializer;
+    }
+
+    public Serializer<V, ?> getValueSerializer() {
+        return valueSerializer;
+    }
+
+    public List<String> getDimensionNames() {
+        return dimensionNames;
+    }
+
     /**
      * Builder class to build Cache config related parameters.
      * @param <K> Type of key.
@@ -86,6 +108,11 @@ public static class Builder<K, V> {
 
         private ToLongBiFunction<ICacheKey<K>, V> weigher;
 
+        private Serializer<K, ?> keySerializer;
+        private Serializer<V, ?> valueSerializer;
+
+        private List<String> dimensionNames;
+
         public Builder() {}
 
         public Builder<K, V> setSettings(Settings settings) {
@@ -113,6 +140,21 @@ public Builder<K, V> setWeigher(ToLongBiFunction<ICacheKey<K>, V> weigher) {
             return this;
         }
 
+        public Builder<K, V> setKeySerializer(Serializer<K, ?> keySerializer) {
+            this.keySerializer = keySerializer;
+            return this;
+        }
+
+        public Builder<K, V> setValueSerializer(Serializer<V, ?> valueSerializer) {
+            this.valueSerializer = valueSerializer;
+            return this;
+        }
+
+        public Builder<K, V> setDimensionNames(List<String> dimensionNames) {
+            this.dimensionNames = dimensionNames;
+            return this;
+        }
+
         public CacheConfig<K, V> build() {
             return new CacheConfig<>(this);
         }