---
apiVersion: v1
kind: ConfigMap
metadata:
  name: osd-rebalance-infra-nodes
  namespace: openshift-monitoring
data:
  entrypoint: |-
    #!/bin/bash

    CLUSTER_READY=$(oc get job -n openshift-monitoring osd-cluster-ready -o jsonpath='{.status.succeeded}')

    if [ "$CLUSTER_READY" != "1" ]; then
        echo "INFO: 'osd-cluster-ready' job hasn't succeeded yet, exiting..."
        exit 0
    fi

    echo "INFO: Attempting infra node pod rebalancing..."

    AZ_COUNT=$(oc get nodes -o yaml | grep -E "^\s+failure-domain.beta.kubernetes.io/zone" | sort | uniq | wc -l)

    echo "INFO: Number of AZs ${AZ_COUNT}"

    INFRA_NODE_COUNT=$(oc get nodes -l node-role.kubernetes.io=infra --no-headers | wc -l)

    echo "INFO: Number of INFRA NODES ${INFRA_NODE_COUNT}"

    REBALANCE_PODS=false
    PENDING_PODS=false

    rebalancePods() {
        APP=$1
        NS=$2
        LABEL=$3
        VOLUME_NAME=$4
        ALL_POD_NODES=$(oc get pods -n $NS -l $LABEL=$APP -o wide --no-headers | awk '{print $7}' | sort | wc -l)
        UNIQUE_POD_NODES=$(oc get pods -n $NS -l $LABEL=$APP -o wide --no-headers | awk '{print $7}' | sort | uniq | wc -l)
        echo "INFO: $APP pods needing rebalance: $(( $UNIQUE_POD_NODES == $INFRA_NODE_COUNT ? 0 : ALL_POD_NODES - UNIQUE_POD_NODES ))"

        LAST_POD_NODE_NAME=""
        for POD in $( oc get pods -n $NS -l $LABEL=$APP -o jsonpath='{.items[*].metadata.name}' ) ; do
            POD_NODE_NAME=$(oc get pod -n $NS $POD -o jsonpath='{.spec.nodeName}')
            if [ $UNIQUE_POD_NODES -ne $INFRA_NODE_COUNT ] && [ "${POD_NODE_NAME}" == "${LAST_POD_NODE_NAME}" ]; then
                if [ ${VOLUME_NAME} ] && [ "${AZ_COUNT}" != "1" ]; then
                    PVC=$(oc get pod -n $NS $POD -o jsonpath='{.spec.volumes[?(@.name=="'$VOLUME_NAME'")].persistentVolumeClaim.claimName}')
                    echo "INFO: Deleting PVC $PVC"
                    oc delete pvc --wait=true -n $NS $PVC
                fi
                echo "INFO: Deleting pod $POD"
                oc delete pod -n $NS $POD --wait=true
                REBALANCE_PODS=true
            fi
            LAST_POD_NODE_NAME=$POD_NODE_NAME
        done
    }

    checkPendingPods() {
        APP=$1
        NS=$2
        LABEL=$3
        for POD in $( oc get pods -n $NS -l $LABEL=$APP -o jsonpath='{.items[*].metadata.name}' ) ; do
            POD_STATUS_PHASE=$(oc get pods -n $NS $POD -o jsonpath='{.status.phase}')
            if [ "${POD_STATUS_PHASE}" == "Pending" ]; then
                echo "INFO: Deleting pod $POD"
                oc delete pod -n $NS $POD --wait=true
                PENDING_PODS=true
            fi
        done
    }

    waitRunningPods() {
        APP=$1
        NS=$2
        LABEL=$3
        for POD in $( oc get pods -n $NS -l $LABEL=$APP -o jsonpath='{.items[*].metadata.name}' ) ; do
            echo "INFO: Waiting for $POD to be Running..."
            while [ ! -z "$(oc get pod -n $NS $POD -o jsonpath='{.status.phase}' 2>/dev/null)" ] && [ "$(oc get pod -n $NS $POD -o jsonpath='{.status.phase}' 2>/dev/null)" != "Running" ];
            do
                sleep 1
            done
        done
    }

    echo "INFO: Rebalancing prometheus pods..."
    rebalancePods prometheus openshift-monitoring app prometheus-data

    echo "INFO: Rebalancing alertmanager pods..."
    rebalancePods alertmanager openshift-monitoring app alertmanager-data

    echo "INFO: Rebalancing splunk-heavy-forwarder pods..."
    rebalancePods splunk-heavy-forwarder openshift-security name

    if $REBALANCE_PODS; then
        echo "INFO: Restarting prometheus operator..."
        OPERATOR_POD=$( oc get pod -n openshift-monitoring -l app.kubernetes.io/name=prometheus-operator -o jsonpath='{.items[*].metadata.name}' )
        oc delete pod -n openshift-monitoring $OPERATOR_POD --wait=true
        sleep 10
    fi

    echo "INFO: Check pending prometheus pods..."
    checkPendingPods prometheus openshift-monitoring app

    echo "INFO: Check pending alertmanager pods..."
    checkPendingPods alertmanager openshift-monitoring app

    echo "INFO: Check pending splunk-heavy-forwarder pods..."
    checkPendingPods splunk-heavy-forwarder openshift-security name

    if $PENDING_PODS; then
        echo "INFO: Restarting prometheus operator..."
        OPERATOR_POD=$( oc get pod -n openshift-monitoring -l app.kubernetes.io/name=prometheus-operator -o jsonpath='{.items[*].metadata.name}' )
        oc delete pod -n openshift-monitoring $OPERATOR_POD --wait=true
        sleep 10
    fi

    echo "INFO: Wait for running prometheus pods..."
    waitRunningPods prometheus openshift-monitoring app

    echo "INFO: Wait for running alertmanager pods..."
    waitRunningPods alertmanager openshift-monitoring app

    echo "INFO: Wait for running splunk-heavy-forwarder pods..."
    waitRunningPods splunk-heavy-forwarder openshift-security name