From 91178c31bf449242aa41a61fd67ced4e2a49a624 Mon Sep 17 00:00:00 2001
From: Steve Wilkerson <sw5822@att.com>
Date: Sat, 4 May 2019 06:32:11 -0600
Subject: [PATCH] Elasticsearch: Add data node specific entrypoint override

This updates the Elasticsearch chart to include a specific start
script for the Elasticsearch data nodes that includes a trap on
signals that removes a data node from allocation eligible nodes
before shutting down.  This results in all shards being moved from
a node on shut down to alleviate issues with planned down nodes,
such as during upgrade scenarios

Change-Id: I22f4957f90e4113831a8ddf48691cb14f811c1e5
---
 .../templates/bin/_elasticsearch.sh.tpl       | 44 +++++++++++++++++++
 elasticsearch/templates/statefulset-data.yaml | 24 ++++++----
 elasticsearch/values.yaml                     |  2 +-
 3 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/elasticsearch/templates/bin/_elasticsearch.sh.tpl b/elasticsearch/templates/bin/_elasticsearch.sh.tpl
index 780ec6e76..0cc1353a8 100644
--- a/elasticsearch/templates/bin/_elasticsearch.sh.tpl
+++ b/elasticsearch/templates/bin/_elasticsearch.sh.tpl
@@ -27,4 +27,48 @@ function stop () {
   kill -TERM 1
 }
 
+function allocate_data_node () {
+  CLUSTER_SETTINGS=$(curl -K- <<< "--user ${ELASTICSEARCH_USERNAME}:${ELASTICSEARCH_PASSWORD}" \
+    "${ELASTICSEARCH_ENDPOINT}/_cluster/settings")
+  if echo "${CLUSTER_SETTINGS}" | grep -E "${NODE_NAME}"; then
+    echo "Activate node ${NODE_NAME}"
+    curl -K- <<< "--user ${ELASTICSEARCH_USERNAME}:${ELASTICSEARCH_PASSWORD}" -XPUT -H 'Content-Type: application/json' \
+     "${ELASTICSEARCH_ENDPOINT}/_cluster/settings" -d "{
+      \"transient\" :{
+          \"cluster.routing.allocation.exclude._name\" : null
+      }
+    }"
+  fi
+  echo "Node ${NODE_NAME} is ready to be used"
+}
+
+function start_data_node () {
+  ulimit -l unlimited
+  allocate_data_node &
+  /docker-entrypoint.sh elasticsearch &
+  function drain_data_node () {
+    echo "Prepare to migrate data off node ${NODE_NAME}"
+    echo "Move all data from node ${NODE_NAME}"
+    curl -K- <<< "--user ${ELASTICSEARCH_USERNAME}:${ELASTICSEARCH_PASSWORD}" -XPUT -H 'Content-Type: application/json' \
+     "${ELASTICSEARCH_ENDPOINT}/_cluster/settings" -d "{
+      \"transient\" :{
+          \"cluster.routing.allocation.exclude._name\" : \"${NODE_NAME}\"
+      }
+    }"
+    echo ""
+    while true ; do
+      echo -e "Wait for node ${NODE_NAME} to become empty"
+      SHARDS_ALLOCATION=$(curl -K- <<< "--user ${ELASTICSEARCH_USERNAME}:${ELASTICSEARCH_PASSWORD}" \
+        -XGET "${ELASTICSEARCH_ENDPOINT}/_cat/shards")
+      if ! echo "${SHARDS_ALLOCATION}" | grep -E "${NODE_NAME}"; then
+        break
+      fi
+      sleep 5
+    done
+    echo "Node ${NODE_NAME} is ready to shutdown"
+  }
+  trap drain_data_node TERM EXIT HUP INT
+  wait
+}
+
 $COMMAND
diff --git a/elasticsearch/templates/statefulset-data.yaml b/elasticsearch/templates/statefulset-data.yaml
index 34fb1890c..a585266c6 100644
--- a/elasticsearch/templates/statefulset-data.yaml
+++ b/elasticsearch/templates/statefulset-data.yaml
@@ -17,6 +17,8 @@ limitations under the License.
 {{- if .Values.manifests.statefulset_data }}
 {{- $envAll := . }}
 
+{{- $esUserSecret := .Values.secrets.elasticsearch.user }}
+
 {{- $mounts_elasticsearch := .Values.pod.mounts.elasticsearch.elasticsearch }}
 
 {{- $serviceAccountName := printf "%s-%s" .Release.Name "elasticsearch-data" }}
@@ -106,13 +108,7 @@ spec:
 {{ dict "envAll" $envAll "application" "data" "container" "elasticsearch_data" | include "helm-toolkit.snippets.kubernetes_container_security_context" | indent 10 }}
           command:
             - /tmp/elasticsearch.sh
-            - start
-          lifecycle:
-            preStop:
-              exec:
-                command:
-                  - /tmp/elasticsearch.sh
-                  - stop
+            - start_data_node
           ports:
             - name: transport
               containerPort: {{ tuple "elasticsearch" "internal" "discovery" . | include "helm-toolkit.endpoints.endpoint_port_lookup" }}
@@ -129,7 +125,19 @@ spec:
             - name: NODE_NAME
               valueFrom:
                 fieldRef:
-                  fieldPath: metadata.name
+                      fieldPath: metadata.name
+            - name: ELASTICSEARCH_USERNAME
+              valueFrom:
+                secretKeyRef:
+                  name: {{ $esUserSecret }}
+                  key: ELASTICSEARCH_USERNAME
+            - name: ELASTICSEARCH_PASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: {{ $esUserSecret }}
+                  key: ELASTICSEARCH_PASSWORD
+            - name: ELASTICSEARCH_ENDPOINT
+              value: {{ tuple "elasticsearch" "internal" "http" . | include "helm-toolkit.endpoints.host_and_port_endpoint_uri_lookup" }}
             - name: NODE_MASTER
               value: "false"
             - name: NODE_INGEST
diff --git a/elasticsearch/values.yaml b/elasticsearch/values.yaml
index 1fe09a089..844cb8103 100644
--- a/elasticsearch/values.yaml
+++ b/elasticsearch/values.yaml
@@ -217,7 +217,7 @@ pod:
       master:
         timeout: 600
       data:
-        timeout: 600
+        timeout: 1200
       client:
         timeout: 600
       prometheus_elasticsearch_exporter: