apiVersion: kyverno.io/v1
kind: Policy
metadata:
  name: enforce-topology-spread
  namespace: rossum
  annotations:
    policies.kyverno.io/title: Spread pods based on zone topology
    policies.kyverno.io/subject: Pod
spec:
  # Depending on the desired outcome for the applications, there could be multiple solutions
  # to the problem, naturally, the most complex and difficult was chosen in order to:
  # 1. Preserve existing topology spread if defined
  # 2. Inject topology spread if not defined
  # 3. Override any settings of topology spread based on zones
  #
  # see https://github.com/kyverno/kyverno/issues/10655 for details why it cannot
  # be achieved using the merge patch
  rules:
    # Check if existing zone topology spread is defined and overwrite it
    - name: enforce-zone-topology-spread-configuration
      match:
        any:
          - resources:
              kinds:
                - Deployment
                - StatefulSet
              selector:
                matchLabels:
                  # Additional validation policies would be needed to ensure the label is present on every resource
                  app.kubernetes.io/name: "*"
              operations:
                - CREATE
                - UPDATE
      mutate:
        foreach:
          - list: "request.object.spec.template.spec.topologySpreadConstraints || []"
            # Use precondition to mutate
            preconditions:
              any:
                - key: "{{ element.topologyKey }}"
                  operator: Equals
                  value: topology.kubernetes.io/zone
            patchesJson6902: |-
              - path: /spec/template/spec/topologySpreadConstraints/{{elementIndex}}
                op: replace
                value:
                  maxSkew: 1
                  topologyKey: topology.kubernetes.io/zone
                  whenUnsatisfiable: ScheduleAnyway
                  labelSelector:
                    matchLabels:
                      app.kubernetes.io/name: '{{request.object.spec.selector.matchLabels."app.kubernetes.io/name"}}'
    # Check if the zone topology spread is not defined and inject it
    - name: inject-zone-topology-spread
      match:
        any:
          - resources:
              kinds:
                - Deployment
                - StatefulSet
              selector:
                matchLabels:
                  # Additional validation policies would be needed to ensure the label is present on every resource
                  app.kubernetes.io/name: "*"
              operations:
                - CREATE
                - UPDATE
      mutate:
        foreach:
          - list: "request.object.spec.template.spec.topologySpreadConstraints || []"
            # Use precondition to mutate
            preconditions:
              any:
                - key: "{{ request.object.spec.template.spec.topologySpreadConstraints[].topologyKey }}"
                  operator: AllNotIn
                  value:
                    - topology.kubernetes.io/zone
            patchesJson6902: |-
              - path: /spec/template/spec/topologySpreadConstraints/0
                op: add
                value:
                  maxSkew: 1
                  topologyKey: topology.kubernetes.io/zone
                  whenUnsatisfiable: ScheduleAnyway
                  labelSelector:
                    matchLabels:
                      app.kubernetes.io/name: '{{request.object.spec.selector.matchLabels."app.kubernetes.io/name"}}'
    # Create topology spread if it does not exist
    - name: create-topology-spread
      match:
        any:
          - resources:
              kinds:
                - Deployment
                - StatefulSet
              selector:
                matchLabels:
                  # Additional validation policies would be needed to ensure the label is present on every resource
                  app.kubernetes.io/name: "*"
              operations:
                - CREATE
                - UPDATE
      mutate:
        patchStrategicMerge:
          spec:
            template:
              spec:
                # Ensure the zone topology spread is present if undefined
                +(topologySpreadConstraints):
                  - maxSkew: 1
                    topologyKey: topology.kubernetes.io/zone
                    # Depending on the workload and requirements, ScheduleAnyway or DoNotSchedule might be chosen
                    whenUnsatisfiable: ScheduleAnyway
                    labelSelector:
                      matchLabels:
                        app.kubernetes.io/name: '{{request.object.spec.selector.matchLabels."app.kubernetes.io/name"}}'