cockroachdb · scottcrossen · Jul 18, 2021
diff --git a/cockroachdb/README.md b/cockroachdb/README.md
@@ -307,6 +307,9 @@ For details see the [`values.yaml`](values.yaml) file.
 | `conf.locality`                                           | Locality attribute for this deployment                          | `""`                                                  |
 | `conf.single-node`                                        | Disable CockroachDB clustering (standalone mode)                | `no`                                                  |
 | `conf.sql-audit-dir`                                      | Directory for SQL audit log                                     | `""`                                                  |
+| `conf.useHostLocality.enabled`                            | Use host node labels to set region, zone, host locality keys    | `false`                                               |
+| `conf.useHostLocality.image.repository`                   | The init-container image name to query the kubernetes API       | `scottcrossen/kube-node-labels`                       |
+| `conf.useHostLocality.image`                              | The init-container image tag to query the kubernetes API        | `1.0.0`                                               |
 | `conf.port`                                               | CockroachDB primary serving port in Pods                        | `26257`                                               |
 | `conf.http-port`                                          | CockroachDB HTTP port in Pods                                   | `8080`                                                |
 | `conf.path`                                               | CockroachDB data directory mount path                           | `cockroach-data`                                      |
@@ -352,6 +355,8 @@ For details see the [`values.yaml`](values.yaml) file.
 | `service.public.type`                                     | Public Service type                                             | `ClusterIP`                                           |
 | `service.public.labels`                                   | Additional labels of public Service                             | `{"app.kubernetes.io/component": "cockroachdb"}`      |
 | `service.public.annotations`                              | Additional annotations of public Service                        | `{}`                                                  |
+| `statefulset.serviceAccount.create`                       | Whether to create a new RBAC service account                    | `yes`                                                 |
+| `statefulset.serviceAccount.name`                         | Name of RBAC service account to use                             | `""`                                                  |
 | `service.discovery.labels`                                | Additional labels of discovery Service                          | `{"app.kubernetes.io/component": "cockroachdb"}`      |
 | `service.discovery.annotations`                           | Additional annotations of discovery Service                     | `{}`                                                  |
 | `ingress.enabled`                                         | Enable ingress resource for CockroachDB                         | `false`                                               |
@@ -380,8 +385,6 @@ For details see the [`values.yaml`](values.yaml) file.
 | `init.tolerations`                                        | Node taints to tolerate by init Job Pod                         | `[]`                                                  |
 | `init.resources`                                          | Resource requests and limits for the Pod of init Job            | `{}`                                                  |
 | `tls.enabled`                                             | Whether to run securely using TLS certificates                  | `no`                                                  |
-| `tls.serviceAccount.create`                               | Whether to create a new RBAC service account                    | `yes`                                                 |
-| `tls.serviceAccount.name`                                 | Name of RBAC service account to use                             | `""`                                                  |
 | `tls.certs.provided`                                      | Bring your own certs scenario, i.e certificates are provided    | `no`                                                  |
 | `tls.certs.clientRootSecret`                              | If certs are provided, secret name for client root cert         | `cockroachdb-root`                                    |
 | `tls.certs.nodeSecret`                                    | If certs are provided, secret name for node cert                | `cockroachdb-node`                                    |

diff --git a/cockroachdb/templates/_helpers.tpl b/cockroachdb/templates/_helpers.tpl
@@ -33,11 +33,11 @@ Create chart name and version as used by the chart label.
 {{/*
 Create the name of the ServiceAccount to use.
 */}}
-{{- define "cockroachdb.tls.serviceAccount.name" -}}
-{{- if .Values.tls.serviceAccount.create -}}
-    {{- default (include "cockroachdb.fullname" .) .Values.tls.serviceAccount.name -}}
+{{- define "cockroachdb.statefulset.serviceAccount.name" -}}
+{{- if .Values.statefulset.serviceAccount.create -}}
+    {{- default (include "cockroachdb.fullname" .) .Values.statefulset.serviceAccount.name -}}
 {{- else -}}
-    {{- default "default" .Values.tls.serviceAccount.name -}}
+    {{- default "default" .Values.statefulset.serviceAccount.name -}}
 {{- end -}}
 {{- end -}}
 

diff --git a/cockroachdb/templates/clusterrole.yaml b/cockroachdb/templates/clusterrole.yaml
@@ -1,4 +1,4 @@
-{{- if and .Values.tls.enabled (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager) }}
+{{- if or (and .Values.tls.enabled (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager)) .Values.conf.useHostLocality.enabled }}
 kind: ClusterRole
 apiVersion: rbac.authorization.k8s.io/v1
 metadata:
@@ -13,7 +13,14 @@ metadata:
     {{- toYaml . | nindent 4 }}
   {{- end }}
 rules:
+  {{- if and .Values.tls.enabled (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager) }}
   - apiGroups: ["certificates.k8s.io"]
     resources: ["certificatesigningrequests"]
     verbs: ["create", "get", "watch"]
-{{- end }}
+  {{- end }}
+  {{- if .Values.conf.useHostLocality.enabled }}
+  - apiGroups: [""]
+    resources: ["nodes"]
+    verbs: ["get"]
+  {{- end }}
+{{- end }}
diff --git a/cockroachdb/templates/clusterrolebinding.yaml b/cockroachdb/templates/clusterrolebinding.yaml
@@ -1,4 +1,4 @@
-{{- if and .Values.tls.enabled (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager) }}
+{{- if or (and .Values.tls.enabled (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager)) .Values.conf.useHostLocality.enabled }}
 kind: ClusterRoleBinding
 apiVersion: rbac.authorization.k8s.io/v1
 metadata:
@@ -18,6 +18,6 @@ roleRef:
   name: {{ template "cockroachdb.fullname" . }}
 subjects:
   - kind: ServiceAccount
-    name: {{ template "cockroachdb.tls.serviceAccount.name" . }}
+    name: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
     namespace: {{ .Release.Namespace | quote }}
-{{- end }}
+{{- end }}
diff --git a/cockroachdb/templates/job.init.yaml b/cockroachdb/templates/job.init.yaml
@@ -46,7 +46,7 @@ spec:
       {{- end }}
     {{- end }}
     {{- if .Values.tls.enabled }}
-      serviceAccountName: {{ template "cockroachdb.tls.serviceAccount.name" . }}
+      serviceAccountName: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
       initContainers:
         - name: copy-certs
           image: "busybox"

diff --git a/cockroachdb/templates/networkpolicy.yaml b/cockroachdb/templates/networkpolicy.yaml
@@ -2,7 +2,7 @@
 kind: NetworkPolicy
 apiVersion: {{ template "cockroachdb.networkPolicy.apiVersion" . }}
 metadata:
-  name: {{ template "cockroachdb.tls.serviceAccount.name" . }}
+  name: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
   namespace: {{ .Release.Namespace | quote }}
   labels:
     helm.sh/chart: {{ template "cockroachdb.chart" . }}

diff --git a/cockroachdb/templates/rolebinding.yaml b/cockroachdb/templates/rolebinding.yaml
@@ -18,6 +18,6 @@ roleRef:
   name: {{ template "cockroachdb.fullname" . }}
 subjects:
   - kind: ServiceAccount
-    name: {{ template "cockroachdb.tls.serviceAccount.name" . }}
+    name: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
     namespace: {{ .Release.Namespace | quote }}
 {{- end }}
diff --git a/cockroachdb/templates/serviceaccount.yaml b/cockroachdb/templates/serviceaccount.yaml
@@ -1,8 +1,8 @@
-{{- if and .Values.tls.enabled .Values.tls.serviceAccount.create }}
+{{- if and (or .Values.tls.enabled .Values.conf.useHostLocality.enabled) .Values.statefulset.serviceAccount.create }}
 kind: ServiceAccount
 apiVersion: v1
 metadata:
-  name: {{ template "cockroachdb.tls.serviceAccount.name" . }}
+  name: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
   namespace: {{ .Release.Namespace | quote }}
   labels:
     helm.sh/chart: {{ template "cockroachdb.chart" . }}

diff --git a/cockroachdb/templates/statefulset.yaml b/cockroachdb/templates/statefulset.yaml
@@ -50,9 +50,10 @@ spec:
         - name: {{ template "cockroachdb.fullname" . }}.self-signed-certs.registry
       {{- end }}
     {{- end }}
+    {{- if or .Values.tls.enabled .Values.conf.useHostLocality.enabled }}
+      serviceAccountName: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
+    {{- end }}
     {{- if .Values.tls.enabled }}
-      serviceAccountName: {{ template "cockroachdb.tls.serviceAccount.name" . }}
-      {{- if .Values.tls.enabled }}
       initContainers:
         - name: copy-certs
           image: "busybox"
@@ -71,7 +72,21 @@ spec:
               mountPath: /cockroach-certs/
             - name: certs-secret
               mountPath: /certs/
-      {{- end }}
+    {{- end }}
+    {{- if .Values.conf.useHostLocality.enabled }}
+        - name: get-node-info
+          image: {{ .Values.conf.useHostLocality.image.repository }}:{{ .Values.conf.useHostLocality.image.tag }}
+          imagePullPolicy: IfNotPresent
+          env:
+            - name: NODE
+              valueFrom:
+                fieldRef:
+                  fieldPath: spec.nodeName
+            - name: OUTPUT_DIR
+              value: /output
+          volumeMounts:
+            - name: node-data
+              mountPath: /output
     {{- end }}
     {{- if or .Values.statefulset.nodeAffinity .Values.statefulset.podAffinity .Values.statefulset.podAntiAffinity }}
       affinity:
@@ -202,8 +217,12 @@ spec:
               --max-offset={{ . }}
             {{- end }}
               --max-sql-memory={{ index .Values.conf `max-sql-memory` }}
-            {{- with .Values.conf.locality }}
+            {{- if .Values.conf.useHostLocality.enabled }}
+              --locality=$(/node-data/topology.sh){{ with .Values.conf.locality }},{{ . }}{{ end }}
+            {{- else }}
+              {{- with .Values.conf.locality }}
               --locality={{ . }}
+              {{- end }}
             {{- end }}
             {{- with index .Values.conf `sql-audit-dir` }}
               --sql-audit-dir={{ . }}
@@ -257,6 +276,10 @@ spec:
               mountPath: /cockroach/log-config
               readOnly: true
           {{- end }}
+          {{- if .Values.conf.useHostLocality.enabled }}
+            - name: node-data
+              mountPath: /node-data
+          {{- end }}
           livenessProbe:
           {{- if .Values.statefulset.customLivenessProbe }}
             {{ toYaml .Values.statefulset.customLivenessProbe | nindent 12 }}
@@ -339,6 +362,10 @@ spec:
           secret:
             secretName: {{ template "cockroachdb.fullname" . }}-log-config
       {{- end }}
+      {{- if .Values.conf.useHostLocality.enabled }}
+        - name: node-data
+          emptyDir: {}
+      {{- end }}
 {{- if .Values.storage.persistentVolume.enabled }}
   volumeClaimTemplates:
     - metadata:

diff --git a/cockroachdb/values.yaml b/cockroachdb/values.yaml
@@ -122,6 +122,15 @@ conf:
   #   locality: planet=earth,province=manitoba,colo=secondary,power=3
   locality: ""
 
+  # Queries the kubernetes API to get the region, zone, and host of the
+  # instance a given pod is running on and adds that to the front of the
+  # locality key set
+  useHostLocality:
+    enabled: false
+    image:
+      repository: scottcrossen/kube-node-labels
+      tag: 1.1.0
+
   # Run CockroachDB instances in standalone mode with replication disabled
   # (replication factor = 1).
   # Enabling this option makes the following values to be ignored:
@@ -261,6 +270,13 @@ statefulset:
     # initialDelaySeconds: 30
     # periodSeconds: 5
 
+  serviceAccount:
+    # Specifies whether this ServiceAccount should be created.
+    create: true
+    # The name of this ServiceAccount to use.
+    # If not set and `create` is `true`, then a name is auto-generated.
+    name: ""
+
 service:
   ports:
     # You can set a different external and internal gRPC ports and their name.
@@ -423,12 +439,6 @@ init:
 # Whether to run securely using TLS certificates.
 tls:
   enabled: true
-  serviceAccount:
-    # Specifies whether this ServiceAccount should be created.
-    create: true
-    # The name of this ServiceAccount to use.
-    # If not set and `create` is `true`, then a name is auto-generated.
-    name: ""
   certs:
     # Bring your own certs scenario. If provided, tls.init section will be ignored.
     provided: false