From 04de28157be577514f3fd8983b9a0074db9797cc Mon Sep 17 00:00:00 2001
From: Gustavo Ciria <ciria@devgrid.co.uk>
Date: Tue, 3 Sep 2024 09:06:22 -0300
Subject: [PATCH] feat(helm): Added support for affinity, readinessProbe and
 custom probes

---
 .../helm/retina/templates/daemonset.yaml      | 22 ++++++++++--
 .../controller/helm/retina/values.yaml        | 34 +++++++++++++++++--
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/deploy/legacy/manifests/controller/helm/retina/templates/daemonset.yaml b/deploy/legacy/manifests/controller/helm/retina/templates/daemonset.yaml
index cc99d56f9f4..af4b4e91ea5 100644
--- a/deploy/legacy/manifests/controller/helm/retina/templates/daemonset.yaml
+++ b/deploy/legacy/manifests/controller/helm/retina/templates/daemonset.yaml
@@ -44,12 +44,24 @@ spec:
             mountPath: /retina/config
       containers:
         - name: {{ include "retina.name" . }} 
+          readinessProbe:
+            httpGet:
+              path: /metrics
+              port: {{ .Values.retinaPort }}
+            initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds | default "30" }}
+            periodSeconds: {{ .Values.readinessProbe.periodSeconds | default "30" }}
+            timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds | default "1" }}
+            failureThreshold: {{ .Values.readinessProbe.failureThreshold | default "3" }}
+            successThreshold: {{ .Values.readinessProbe.successThreshold | default "1" }}
           livenessProbe:
             httpGet:
               path: /metrics
               port: {{ .Values.retinaPort }}
-            initialDelaySeconds: 30
-            periodSeconds: 30
+            initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds | default "30" }}
+            periodSeconds: {{ .Values.livenessProbe.periodSeconds | default "30" }}
+            timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds | default "1" }}
+            failureThreshold: {{ .Values.livenessProbe.failureThreshold | default "3" }}
+            successThreshold: {{ .Values.livenessProbe.successThreshold | default "1" }}
           image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
           imagePullPolicy: {{ .Values.image.pullPolicy }}
           {{- if .Values.daemonset.container.retina.command }}
@@ -129,6 +141,9 @@ spec:
         hostPath: 
           path: /sys/class/infiniband
       {{- end }}
+      {{- if .Values.affinity }}
+      affinity: {{- toYaml .Values.affinity | nindent 8 }}
+      {{- end }}
       nodeSelector:
         kubernetes.io/os: linux
         {{- if .Values.nodeSelector }}
@@ -213,6 +228,9 @@ spec:
               mountPath: {{ $mountPath }}
           {{- end }}
         {{- end }}
+      {{- if .Values.affinity }}
+      affinity: {{- toYaml .Values.affinity | nindent 8 }}
+      {{- end }}
       nodeSelector:
         kubernetes.io/os: windows
         {{- if .Values.nodeSelector }}
diff --git a/deploy/legacy/manifests/controller/helm/retina/values.yaml b/deploy/legacy/manifests/controller/helm/retina/values.yaml
index 4dd1c1734bd..3d7b01183bf 100644
--- a/deploy/legacy/manifests/controller/helm/retina/values.yaml
+++ b/deploy/legacy/manifests/controller/helm/retina/values.yaml
@@ -154,6 +154,37 @@ nodeSelector: {}
 ## Ref: https://kubernetes.io/docs/concepts/configuration/taint-and-toleration/
 ##
 tolerations: []
+## @param affinity [object] Affinity rules for pod assignment
+## Ref: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#affinity-and-anti-affinity
+##
+affinity: []
+
+## @param readinessProbe.initialDelaySeconds [array] Initial delay seconds for readinessProbe
+## @param readinessProbe.periodSeconds [array] Period seconds for readinessProbe
+## @param readinessProbe.timeoutSeconds [array] Timeout seconds for readinessProbe
+## @param readinessProbe.failureThreshold [array] Failure threshold for readinessProbe
+## @param readinessProbe.successThreshold [array] Success threshold for readinessProbe
+## Ref: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-probes/#configure-probes
+##
+readinessProbe: {}
+  # initialDelaySeconds: 30
+  # periodSeconds: 10
+  # timeoutSeconds: 15
+  # failureThreshold: 5
+  # successThreshold: 1
+## @param livenessProbe.initialDelaySeconds [array] Initial delay seconds for livenessProbe
+## @param livenessProbe.periodSeconds [array] Period seconds for livenessProbe
+## @param livenessProbe.timeoutSeconds [array] Timeout seconds for livenessProbe
+## @param livenessProbe.failureThreshold [array] Failure threshold for livenessProbe
+## @param livenessProbe.successThreshold [array] Success threshold for livenessProbe
+## Ref: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-probes/#configure-probes
+##
+livenessProbe: {}
+  # initialDelaySeconds: 30
+  # periodSeconds: 10
+  # timeoutSeconds: 15
+  # failureThreshold: 5
+  # successThreshold: 1
 
 metrics:
   ## Prometheus Service Monitor
@@ -190,5 +221,4 @@ metrics:
     tlsConfig: {}
     ## @param metrics.podMonitor.relabelings [array] Prometheus relabeling rules
     ##
-    relabelings: []
-
+    relabelings: []
\ No newline at end of file