From 881593eea9de0c6d6d726c163996bb387081adb0 Mon Sep 17 00:00:00 2001
From: Jason Tedor <jason@tedor.me>
Date: Thu, 8 Oct 2015 17:56:43 -0400
Subject: [PATCH] Safe locking and unlocking of segments during invalidation

---
 .../org/elasticsearch/common/cache/Cache.java | 39 ++++++++++++-------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/common/cache/Cache.java b/core/src/main/java/org/elasticsearch/common/cache/Cache.java
index 92b3b899e46..225b8fa5ea5 100644
--- a/core/src/main/java/org/elasticsearch/common/cache/Cache.java
+++ b/core/src/main/java/org/elasticsearch/common/cache/Cache.java
@@ -249,7 +249,8 @@ public class Cache<K, V> {
         }
     }
 
-    private final CacheSegment<K, V>[] segments = new CacheSegment[256];
+    public static final int NUMBER_OF_SEGMENTS = 256;
+    private final CacheSegment<K, V>[] segments = new CacheSegment[NUMBER_OF_SEGMENTS];
 
     {
         for (int i = 0; i < segments.length; i++) {
@@ -368,20 +369,32 @@ public class Cache<K, V> {
      */
     public void invalidateAll() {
         Entry<K, V> h;
-        Arrays.stream(segments).forEach(segment -> segment.segmentLock.writeLock().lock());
-        try (ReleasableLock ignored = lruLock.acquire()) {
-            h = head;
-            Arrays.stream(segments).forEach(segment -> segment.map = new HashMap<>());
-            Entry<K, V> current = head;
-            while (current != null) {
-                current.state = State.DELETED;
-                current = current.after;
+
+        boolean[] haveSegmentLock = new boolean[NUMBER_OF_SEGMENTS];
+        try {
+            for (int i = 0; i < NUMBER_OF_SEGMENTS; i++) {
+                segments[i].segmentLock.writeLock().lock();
+                haveSegmentLock[i] = true;
+            }
+            try (ReleasableLock ignored = lruLock.acquire()) {
+                h = head;
+                Arrays.stream(segments).forEach(segment -> segment.map = new HashMap<>());
+                Entry<K, V> current = head;
+                while (current != null) {
+                    current.state = State.DELETED;
+                    current = current.after;
+                }
+                head = tail = null;
+                count = 0;
+                weight = 0;
+            }
+        } finally {
+            for (int i = 0; i < NUMBER_OF_SEGMENTS; i++) {
+                if (haveSegmentLock[i]) {
+                    segments[i].segmentLock.writeLock().unlock();
+                }
             }
-            head = tail = null;
-            count = 0;
-            weight = 0;
         }
-        Arrays.stream(segments).forEach(segment -> segment.segmentLock.writeLock().unlock());
         while (h != null) {
             removalListener.onRemoval(new RemovalNotification<>(h.key, h.value, RemovalNotification.RemovalReason.INVALIDATED));
             h = h.after;