From 8635d264ae9e4c2d16e04a0d8b3a6a7db094e0bf Mon Sep 17 00:00:00 2001
From: Alexander Reelsen <alexander@reelsen.net>
Date: Thu, 28 Jan 2016 18:30:27 +0100
Subject: [PATCH] Shield: Give native stores possibility to exit poller loop

Similar to the lifecycle services, stopping the shield lifecycle should
also ensure that the poller threads are stopped, which is tricky, in case
they run through huge user/role lists.

Original commit: elastic/x-pack-elasticsearch@7a48f19853fcb79bc0909864f807d7e0df209302
---
 .../authc/esnative/ESNativeUsersStore.java      | 17 +++++++++++++++++
 .../authz/esnative/ESNativeRolesStore.java      | 11 +++++++++++
 2 files changed, 28 insertions(+)

diff --git a/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authc/esnative/ESNativeUsersStore.java b/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authc/esnative/ESNativeUsersStore.java
index 342e0d41ae0..a01a696783d 100644
--- a/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authc/esnative/ESNativeUsersStore.java
+++ b/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authc/esnative/ESNativeUsersStore.java
@@ -532,6 +532,9 @@ public class ESNativeUsersStore extends AbstractComponent implements ClusterStat
 
         @Override
         public void doRun() {
+            if (isStopped()) {
+                return;
+            }
             if (shieldIndexExists == false) {
                 logger.trace("cannot poll for user changes since shield admin index [{}] does not exist", ShieldTemplateService.SHIELD_ADMIN_INDEX_NAME);
                 return;
@@ -568,6 +571,11 @@ public class ESNativeUsersStore extends AbstractComponent implements ClusterStat
                 }
             }
 
+            // exit before comparing with known users
+            if (isStopped()) {
+                return;
+            }
+
             // we now have a list of users that were in our version map and have been deleted
             Iterator<ObjectCursor<String>> userIter = knownUsers.iterator();
             while (userIter.hasNext()) {
@@ -619,6 +627,10 @@ public class ESNativeUsersStore extends AbstractComponent implements ClusterStat
 
                 boolean keepScrolling = response.getHits().getHits().length > 0;
                 while (keepScrolling) {
+                    if (isStopped()) {
+                        // instead of throwing an exception we return an empty map so nothing is processed and we exit early
+                        return new ObjectLongHashMap<>();
+                    }
                     for (SearchHit hit : response.getHits().getHits()) {
                         String username = hit.id();
                         long version = hit.version();
@@ -638,6 +650,11 @@ public class ESNativeUsersStore extends AbstractComponent implements ClusterStat
             }
             return map;
         }
+
+        private boolean isStopped() {
+            State state = state();
+            return state == State.STOPPED || state == State.STOPPING;
+        }
     }
 
     public static class Fields {
diff --git a/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authz/esnative/ESNativeRolesStore.java b/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authz/esnative/ESNativeRolesStore.java
index e6a4c0ae45d..27ec1810185 100644
--- a/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authz/esnative/ESNativeRolesStore.java
+++ b/elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/authz/esnative/ESNativeRolesStore.java
@@ -460,6 +460,9 @@ public class ESNativeRolesStore extends AbstractComponent implements RolesStore,
 
         @Override
         protected void doRun() throws Exception {
+            if (isStopped()) {
+                return;
+            }
             if (shieldIndexExists == false) {
                 logger.trace("cannot poll for role changes since shield admin index [{}] does not exist", ShieldTemplateService.SHIELD_ADMIN_INDEX_NAME);
                 return;
@@ -485,6 +488,9 @@ public class ESNativeRolesStore extends AbstractComponent implements RolesStore,
 
                 boolean keepScrolling = response.getHits().getHits().length > 0;
                 while (keepScrolling) {
+                    if (isStopped()) {
+                        return;
+                    }
                     for (final SearchHit hit : response.getHits().getHits()) {
                         final String roleName = hit.getId();
                         final long version = hit.version();
@@ -530,6 +536,11 @@ public class ESNativeRolesStore extends AbstractComponent implements RolesStore,
         public void onFailure(Throwable t) {
             logger.error("error occurred while checking the native roles for changes", t);
         }
+
+        private boolean isStopped() {
+            State state = state();
+            return state == State.STOPPED || state == State.STOPPING;
+        }
     }
 
     private static class RoleAndVersion {