3 files changed, 77 insertions, 34 deletions

diff --git a/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/CacheKeyFactory.java b/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/CacheKeyFactory.java
index 1b444ba3c..bf4659e89 100644
--- a/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/CacheKeyFactory.java
+++ b/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/CacheKeyFactory.java
@@ -27,5 +27,12 @@ public interface CacheKeyFactory {
     /**
      * Construct key accordingly to provided {@code InstanceIdentifier<?>}
      */
+    @Nonnull
     String createKey(@Nonnull final InstanceIdentifier<?> actualContextIdentifier);
+
+    /**
+     * Returns type of data, for which is this factory creating keys
+     */
+    @Nonnull
+    Class<?> getCachedDataType();
 }
diff --git a/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/DumpCacheManager.java b/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/DumpCacheManager.java
index f1b265dec..adcd32e0c 100644
--- a/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/DumpCacheManager.java
+++ b/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/DumpCacheManager.java
@@ -17,6 +17,8 @@
 package io.fd.honeycomb.translate.util.read.cache;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
+import static java.util.Objects.nonNull;
 
 import com.google.common.base.Optional;
 import io.fd.honeycomb.translate.ModificationCache;
@@ -41,11 +43,13 @@ public final class DumpCacheManager<T, U> {
     private final EntityDumpExecutor<T, U> dumpExecutor;
     private final EntityDumpPostProcessingFunction<T> postProcessor;
     private final CacheKeyFactory cacheKeyFactory;
+    private final Class<?> acceptOnly;
 
     private DumpCacheManager(DumpCacheManagerBuilder<T, U> builder) {
         this.dumpExecutor = builder.dumpExecutor;
         this.postProcessor = builder.postProcessingFunction;
         this.cacheKeyFactory = builder.cacheKeyFactory;
+        this.acceptOnly = builder.acceptOnly;
     }
 
     /**
@@ -79,6 +83,12 @@ public final class DumpCacheManager<T, U> {
             cache.put(entityKey, dump);
             return Optional.of(dump);
         } else {
+            // if specified, check whether data returned from cache can be used as result of this dump manager
+            // used as a secondary check if cache does not have any data of different type stored under the same key
+            checkState(acceptOnly.isInstance(dump),
+                    "This dump manager accepts only %s as data, but %s was returned from cache",
+                    acceptOnly, dump.getClass());
+
             LOG.debug("Cached instance of dump was found for KEY[{}]", entityKey);
             return Optional.of(dump);
         }
@@ -86,18 +96,14 @@ public final class DumpCacheManager<T, U> {
 
     public static final class DumpCacheManagerBuilder<T, U> {
 
-        private static final CacheKeyFactory DEFAULT_CACHE_KEY_FACTORY_INSTANCE = new IdentifierCacheKeyFactory();
-
         private EntityDumpExecutor<T, U> dumpExecutor;
         private EntityDumpPostProcessingFunction<T> postProcessingFunction;
         private CacheKeyFactory cacheKeyFactory;
+        private Class<?> acceptOnly;
 
         public DumpCacheManagerBuilder() {
             // for cases when user does not set specific post-processor
             postProcessingFunction = new NoopDumpPostProcessingFunction<T>();
-
-            //use no additional scopes version by default
-            cacheKeyFactory = DEFAULT_CACHE_KEY_FACTORY_INSTANCE;
         }
 
         public DumpCacheManagerBuilder<T, U> withExecutor(@Nonnull final EntityDumpExecutor<T, U> executor) {
@@ -111,17 +117,37 @@ public final class DumpCacheManager<T, U> {
             return this;
         }
 
+        /**
+         * Key providing unique(type-aware) keys.
+         */
         public DumpCacheManagerBuilder<T, U> withCacheKeyFactory(@Nonnull final CacheKeyFactory cacheKeyFactory) {
             this.cacheKeyFactory = cacheKeyFactory;
             return this;
         }
 
+        /**
+         * If modification returns object of different type that this, throw exception to prevent processing data
+         * of different type.
+         */
+        public DumpCacheManagerBuilder<T, U> acceptOnly(@Nonnull final Class<?> acceptOnly) {
+            this.acceptOnly = acceptOnly;
+            return this;
+        }
+
         public DumpCacheManager<T, U> build() {
             checkNotNull(dumpExecutor, "Dump executor cannot be null");
             checkNotNull(postProcessingFunction,
                     "Dump post-processor cannot be null cannot be null, default implementation is used when not set explicitly");
-            checkNotNull(cacheKeyFactory,
-                    "Cache key factory cannot be null, default non-extended implementation is used when not set explicitly");
+
+            if (acceptOnly != null) {
+                cacheKeyFactory = new TypeAwareIdentifierCacheKeyFactory(acceptOnly);
+            } else if (cacheKeyFactory != null) {
+                acceptOnly = cacheKeyFactory.getCachedDataType();
+            } else {
+                throw new IllegalStateException(
+                        "Invalid combination - either acceptOnly type must be defined[defined=" + nonNull(acceptOnly) +
+                                "], or type-aware cache key factory[defined=" + nonNull(cacheKeyFactory) + "]");
+            }
 
             return new DumpCacheManager<>(this);
         }
diff --git a/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/IdentifierCacheKeyFactory.java b/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/TypeAwareIdentifierCacheKeyFactory.java
index 51f47e137..ba4e7e493 100644
--- a/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/IdentifierCacheKeyFactory.java
+++ b/infra/translate-utils/src/main/java/io/fd/honeycomb/translate/util/read/cache/TypeAwareIdentifierCacheKeyFactory.java
@@ -32,31 +32,51 @@ import org.opendaylight.yangtools.yang.binding.InstanceIdentifier;
 /**
  * Factory providing cache keys to easier switching between scopes of caching
  */
-public final class IdentifierCacheKeyFactory implements CacheKeyFactory {
+public final class TypeAwareIdentifierCacheKeyFactory implements CacheKeyFactory {
 
     private static final String KEY_PARTS_SEPARATOR = "|";
 
-    // should be Set<Class<? extends DataObject & Identificable<?>>>, but that's not possible for wildcards
+    // should be Set<Class<? extends DataObject & Identifiable<?>>>, but that's not possible for wildcards
     private final Set<Class<? extends DataObject>> additionalKeyTypes;
+    // factory must be aware of type of data, to prevent creating same key for same identifier but different data
+    private final Class<?> type;
 
     /**
      * Construct simple cache key factory
      */
-    public IdentifierCacheKeyFactory() {
-        this(Collections.emptySet());
+    public TypeAwareIdentifierCacheKeyFactory(@Nonnull final Class<?> type) {
+        this(type, Collections.emptySet());
     }
 
     /**
      * @param additionalKeyTypes Additional types from path of cached type, that are specifying scope
      */
-    public IdentifierCacheKeyFactory(@Nonnull final Set<Class<? extends DataObject>> additionalKeyTypes) {
+    public TypeAwareIdentifierCacheKeyFactory(@Nonnull final Class<?> type,
+                                              @Nonnull final Set<Class<? extends DataObject>> additionalKeyTypes) {
         // verify that all are non-null and identifiable
+        this.type = checkNotNull(type, "Type cannot be null");
         this.additionalKeyTypes = checkNotNull(additionalKeyTypes, "Additional key types can't be null").stream()
-                .map(IdentifierCacheKeyFactory::verifyNotNull)
-                .map(IdentifierCacheKeyFactory::verifyIsIdentifiable)
+                .map(TypeAwareIdentifierCacheKeyFactory::verifyNotNull)
+                .map(TypeAwareIdentifierCacheKeyFactory::verifyIsIdentifiable)
                 .collect(Collectors.toSet());
     }
 
+    private static String bindKeyString(IdentifiableItem identifiableItem) {
+        return String.format("%s[%s]", identifiableItem.getType().getTypeName(), identifiableItem.getKey());
+    }
+
+    private static Class<? extends DataObject> verifyNotNull(final Class<? extends DataObject> type) {
+        return checkNotNull(type, "Cannot use null as key");
+    }
+
+    /**
+     * Initial check if provided scope variables are identifiable aka. can be used to create unique cache key
+     */
+    private static Class<? extends DataObject> verifyIsIdentifiable(final Class<? extends DataObject> type) {
+        checkArgument(Identifiable.class.isAssignableFrom(type), "Type %s is not Identifiable", type);
+        return type;
+    }
+
     @Override
     public String createKey(@Nonnull final InstanceIdentifier<?> actualContextIdentifier) {
 
@@ -64,18 +84,25 @@ public final class IdentifierCacheKeyFactory implements CacheKeyFactory {
 
         // easiest case when only simple key is needed
         if (additionalKeyTypes.isEmpty()) {
-            return actualContextIdentifier.getTargetType().toString();
+            return String
+                    .join(KEY_PARTS_SEPARATOR, type.getTypeName(), actualContextIdentifier.getTargetType().toString());
         }
 
         checkArgument(isUniqueKeyConstructable(actualContextIdentifier),
                 "Unable to construct unique key, required key types : %s, provided paths : %s", additionalKeyTypes,
                 actualContextIdentifier.getPathArguments());
 
+        // joins unique key in form : type | additional keys | actual context
         return String
-                .join(KEY_PARTS_SEPARATOR, additionalKeys(actualContextIdentifier),
+                .join(KEY_PARTS_SEPARATOR, type.getTypeName(), additionalKeys(actualContextIdentifier),
                         actualContextIdentifier.getTargetType().toString());
     }
 
+    @Override
+    public Class<?> getCachedDataType() {
+        return type;
+    }
+
     /**
      * Verifies that all requested key parts have keys
      */
@@ -94,29 +121,12 @@ public final class IdentifierCacheKeyFactory implements CacheKeyFactory {
         return pathArgument instanceof IdentifiableItem;
     }
 
-
     private String additionalKeys(final InstanceIdentifier<?> actualContextIdentifier) {
         return StreamSupport.stream(actualContextIdentifier.getPathArguments().spliterator(), false)
                 .filter(this::isAdditionalScope)
                 .filter(this::isIdentifiable)
                 .map(IdentifiableItem.class::cast)
-                .map(IdentifierCacheKeyFactory::bindKeyString)
+                .map(TypeAwareIdentifierCacheKeyFactory::bindKeyString)
                 .collect(Collectors.joining(KEY_PARTS_SEPARATOR));
     }
-
-    private static String bindKeyString(IdentifiableItem identifiableItem) {
-        return String.format("%s[%s]", identifiableItem.getType().getTypeName(), identifiableItem.getKey());
-    }
-
-    private static Class<? extends DataObject> verifyNotNull(final Class<? extends DataObject> type) {
-        return checkNotNull(type, "Cannot use null as key");
-    }
-
-    /**
-     * Initial check if provided scope variables are identifiable aka. can be used to create unique cache key
-     */
-    private static Class<? extends DataObject> verifyIsIdentifiable(final Class<? extends DataObject> type) {
-        checkArgument(Identifiable.class.isAssignableFrom(type), "Type %s is not Identifiable", type);
-        return type;
-    }
 }