Kafka Client Quota Guard

id: kafka-client-quota-guard
namespace: company.team

description: |
  Keep Kafka client-id quotas on the declared policy: describe every client quota on the
  cluster, report policy clients with no quota or a different one, and any client allowed
  more than the ceiling (an ad-hoc notebook allowed 100 MB/s can starve every producer on
  the broker). Alert Slack, and optionally apply the policy values and clamp over-ceiling
  rates with QuotaAlter, then describe the quotas again to verify. Off by default.

inputs:
  - id: bootstrap_servers
    type: STRING
    displayName: Kafka bootstrap servers
    defaults: "localhost:9092"

  - id: policy
    type: JSON
    displayName: Quota policy
    description: Required quotas per client-id, in bytes per second. Keys among
      producer_byte_rate, consumer_byte_rate and request_percentage.
    defaults: >-
      {"etl-loader": {"producer_byte_rate": 5242880, "consumer_byte_rate":
      10485760},
       "billing-service": {"producer_byte_rate": 524288, "consumer_byte_rate": 1048576},
       "analytics-reader": {"consumer_byte_rate": 20971520}}

  - id: ceiling_byte_rate
    type: INT
    displayName: Ceiling (bytes/s)
    description: No client, in the policy or not, may have a producer or consumer
      byte rate above this.
    defaults: 52428800

  - id: auto_remediate
    type: BOOL
    displayName: Auto remediate
    defaults: false

  - id: dry_run
    type: BOOL
    displayName: Dry run
    description: Must be false, together with auto_remediate, for QuotaAlter to run.
    defaults: true

tasks:
  - id: describe
    type: io.kestra.plugin.kafka.QuotaDescribe
    description: Read-only. Every quota entity on the cluster with its configured rates.
    properties:
      bootstrap.servers: "{{ inputs.bootstrap_servers }}"

  - id: findings
    type: io.kestra.plugin.core.output.OutputValues
    description: One finding per client to fix, with the rates to set. Rates already
      right are left out.
    values:
      items: >-
        {{ {'policy': inputs.policy, 'live': outputs.describe.quotas ?? [],
        'ceil': inputs.ceiling_byte_rate} | jq('
           (.live | map(select(.entity["client-id"] != null) | {(.entity["client-id"]): .quotas}) | add // {}) as $live
           | .ceil as $ceil
           | ([.policy | to_entries[] | . as $p | ($live[$p.key] // null) as $cur
               | {client: $p.key, current: $cur, set: ($p.value | with_entries(select(($cur // {})[.key] != .value)))}
               | select(.set | length > 0)
               | . + {reason: (if .current == null then "NO_QUOTA" else "DRIFT" end)}])
           + ([$live | to_entries[] | select(.key as $c | (["producer_byte_rate", "consumer_byte_rate"] | map(. as $r | ($live[$c][$r] // 0) > $ceil) | any))]
              | map(. as $e | {client: .key, current: .value, reason: "OVER_CEILING",
                    set: (["producer_byte_rate", "consumer_byte_rate"] | map(select(($e.value[.] // 0) > $ceil)) | map({(.): $ceil}) | add)}))
           | group_by(.client) | map({client: .[0].client, current: .[0].current, reason: (map(.reason) | unique | join("+")), set: (map(.set) | add)})') | first | toJson }}

  - id: summary
    type: io.kestra.plugin.core.log.Log
    message: >-
      Kafka client quotas on {{ inputs.bootstrap_servers }}: {{
      outputs.describe.quotas | length }} quota entities, {{
      fromJson(outputs.findings.values.items) | length }} clients to fix{{
      fromJson(outputs.findings.values.items) | jq('if length == 0 then "" else
      ": " + (map("\(.client) \(.reason) \(.current // {} | tojson) -> \(.set |
      tojson)") | join("; ")) end') | first }}.

  - id: route
    type: io.kestra.plugin.core.flow.If
    condition: "{{ fromJson(outputs.findings.values.items) | length > 0 }}"
    then:
      - id: alert
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {{ {'text': 'Kafka client quotas off policy on ' ~ inputs.bootstrap_servers ~ ': ' ~ (fromJson(outputs.findings.values.items) | jq('map("\(.client) (\(.reason))") | join(", ")') | first) ~ '.' ~ (inputs.auto_remediate and not inputs.dry_run ? ' Applying the policy now.' : ' No change applied (auto_remediate=' ~ inputs.auto_remediate ~ ', dry_run=' ~ inputs.dry_run ~ ').')} | toJson }}

      - id: remediate_gate
        type: io.kestra.plugin.core.flow.If
        condition: "{{ inputs.auto_remediate and not inputs.dry_run }}"
        then:
          - id: remediate
            type: io.kestra.plugin.core.flow.Loop
            description: One QuotaAlter per client, setting only the rates that are wrong. A
              rate left empty is not touched.
            values: "{{ outputs.findings.values.items }}"
            tasks:
              - id: alter
                type: io.kestra.plugin.kafka.QuotaAlter
                properties:
                  bootstrap.servers: "{{ inputs.bootstrap_servers }}"
                entityClientId: "{{ fromJson(item.value).client }}"
                producerByteRate: "{{ fromJson(item.value).set | jq('.producer_byte_rate') |
                  first }}"
                consumerByteRate: "{{ fromJson(item.value).set | jq('.consumer_byte_rate') |
                  first }}"
                requestPercentage: "{{ fromJson(item.value).set | jq('.request_percentage') |
                  first }}"

          - id: verify
            type: io.kestra.plugin.kafka.QuotaDescribe
            properties:
              bootstrap.servers: "{{ inputs.bootstrap_servers }}"

          - id: verify_result
            type: io.kestra.plugin.core.flow.If
            condition: >-
              {{ {'want': fromJson(outputs.findings.values.items), 'live':
              outputs.verify.quotas} | jq('(.live |
              map(select(.entity["client-id"] != null) |
              {(.entity["client-id"]): .quotas}) | add // {}) as $l
                 | [.want[] | .client as $c | .set | to_entries[] | select(($l[$c][.key] // null) != .value)] | length') | first > 0 }}
            then:
              - id: remediation_failed
                type: io.kestra.plugin.core.execution.Fail
                errorMessage: "QuotaAlter ran, but the re-describe shows quotas that still
                  differ."
            else:
              - id: remediated
                type: io.kestra.plugin.core.log.Log
                message: "Applied and verified on {{ fromJson(outputs.findings.values.items) |
                  length }} clients."
    else:
      - id: compliant
        type: io.kestra.plugin.core.log.Log
        message: "Every policy client has its quotas and no client is above {{
          inputs.ceiling_byte_rate }} bytes/s."

errors:
  - id: alert_flow_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {{ {'text': 'Kafka client quota guard failed on ' ~ inputs.bootstrap_servers ~ ' (execution ' ~ execution.id ~ '). Quota state unknown, check the logs.'} | toJson }}

triggers:
  - id: hourly
    type: io.kestra.plugin.core.trigger.Schedule
    description: Shipped disabled, and always in safe mode.
    cron: "15 * * * *"
    disabled: true
    inputs:
      auto_remediate: false
      dry_run: true