Expose node locality to pod

scottcrossen · scottcrossen · commit b9946d21e3ea · 2021-08-14T18:28:02.000-07:00
diff --git a/cockroachdb/Chart.yaml b/cockroachdb/Chart.yaml
@@ -1,7 +1,7 @@
 apiVersion: v1
 name: cockroachdb
 home: https://www.cockroachlabs.com
-version: 6.0.8
+version: 6.1.8
 appVersion: 21.1.7
 description: CockroachDB is a scalable, survivable, strongly-consistent SQL database.
 icon: https://raw.githubusercontent.com/cockroachdb/cockroach/master/docs/media/cockroach_db.png
diff --git a/cockroachdb/README.md b/cockroachdb/README.md
@@ -300,6 +300,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`                                 |
@@ -345,6 +348,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`                                          |
@@ -368,12 +373,11 @@ For details see the [`values.yaml`](values.yaml) file.
 | `init.labels`                                             | Additional labels of init Job and its Pod                       | `{"app.kubernetes.io/component": "init"}`        |
 | `init.annotations`                                        | Additional labels of the Pod of init Job                        | `{}`                                             |
 | `init.affinity`                                           | [Affinity rules][2] of init Job Pod                             | `{}`                                             |
+| `init.force`                                              | Create the init pod even if the join URL is set                 | `false`                                          |
 | `init.nodeSelector`                                       | Node labels for init Job Pod assignment                         | `{}`                                             |
 | `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
@@ -1,4 +1,4 @@
-{{- if and (eq (len .Values.conf.join) 0) (not (index .Values.conf `single-node`)) }}
+{{- if and (or .Values.init.force (eq (len .Values.conf.join) 0)) (not (index .Values.conf `single-node`)) }}
 kind: Job
 apiVersion: batch/v1
 metadata:
@@ -15,9 +15,9 @@ metadata:
   {{- with .Values.labels }}
     {{- toYaml . | nindent 4 }}
   {{- end }}
-  annotations:
-    helm.sh/hook: post-install,post-upgrade
-    helm.sh/hook-delete-policy: hook-succeeded,before-hook-creation
+  {{- with .Values.init.jobAnnotations }}
+  annotations: {{- toYaml . | nindent 4 }}
+  {{- end }}
 spec:
   template:
     metadata:
@@ -43,7 +43,7 @@ spec:
       {{- end }}
     {{- end }}
     {{- if and .Values.tls.enabled (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager) }}
-      serviceAccountName: {{ template "cockroachdb.tls.serviceAccount.name" . }}
+      serviceAccountName: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
       initContainers:
         # The init-certs container sends a CSR (certificate signing request) to
         # the Kubernetes cluster.
@@ -77,7 +77,7 @@ spec:
               mountPath: /cockroach-certs/
     {{- end }}
     {{- if or .Values.tls.certs.certManager (and .Values.tls.enabled (.Values.tls.certs.provided))}}
-      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,10 +50,12 @@ spec:
         - name: {{ template "cockroachdb.fullname" . }}.init-certs.registry
       {{- end }}
     {{- end }}
+    {{- if or .Values.tls.enabled .Values.conf.useHostLocality.enabled }}
+      serviceAccountName: {{ template "cockroachdb.statefulset.serviceAccount.name" . }}
+    {{- end }}
+      initContainers:
     {{- if .Values.tls.enabled }}
-      serviceAccountName: {{ template "cockroachdb.tls.serviceAccount.name" . }}
       {{- if and (not .Values.tls.certs.provided) (not .Values.tls.certs.certManager) }}
-      initContainers:
         # The init-certs container sends a CSR (certificate signing request) to
         # the Kubernetes cluster.
         # You can see pending requests using:
@@ -109,6 +111,21 @@ spec:
               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:
       {{- with .Values.statefulset.nodeAffinity }}
@@ -239,8 +256,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={{ . }}
@@ -284,6 +305,10 @@ spec:
               mountPath: {{ printf "/etc/cockroach/secrets/%s" . | quote }}
               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 }}
@@ -357,6 +382,10 @@ spec:
           secret:
             secretName: {{ . | quote }}
       {{- 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
@@ -109,6 +109,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:
@@ -248,6 +257,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.
@@ -346,6 +362,11 @@ init:
   labels:
     app.kubernetes.io/component: init
 
+  # Annotations to add to the init job
+  jobAnnotations:
+    helm.sh/hook: post-install,post-upgrade
+    helm.sh/hook-delete-policy: hook-succeeded,before-hook-creation
+
   # Additional annotations to apply to the Pod of this Job.
   annotations: {}
 
@@ -374,16 +395,12 @@ init:
     #   cpu: "10m"
     #   memory: "128Mi"
 
+  # Create the init job even if this has the join URL set
+  force: false
 
 # Whether to run securely using TLS certificates.
 tls:
   enabled: false
-  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
