#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: celeborn-worker
  labels:
    helm.sh/chart: {{ .Chart.Name }}-{{ .Chart.Version }}
    app.kubernetes.io/name: {{ .Chart.Name }}
    app.kubernetes.io/instance: {{ .Release.Name }}
    app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
    app.kubernetes.io/managed-by: {{ .Release.Service }}
    app.kubernetes.io/role: worker
    {{- include "celeborn.labels" . | nindent 4 }}
spec:
  selector:
    matchLabels:
      app.kubernetes.io/name: {{ .Chart.Name }}
      app.kubernetes.io/instance: {{ .Release.Name }}
      app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
      app.kubernetes.io/role: worker
      {{- include "celeborn.selectorLabels" . | nindent 6 }}
  serviceName: celeborn-worker
  replicas: {{ .Values.workerReplicas }}
  template:
    metadata:
      {{- with .Values.podAnnotations }}
      annotations:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      labels:
        app.kubernetes.io/name: {{ .Chart.Name }}
        app.kubernetes.io/instance: {{ .Release.Name }}
        app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
        app.kubernetes.io/role: worker
        app.kubernetes.io/tag: {{ .Values.image.tag }}
        {{- include "celeborn.selectorLabels" . | nindent 8 }}
    spec:
      {{- with .Values.affinity.worker }}
      affinity:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      {{- with .Values.imagePullSecrets }}
      imagePullSecrets:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      {{- with .Values.nodeSelector }}
      nodeSelector:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      {{- with .Values.tolerations }}
      tolerations:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      {{- with .Values.securityContext }}
      securityContext:
        {{- toYaml . | nindent 8 }}
      {{- end }}
    {{- if .Values.volumeUser }}
      {{- $userId := .Values.volumeUser }}
      {{- $userGroup := .Values.volumeGroup }}
      initContainers:  
    {{- $dirs := get .Values.celeborn "celeborn.worker.storage.dirs" | splitList ","}}
    {{- range $index, $dir := $dirs }}
      - name: volume-permission-{{ $index }}
        image: public.ecr.aws/docker/library/busybox
        # grant volume access to hadoop user  
        command: ['sh', '-c', 'mkdir {{ $dir }} ;chown -R {{ $userId}}:{{ $userGroup }} {{ $dir }}']  
        volumeMounts:
        - mountPath: {{ $dir }} 
          name: vol-{{ $index }}
    {{- end }} 
    {{- end }}    
      containers:
      - name: {{ .Chart.Name }}
        image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
        imagePullPolicy: {{ .Values.image.pullPolicy | default "IfNotPresent"}}
        command:
          - "/usr/bin/tini"
          - "--"
          - "/bin/sh"
          - '-c'
          {{- $namespace := .Release.Namespace }}
          - "until {{ range until (.Values.masterReplicas |int) }}nslookup celeborn-master-{{ . }}.celeborn-master-svc.{{ $namespace }}.svc.{{ $.Values.cluster.name }}.local && {{ end }}true; do echo waiting for master; sleep 2; done && exec /opt/celeborn/sbin/start-worker.sh"
        resources:
          {{- toYaml .Values.resources.worker | nindent 12 }}
        ports:
          - containerPort: {{ get .Values.celeborn "celeborn.worker.metrics.prometheus.port" | default 9096 }}
            name: metrics
            protocol: TCP
        volumeMounts:
          {{- if .Values.configmap }}
          - mountPath: /opt/celeborn/conf
            name: {{ include "celeborn.fullname" . }}-volume
            readOnly: true
          {{- end }}
          {{- $dirs := get .Values.celeborn "celeborn.worker.storage.dirs" | splitList "," -}}
          {{- range $index, $dir := $dirs }}
          - mountPath: {{ $dir }} 
            name: vol-{{ $index }}
          {{- end }}
        env:
          {{- range $key, $val := .Values.environments }}
          - name: {{ $key }}
            value: {{ $val | quote }}
          {{- end}}
      terminationGracePeriodSeconds: 30
      volumes:
        - configMap:
            name: {{ .Values.configmap }}
          name: {{ include "celeborn.fullname" . }}-volume
        {{- $vols := get .Values.celeborn "celeborn.volumes.hostpath" | splitList ","}}   
        {{- range $index, $vol := $vols }}
        - hostPath:
            path: {{ $vol }}
            type: DirectoryOrCreate
          name: vol-{{ $index }}
        {{- end }}