From 6f66053ca7e057b47379c9eb05a665b8b6688c51 Mon Sep 17 00:00:00 2001
From: Theron Voran <tvoran@users.noreply.github.com>
Date: Thu, 30 Jan 2020 09:39:08 -0800
Subject: [PATCH] Adding sleep in the preStop lifecycle step (#188)

Aims to make vault pod termination more graceful with respect to user
requests.
---
 templates/server-statefulset.yaml |  8 +++++++-
 test/unit/server-statefulset.bats | 21 +++++++++++++++++++++
 values.yaml                       |  3 +++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/templates/server-statefulset.yaml b/templates/server-statefulset.yaml
index c89175df2..abde79db5 100644
--- a/templates/server-statefulset.yaml
+++ b/templates/server-statefulset.yaml
@@ -127,7 +127,13 @@ spec:
             # from Consul (zombie services).
             preStop:
               exec:
-                command: ["/bin/sh","-c","kill -SIGTERM $(pidof vault)"]
+                command: [
+                  "/bin/sh", "-c",
+                  # Adding a sleep here to give the pod eviction a
+                  # chance to propagate, so requests will not be made
+                  # to this pod while it's terminating
+                  "sleep {{ .Values.server.preStopSleepSeconds }} && kill -SIGTERM $(pidof vault)",
+                ]
         {{- if .Values.server.extraContainers }}
           {{ toYaml .Values.server.extraContainers | nindent 8}}
         {{- end }}
diff --git a/test/unit/server-statefulset.bats b/test/unit/server-statefulset.bats
index cfc0c4b77..60b54c88e 100755
--- a/test/unit/server-statefulset.bats
+++ b/test/unit/server-statefulset.bats
@@ -841,3 +841,24 @@ load _helpers
        yq -r '.spec.template.spec.containers[0].args[0]' | tee /dev/stderr)
   [[ "${actual}" = *"foobar"* ]]
 }
+
+#--------------------------------------------------------------------
+# preStop
+@test "server/standalone-StatefulSet: preStop sleep duration default" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -x templates/server-statefulset.yaml \
+      . | tee /dev/stderr |
+       yq -r '.spec.template.spec.containers[0].lifecycle.preStop.exec.command[2]' | tee /dev/stderr)
+  [[ "${actual}" = "sleep 5 &&"* ]]
+}
+
+@test "server/standalone-StatefulSet: preStop sleep duration 10" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -x templates/server-statefulset.yaml \
+      --set 'server.preStopSleepSeconds=10' \
+      . | tee /dev/stderr |
+       yq -r '.spec.template.spec.containers[0].lifecycle.preStop.exec.command[2]' | tee /dev/stderr)
+  [[ "${actual}" = "sleep 10 &&"* ]]
+}
diff --git a/values.yaml b/values.yaml
index 3fee150d6..54330266b 100644
--- a/values.yaml
+++ b/values.yaml
@@ -135,6 +135,9 @@ server:
     path: "/v1/sys/health?standbyok=true"
     initialDelaySeconds: 60
 
+  # Used to set the sleep time during the preStop step
+  preStopSleepSeconds: 5
+
   # extraEnvironmentVars is a list of extra enviroment variables to set with the stateful set. These could be
   # used to include variables required for auto-unseal.
   extraEnvironmentVars: {}