From 6bd1179a17ad12f01983d7bc94e25468ae74a7e1 Mon Sep 17 00:00:00 2001
From: jjwong <64273338+WJay-tec@users.noreply.github.com>
Date: Sat, 9 Jul 2022 02:42:03 +0800
Subject: [PATCH] Adding podDistruptionBudget to connect injector (#1316)

* Adding podDistruptionBudget to connect injector
---
 charts/consul/templates/_helpers.tpl          |  14 ++
 .../connect-injector-disruptionbudget.yaml    |  26 +++
 .../connect-injector-disruptionbudget.bats    | 161 ++++++++++++++++++
 charts/consul/values.yaml                     |  16 ++
 4 files changed, 217 insertions(+)
 create mode 100644 charts/consul/templates/connect-injector-disruptionbudget.yaml
 create mode 100755 charts/consul/test/unit/connect-injector-disruptionbudget.bats

diff --git a/charts/consul/templates/_helpers.tpl b/charts/consul/templates/_helpers.tpl
index 56f4fbf128..b667bdb460 100644
--- a/charts/consul/templates/_helpers.tpl
+++ b/charts/consul/templates/_helpers.tpl
@@ -197,6 +197,20 @@ Add a special case for replicas=1, where it should default to 0 as well.
 {{- end -}}
 {{- end -}}
 
+{{- define "consul.pdb.connectInject.maxUnavailable" -}}
+{{- if eq (int .Values.connectInject.replicas) 1 -}}
+{{ 0 }}
+{{- else if .Values.connectInject.disruptionBudget.maxUnavailable -}}
+{{ .Values.server.disruptionBudget.maxUnavailable -}}
+{{- else -}}
+{{- if eq (int .Values.connectInject.replicas) 3 -}}
+{{- 1 -}}
+{{- else -}}
+{{- sub (div (int .Values.connectInject.replicas) 2) 1 -}}
+{{- end -}}
+{{- end -}}
+{{- end -}}
+
 {{/*
 Inject extra environment vars in the format key:value, if populated
 */}}
diff --git a/charts/consul/templates/connect-injector-disruptionbudget.yaml b/charts/consul/templates/connect-injector-disruptionbudget.yaml
new file mode 100644
index 0000000000..08f1401fbe
--- /dev/null
+++ b/charts/consul/templates/connect-injector-disruptionbudget.yaml
@@ -0,0 +1,26 @@
+{{- if (and .Values.connectInject.disruptionBudget.enabled (or (and (ne (.Values.connectInject.enabled | toString) "-") .Values.connectInject.enabled) (and (eq (.Values.connectInject.enabled | toString) "-") .Values.global.enabled))) }}
+# PodDisruptionBudget to prevent degrading the connectInject cluster through
+# voluntary cluster changes.
+{{- if .Capabilities.APIVersions.Has "policy/v1/PodDisruptionBudget" }}
+apiVersion: policy/v1
+{{- else }}
+apiVersion: policy/v1beta1
+{{- end }}
+kind: PodDisruptionBudget
+metadata:
+  name: {{ template "consul.fullname" . }}-connect-injector
+  namespace: {{ .Release.Namespace }}
+  labels:
+    app: {{ template "consul.name" . }}
+    chart: {{ template "consul.chart" . }}
+    heritage: {{ .Release.Service }}
+    release: {{ .Release.Name }}
+    component: connect-injector
+spec:
+  maxUnavailable: {{ template "consul.pdb.connectInject.maxUnavailable" . }}
+  selector:
+    matchLabels:
+      app: {{ template "consul.name" . }}
+      release: "{{ .Release.Name }}"
+      component: connect-injector
+{{- end }}
diff --git a/charts/consul/test/unit/connect-injector-disruptionbudget.bats b/charts/consul/test/unit/connect-injector-disruptionbudget.bats
new file mode 100755
index 0000000000..ec998d0750
--- /dev/null
+++ b/charts/consul/test/unit/connect-injector-disruptionbudget.bats
@@ -0,0 +1,161 @@
+#!/usr/bin/env bats
+
+load _helpers
+
+@test "connect-injector/DisruptionBudget: disabled by default" {
+  cd `chart_dir`
+  assert_empty helm template \
+      -s templates/connect-injector-disruptionbudget.yaml .
+}
+
+@test "connect-injector/DisruptionBudget: enabled with connectInject=enabled , connectInject.disruptionBudget.enabled=true and global.enabled=true " {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq 'length > 0' | tee /dev/stderr)
+  [ "${actual}" = "true" ]
+}
+
+@test "connect-injector/DisruptionBudget: disabled with connectInject.enabled=false" {
+  cd `chart_dir`
+  assert_empty helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.enabled=false' \
+      .
+}
+
+@test "connect-injector/DisruptionBudget: disabled with connectInject.disruptionBudget.enabled=false" {
+  cd `chart_dir`
+  assert_empty helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.disruptionBudget.enabled=false' \
+      .
+}
+
+@test "connect-injector/DisruptionBudget: disabled with global.enabled=false" {
+  cd `chart_dir`
+  assert_empty helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'global.enabled=false' \
+      .
+}
+
+#--------------------------------------------------------------------
+# maxUnavailable
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=1" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=1' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "0" ]
+}
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=3" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=3' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "1" ]
+}
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=4" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=4' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "1" ]
+}
+
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=5" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=5' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "1" ]
+}
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=6" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=6' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "2" ]
+}
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=7" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=7' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "2" ]
+}
+
+@test "connect-injector/DisruptionBudget: correct maxUnavailable with replicas=8" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml  \
+      --set 'connectInject.replicas=8' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq '.spec.maxUnavailable' | tee /dev/stderr)
+  [ "${actual}" = "3" ]
+}
+
+#--------------------------------------------------------------------
+# apiVersion
+
+@test "connect-injector/DisruptionBudget: uses policy/v1 if supported" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      -s templates/connect-injector-disruptionbudget.yaml \
+      --api-versions 'policy/v1/PodDisruptionBudget' \
+      --set 'global.enabled=true' \
+      --set 'connectInject.enabled=true' \
+      --set 'connectInject.disruptionBudget.enabled=true' \
+      . | tee /dev/stderr |
+      yq -r '.apiVersion' | tee /dev/stderr)
+  [ "${actual}" = "policy/v1" ]
+}
+# NOTE: can't test that it uses policy/v1beta1 if policy/v1 is *not* supported
+# because the supported API versions depends on the Helm version and there's
+# no flag to *remove* an API version so some Helm versions will always have
+# policy/v1 support and will always use that API version.
+
diff --git a/charts/consul/values.yaml b/charts/consul/values.yaml
index 1c756a951d..2079c41ad5 100644
--- a/charts/consul/values.yaml
+++ b/charts/consul/values.yaml
@@ -1860,6 +1860,22 @@ connectInject:
     # Note: This value has no effect if transparent proxy is disabled on the pod.
     defaultOverwriteProbes: true
 
+  # This configures the PodDisruptionBudget (https://kubernetes.io/docs/tasks/run-application/configure-pdb/)
+  # for the service mesh sidecar injector.
+  disruptionBudget: 
+    # This will enable/disable registering a PodDisruptionBudget for the 
+    # service mesh sidecar injector. If this is enabled, it will only register the budget so long as
+    # the service mesh is enabled.
+    enabled: true
+
+    # The maximum number of unavailable pods. By default, this will be
+    # automatically computed based on the `connectInject.replicas` value to be `(n/2)-1`.
+    # If you need to set this to `0`, you will need to add a
+    # --set 'connectInject.disruptionBudget.maxUnavailable=0'` flag to the helm chart installation
+    # command because of a limitation in the Helm templating language.
+    # @type: integer
+    maxUnavailable: null
+
   # Configures metrics for Consul Connect services. All values are overridable
   # via annotations on a per-pod basis.
   metrics: