From 3896457d84ab11690bb27f98bac4c5dc1cf4e673 Mon Sep 17 00:00:00 2001
From: Nik Everett <nik9000@gmail.com>
Date: Fri, 15 Jan 2016 14:06:16 -0500
Subject: [PATCH] Handle closed readers in ShardCoreKeyMap

Fixes an issue caused by trying to add a LeafReader for a closed index to
ShardCoreKeyMap. It add itself to the map half way before throwing
AlreadyClosedException, leaking some references and causesing Elasticsearch
to refuse to shut down cleanly when testing.
---
 .../common/lucene/ShardCoreKeyMap.java        | 19 +++++++++++++++--
 .../common/lucene/ShardCoreKeyMapTests.java   | 21 +++++++++++++++++++
 2 files changed, 38 insertions(+), 2 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 b1271e7338d..92aa02ba002 100644
--- a/core/src/main/java/org/elasticsearch/common/lucene/ShardCoreKeyMap.java
+++ b/core/src/main/java/org/elasticsearch/common/lucene/ShardCoreKeyMap.java
@@ -20,9 +20,11 @@
 package org.elasticsearch.common.lucene;
 
 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;
 
+import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -72,7 +74,7 @@ public final class ShardCoreKeyMap {
                 }
                 final boolean added = objects.add(coreKey);
                 assert added;
-                reader.addCoreClosedListener(ownerCoreCacheKey -> {
+                CoreClosedListener listener = ownerCoreCacheKey -> {
                     assert coreKey == ownerCoreCacheKey;
                     synchronized (ShardCoreKeyMap.this) {
                         coreKeyToShard.remove(ownerCoreCacheKey);
@@ -83,7 +85,20 @@ public final class ShardCoreKeyMap {
                             indexToCoreKey.remove(index);
                         }
                     }
-                });
+                };
+                boolean addedListener = false;
+                try {
+                    reader.addCoreClosedListener(listener);
+                    addedListener = true;
+                } finally {
+                    if (false == addedListener) {
+                        try {
+                            listener.onClose(coreKey);
+                        } catch (IOException e) {
+                            throw new RuntimeException("Blow up trying to recover from failure to add listener", e);
+                        }
+                    }
+                }
             }
         }
     }
diff --git a/core/src/test/java/org/elasticsearch/common/lucene/ShardCoreKeyMapTests.java b/core/src/test/java/org/elasticsearch/common/lucene/ShardCoreKeyMapTests.java
index 61660f96742..0c14e1a0bcb 100644
--- a/core/src/test/java/org/elasticsearch/common/lucene/ShardCoreKeyMapTests.java
+++ b/core/src/test/java/org/elasticsearch/common/lucene/ShardCoreKeyMapTests.java
@@ -22,8 +22,10 @@ package org.elasticsearch.common.lucene;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.RandomIndexWriter;
+import org.apache.lucene.store.AlreadyClosedException;
 import org.apache.lucene.store.Directory;
 import org.elasticsearch.common.lucene.index.ElasticsearchDirectoryReader;
 import org.elasticsearch.index.shard.ShardId;
@@ -55,6 +57,25 @@ public class ShardCoreKeyMapTests extends ESTestCase {
         }
     }
 
+    public void testAddingAClosedReader() throws Exception {
+        LeafReader reader;
+        try (Directory dir = newDirectory();
+                RandomIndexWriter writer = new RandomIndexWriter(random(), dir)) {
+            writer.addDocument(new Document());
+            try (DirectoryReader dirReader = ElasticsearchDirectoryReader.wrap(writer.getReader(), new ShardId("index1", 1))) {
+                reader = dirReader.leaves().get(0).reader();
+            }
+        }
+        ShardCoreKeyMap map = new ShardCoreKeyMap();
+        try {
+            map.add(reader);
+            fail("Expected AlreadyClosedException");
+        } catch (AlreadyClosedException e) {
+            // What we wanted
+        }
+        assertEquals(0, map.size());
+    }
+
     public void testBasics() throws IOException {
         Directory dir1 = newDirectory();
         RandomIndexWriter w1 = new RandomIndexWriter(random(), dir1);