1 files changed, 15 insertions, 11 deletions

diff --git a/core/src/main/java/org/elasticsearch/common/lucene/ShardCoreKeyMap.java b/core/src/main/java/org/elasticsearch/common/lucene/ShardCoreKeyMap.java
index 38e7691f39..73255aef13 100644
--- a/core/src/main/java/org/elasticsearch/common/lucene/ShardCoreKeyMap.java
+++ b/core/src/main/java/org/elasticsearch/common/lucene/ShardCoreKeyMap.java
@@ -19,8 +19,8 @@
 
 package org.elasticsearch.common.lucene;
 
+import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReader;
-import org.apache.lucene.index.LeafReader.CoreClosedListener;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.index.shard.ShardUtils;
 
@@ -46,8 +46,8 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public final class ShardCoreKeyMap {
 
-    private final Map<Object, ShardId> coreKeyToShard;
-    private final Map<String, Set<Object>> indexToCoreKey;
+    private final Map<IndexReader.CacheKey, ShardId> coreKeyToShard;
+    private final Map<String, Set<IndexReader.CacheKey>> indexToCoreKey;
 
     public ShardCoreKeyMap() {
         coreKeyToShard = new ConcurrentHashMap<>();
@@ -63,7 +63,11 @@ public final class ShardCoreKeyMap {
         if (shardId == null) {
             throw new IllegalArgumentException("Could not extract shard id from " + reader);
         }
-        final Object coreKey = reader.getCoreCacheKey();
+        final IndexReader.CacheHelper cacheHelper = reader.getCoreCacheHelper();
+        if (cacheHelper == null) {
+            throw new IllegalArgumentException("Reader " + reader + " does not support caching");
+        }
+        final IndexReader.CacheKey coreKey = cacheHelper.getKey();
 
         if (coreKeyToShard.containsKey(coreKey)) {
             // Do this check before entering the synchronized block in order to
@@ -75,18 +79,18 @@ public final class ShardCoreKeyMap {
         final String index = shardId.getIndexName();
         synchronized (this) {
             if (coreKeyToShard.containsKey(coreKey) == false) {
-                Set<Object> objects = indexToCoreKey.get(index);
+                Set<IndexReader.CacheKey> objects = indexToCoreKey.get(index);
                 if (objects == null) {
                     objects = new HashSet<>();
                     indexToCoreKey.put(index, objects);
                 }
                 final boolean added = objects.add(coreKey);
                 assert added;
-                CoreClosedListener listener = ownerCoreCacheKey -> {
+                IndexReader.ClosedListener listener = ownerCoreCacheKey -> {
                     assert coreKey == ownerCoreCacheKey;
                     synchronized (ShardCoreKeyMap.this) {
                         coreKeyToShard.remove(ownerCoreCacheKey);
-                        final Set<Object> coreKeys = indexToCoreKey.get(index);
+                        final Set<IndexReader.CacheKey> coreKeys = indexToCoreKey.get(index);
                         final boolean removed = coreKeys.remove(coreKey);
                         assert removed;
                         if (coreKeys.isEmpty()) {
@@ -96,7 +100,7 @@ public final class ShardCoreKeyMap {
                 };
                 boolean addedListener = false;
                 try {
-                    reader.addCoreClosedListener(listener);
+                    cacheHelper.addClosedListener(listener);
                     addedListener = true;
 
                     // Only add the core key to the map as a last operation so that
@@ -131,7 +135,7 @@ public final class ShardCoreKeyMap {
      * Get the set of core cache keys associated with the given index.
      */
     public synchronized Set<Object> getCoreKeysForIndex(String index) {
-        final Set<Object> objects = indexToCoreKey.get(index);
+        final Set<IndexReader.CacheKey> objects = indexToCoreKey.get(index);
         if (objects == null) {
             return Collections.emptySet();
         }
@@ -154,9 +158,9 @@ public final class ShardCoreKeyMap {
         if (assertionsEnabled == false) {
             throw new AssertionError("only run this if assertions are enabled");
         }
-        Collection<Set<Object>> values = indexToCoreKey.values();
+        Collection<Set<IndexReader.CacheKey>> values = indexToCoreKey.values();
         int size = 0;
-        for (Set<Object> value : values) {
+        for (Set<IndexReader.CacheKey> value : values) {
             size += value.size();
         }
         return size == coreKeyToShard.size();