Merge pull request #182 from DrDroidLab/prateek/fix/redis_reliability

fix: Improve Redis connection reliability and Kubernetes deployment s…

Author: dimittal

agent/settings.py

@@ -114,6 +114,20 @@ def load_yaml(filepath, native_k8s_connector_mode=False):
 # Celery Configuration Options
 CELERY_BROKER_URL = env.str('CELERY_BROKER_URL', default='redis://localhost:6379/0')
 CELERY_RESULT_BACKEND = env.str('CELERY_RESULT_BACKEND', default='redis://localhost:6379/0')
+
+# Broker connection reliability settings
+CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP = env.bool('CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP', default=True)
+CELERY_BROKER_CONNECTION_MAX_RETRIES = env.int('CELERY_BROKER_CONNECTION_MAX_RETRIES', default=10)
+CELERY_BROKER_CONNECTION_RETRY = True
+CELERY_BROKER_POOL_LIMIT = 10
+CELERY_RESULT_BACKEND_TRANSPORT_OPTIONS = {
+    'master_name': 'mymaster',
+    'socket_keepalive': True,
+    'socket_connect_timeout': 5,
+    'socket_timeout': 5,
+    'retry_on_timeout': True,
+    'health_check_interval': 30,
+}
 CELERY_ACCEPT_CONTENT = ['application/json']
 CELERY_RESULT_SERIALIZER = 'json'
 CELERY_TASK_SERIALIZER = 'json'

agent/tasks.py

@@ -1,4 +1,5 @@
 import logging
+import threading
 
 import requests
 from celery import shared_task
@@ -13,19 +14,59 @@
 K8S_API_TIMEOUT = 30
 HTTP_REQUEST_TIMEOUT = 30
 
+# Kubernetes client cache (singleton pattern)
+_k8s_clients_cache = {}
+_k8s_config_lock = threading.Lock()
+_k8s_config_loaded = False
 
-def get_pod_details(namespace='drdroid'):
-    # Import kubernetes here to avoid loading at module import time
-    from kubernetes import client, config
 
-    try:
-        # Try to load in-cluster config first, then local config
-        try:
-            config.load_incluster_config()
-        except config.ConfigException:
-            config.load_kube_config()
+def _get_k8s_clients():
+    """
+    Get or create singleton Kubernetes API clients.
+    Reuses the same client instances to prevent connection leaks.
+    Thread-safe for Celery workers.
+    """
+    global _k8s_config_loaded, _k8s_clients_cache
+
+    # Fast path: clients already initialized
+    if _k8s_clients_cache:
+        return _k8s_clients_cache['v1'], _k8s_clients_cache['custom']
+
+    # Slow path: initialize clients (thread-safe)
+    with _k8s_config_lock:
+        # Double-check after acquiring lock
+        if _k8s_clients_cache:
+            return _k8s_clients_cache['v1'], _k8s_clients_cache['custom']
+
+        from kubernetes import client, config
+
+        # Load Kubernetes config only once
+        if not _k8s_config_loaded:
+            try:
+                config.load_incluster_config()
+                logger.info("Loaded in-cluster Kubernetes config")
+            except config.ConfigException:
+                config.load_kube_config()
+                logger.info("Loaded local Kubernetes config")
+            _k8s_config_loaded = True
 
-        v1 = client.CoreV1Api()
+        # Create singleton API clients
+        _k8s_clients_cache['v1'] = client.CoreV1Api()
+        _k8s_clients_cache['custom'] = client.CustomObjectsApi()
+
+        logger.info("Initialized Kubernetes API clients (singleton)")
+
+    return _k8s_clients_cache['v1'], _k8s_clients_cache['custom']
+
+
+def get_pod_details(namespace='drdroid'):
+    """
+    Fetch pod details from Kubernetes API.
+    Uses cached API clients to prevent connection leaks.
+    """
+    try:
+        # Get singleton clients (reuses existing instances)
+        v1, custom_api = _get_k8s_clients()
 
         # Get all pods in the namespace with timeout
         pods = v1.list_namespaced_pod(namespace, _request_timeout=K8S_API_TIMEOUT)
@@ -34,7 +75,6 @@ def get_pod_details(namespace='drdroid'):
         # Get pod metrics for CPU/memory usage
         pod_metrics = {}
         try:
-            custom_api = client.CustomObjectsApi()
             metrics = custom_api.list_namespaced_custom_object(
                 group="metrics.k8s.io",
                 version="v1beta1",

helm/charts/celery_beat/templates/deployment.yaml

@@ -25,6 +25,26 @@ spec:
       tolerations:
         {{- toYaml .Values.global.tolerations | nindent 8 }}
       {{- end }}
+      initContainers:
+        - name: wait-for-redis
+          image: busybox:1.36
+          command:
+            - sh
+            - -c
+            - |
+              echo "Waiting for Redis to be ready..."
+              until nc -z redis-service 6379; do
+                echo "Redis is unavailable - sleeping"
+                sleep 2
+              done
+              echo "Redis is ready!"
+          resources:
+            requests:
+              cpu: "10m"
+              memory: "16Mi"
+            limits:
+              cpu: "50m"
+              memory: "32Mi"
       containers:
         - name: celery-beat
           image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
@@ -45,6 +65,10 @@ spec:
               value: {{ .Values.global.DRD_CLOUD_API_HOST | quote }}
             - name: NATIVE_KUBERNETES_API_MODE
               value: {{ .Values.global.NATIVE_KUBERNETES_API_MODE | quote }}
+            - name: CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP
+              value: "true"
+            - name: CELERY_BROKER_CONNECTION_MAX_RETRIES
+              value: "10"
           volumeMounts:
             - name: credentials-volume
               mountPath: /code/credentials/secrets.yaml
@@ -56,6 +80,26 @@ spec:
             limits:
               cpu: "500m"
               memory: "512Mi"
+          livenessProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "test -f /code/celerybeat.pid && ps -p $(cat /code/celerybeat.pid) > /dev/null"
+            initialDelaySeconds: 30
+            periodSeconds: 30
+            timeoutSeconds: 5
+            failureThreshold: 3
+          startupProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "test -f /code/celerybeat.pid"
+            initialDelaySeconds: 15
+            periodSeconds: 5
+            timeoutSeconds: 3
+            failureThreshold: 12
       volumes:
         - name: credentials-volume
           configMap:

helm/charts/celery_worker/templates/deployment.yaml

@@ -25,6 +25,26 @@ spec:
       tolerations:
         {{- toYaml .Values.global.tolerations | nindent 8 }}
       {{- end }}
+      initContainers:
+        - name: wait-for-redis
+          image: busybox:1.36
+          command:
+            - sh
+            - -c
+            - |
+              echo "Waiting for Redis to be ready..."
+              until nc -z redis-service 6379; do
+                echo "Redis is unavailable - sleeping"
+                sleep 2
+              done
+              echo "Redis is ready!"
+          resources:
+            requests:
+              cpu: "10m"
+              memory: "16Mi"
+            limits:
+              cpu: "50m"
+              memory: "32Mi"
       containers:
         - name: celery-worker-scheduler # Lightweight task scheduler
           image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
@@ -49,6 +69,10 @@ spec:
               value: "celery"
             - name: CELERY_WORKER_COUNT
               value: "4"
+            - name: CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP
+              value: "true"
+            - name: CELERY_BROKER_CONNECTION_MAX_RETRIES
+              value: "10"
           volumeMounts:
             - name: credentials-volume
               mountPath: /code/credentials/secrets.yaml
@@ -60,6 +84,26 @@ spec:
             limits:
               cpu: {{ .Values.resources.scheduler.limits.cpu | quote }}
               memory: {{ .Values.resources.scheduler.limits.memory | quote }}
+          startupProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "celery -A agent inspect ping -d celery@$HOSTNAME -t 5"
+            initialDelaySeconds: 30
+            periodSeconds: 10
+            timeoutSeconds: 5
+            failureThreshold: 12
+          livenessProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "celery -A agent inspect ping -d celery@$HOSTNAME -t 5"
+            initialDelaySeconds: 30
+            periodSeconds: 30
+            timeoutSeconds: 10
+            failureThreshold: 3
 
         - name: celery-worker-task-executor # Task executor for high-priority tasks
           image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
@@ -84,6 +128,10 @@ spec:
               value: "exec"
             - name: CELERY_WORKER_COUNT
               value: "4"
+            - name: CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP
+              value: "true"
+            - name: CELERY_BROKER_CONNECTION_MAX_RETRIES
+              value: "10"
           volumeMounts:
             - name: credentials-volume
               mountPath: /code/credentials/secrets.yaml
@@ -95,6 +143,26 @@ spec:
             limits:
               cpu: {{ .Values.resources.taskExecutor.limits.cpu | quote }}
               memory: {{ .Values.resources.taskExecutor.limits.memory | quote }}
+          startupProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "celery -A agent inspect ping -d celery@$HOSTNAME -t 5"
+            initialDelaySeconds: 30
+            periodSeconds: 10
+            timeoutSeconds: 5
+            failureThreshold: 12
+          livenessProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "celery -A agent inspect ping -d celery@$HOSTNAME -t 5"
+            initialDelaySeconds: 30
+            periodSeconds: 30
+            timeoutSeconds: 10
+            failureThreshold: 3
 
         - name: celery-worker-asset-extractor # Task executor for asset extraction tasks, which run rarely and are long-running
           image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
@@ -119,6 +187,10 @@ spec:
               value: "asset_extraction"
             - name: CELERY_WORKER_COUNT
               value: "3"
+            - name: CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP
+              value: "true"
+            - name: CELERY_BROKER_CONNECTION_MAX_RETRIES
+              value: "10"
           volumeMounts:
             - name: credentials-volume
               mountPath: /code/credentials/secrets.yaml
@@ -130,6 +202,26 @@ spec:
             limits:
               cpu: {{ .Values.resources.assetExtractor.limits.cpu | quote }}
               memory: {{ .Values.resources.assetExtractor.limits.memory | quote }}
+          startupProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "celery -A agent inspect ping -d celery@$HOSTNAME -t 5"
+            initialDelaySeconds: 30
+            periodSeconds: 10
+            timeoutSeconds: 5
+            failureThreshold: 12
+          livenessProbe:
+            exec:
+              command:
+                - /bin/sh
+                - -c
+                - "celery -A agent inspect ping -d celery@$HOSTNAME -t 5"
+            initialDelaySeconds: 30
+            periodSeconds: 30
+            timeoutSeconds: 10
+            failureThreshold: 3
 
       volumes:
         - name: credentials-volume

helm/charts/redis/templates/configmap.yaml

@@ -15,7 +15,7 @@ data:
     timeout 300
     
     # Memory management
-    maxmemory 400mb
+    maxmemory 900mb
     maxmemory-policy allkeys-lru
     
     # Persistence settings

helm/charts/redis/templates/deployment.yaml

@@ -14,6 +14,7 @@ spec:
       labels:
         app: redis
     spec:
+      restartPolicy: Always
       serviceAccountName: drd-vpc-agent
       {{- if .Values.global.nodeSelector }}
       nodeSelector:
@@ -30,11 +31,11 @@ spec:
             - containerPort: 6379
           resources:
             requests:
-              memory: "128Mi"
-              cpu: "100m"
+              memory: "256Mi"
+              cpu: "200m"
             limits:
-              memory: "512Mi"
-              cpu: "500m"
+              memory: "1Gi"
+              cpu: "1000m"
           livenessProbe:
             tcpSocket:
               port: 6379

helm/charts/redis/templates/poddisruptionbudget.yaml

@@ -0,0 +1,11 @@
+apiVersion: policy/v1
+kind: PodDisruptionBudget
+metadata:
+  name: redis-pdb
+  annotations:
+    meta.helm.sh/release-name: "{{ .Release.name }}"
+spec:
+  minAvailable: 1
+  selector:
+    matchLabels:
+      app: redis

scripts/start-celery-beat.sh

@@ -6,4 +6,4 @@ set -o nounset
 rm -f './celerybeat.pid'
 
 python manage.py migrate
-celery -A agent beat -l INFO
+celery -A agent beat -l INFO --pidfile=/code/celerybeat.pid