Apache Pinot Consuming Segment Commit Stall and Consumer Lag Guard

id: pinot-consuming-segment-commit-stall-guard
namespace: company.team

description: |
  Audits Apache Pinot real-time tables for stuck consuming segments and commit protocol stalls,
  classifies consumer partition lag against operational thresholds, safely executes an authorized
  pause-and-resume cycle to recycle wedged consumer threads, and verifies that ingestion resumes.

inputs:
  - id: pinot_jdbc_url
    type: STRING
    displayName: Pinot JDBC URL
    description: JDBC connection string for the Apache Pinot Controller (the Pinot
      JDBC driver resolves and routes broker queries through the controller),
      e.g. jdbc:pinot://host:9000.
    defaults: "jdbc:pinot://pinot-controller:9000"

  - id: controller_url
    type: STRING
    displayName: Pinot Controller REST URL
    description: HTTP base URL of the Apache Pinot Controller managing cluster
      metadata and Helix state.
    defaults: "http://pinot-controller:9000"

  - id: table_name
    type: STRING
    displayName: Real-Time Table Name
    description: Name of the Apache Pinot real-time table (including the _REALTIME suffix).
    defaults: "transcript_REALTIME"

  - id: max_commit_stall_seconds
    type: INT
    displayName: Max Commit Stall Threshold (Seconds)
    description: Maximum allowable time a segment may remain in COMMITTING or
      stalled CONSUMING state before escalating.
    defaults: 600

  - id: max_allowed_consumer_lag_records
    type: INT
    displayName: Max Consumer Lag Threshold
    description: Maximum permissible difference between partition upstream high
      water mark and current consumer offset.
    defaults: 50000

  - id: auto_remediate
    type: BOOL
    displayName: Auto Remediate
    description: When true and dry_run is false, triggers a Controller
      pause-and-resume cycle to recycle stuck consumer threads.
    defaults: false

  - id: dry_run
    type: BOOL
    displayName: Dry Run Mode
    description: When true, simulates remediation without executing state-modifying
      requests against the Pinot Controller.
    defaults: true

tasks:
  - id: probe_broker_metrics
    type: io.kestra.plugin.jdbc.pinot.Query
    description: Queries the Apache Pinot broker to determine total active rows and
      latest ingested event timestamp.
    url: "{{ inputs.pinot_jdbc_url }}"
    username: "{{ secret('PINOT_BROKER_USER') }}"
    password: "{{ secret('PINOT_BROKER_PASSWORD') }}"
    sql: "SELECT count(*) AS total_rows, max(event_timestamp) AS max_event_timestamp
      FROM {{ inputs.table_name }};"
    fetchType: FETCH_ONE
    retry:
      type: constant
      interval: PT5S
      maxAttempts: 3

  - id: probe_consuming_segments
    type: io.kestra.plugin.core.http.Request
    description: Fetches consuming segment runtime metadata and partition offset
      progression from the Pinot Controller.
    uri: "{{ inputs.controller_url }}/tables/{{ inputs.table_name
      }}/consumingSegmentsInfo"
    method: GET
    headers:
      Authorization: "Bearer {{ secret('PINOT_CONTROLLER_AUTH_TOKEN') }}"
      Accept: "application/json"
    retry:
      type: constant
      interval: PT5S
      maxAttempts: 3

  - id: evaluate_segment_health
    type: io.kestra.plugin.scripts.python.Script
    description: Classifies consuming segment state into HEALTHY, STALLED_COMMIT, or
      CRITICAL_LAG using the real consumingSegmentsInfo response shape (a map of
      segment name to per-replica consumer state, last-consumed timestamp, and
      per-partition records lag).
    containerImage: python:3.12-slim
    env:
      SEGMENTS_INFO: "{{ outputs.probe_consuming_segments.body | toJson }}"
    script: |
      import json
      import os
      import time

      segments_info = json.loads(os.environ.get("SEGMENTS_INFO", "{}"))
      max_commit_stall_ms = {{ inputs.max_commit_stall_seconds }} * 1000
      max_lag_records = {{ inputs.max_allowed_consumer_lag_records }}
      now_ms = int(time.time() * 1000)

      stalled_segments = set()
      high_lag_segments = set()
      total_consuming_replicas = 0

      for segment_name, replica_infos in segments_info.items():
          if not replica_infos:
              # Helix ideal state expects this segment to be CONSUMING, but no server
              # reported an active consumer thread for it. This is the same signal
              # Pinot's own ConsumingSegmentInfoReader.getIngestionStatus() treats as
              # UNHEALTHY -- the segment's consumer is gone, a classic wedged-commit symptom.
              stalled_segments.add(segment_name)
              continue

          for replica_info in replica_infos:
              consumer_state = replica_info.get("consumerState")

              if consumer_state == "NOT_CONSUMING":
                  stalled_segments.add(segment_name)
                  continue

              if consumer_state != "CONSUMING":
                  continue

              total_consuming_replicas += 1

              last_consumed_ts = replica_info.get("lastConsumedTimestamp") or 0
              if (now_ms - last_consumed_ts) > max_commit_stall_ms:
                  stalled_segments.add(segment_name)

              records_lag_map = (replica_info.get("partitionOffsetInfo") or {}).get("recordsLagMap") or {}
              for lag_value in records_lag_map.values():
                  try:
                      lag = int(lag_value)
                  except (TypeError, ValueError):
                      continue
                  if lag > max_lag_records:
                      high_lag_segments.add(segment_name)
                      break

      if not segments_info:
          status = "NO_CONSUMING_SEGMENTS"
      elif stalled_segments:
          status = "STALLED_COMMIT"
      elif high_lag_segments:
          status = "CRITICAL_LAG"
      else:
          status = "HEALTHY"

      output_payload = {
          "status": status,
          "stalled_commit_count": len(stalled_segments),
          "high_lag_count": len(high_lag_segments),
          "total_consuming_replicas": total_consuming_replicas,
      }
      print(f"Pinot segment health: status={status}, stalled={sorted(stalled_segments)}, high_lag={sorted(high_lag_segments)}")
      print(f'::{{"outputs": {json.dumps(output_payload)}}}::')

  - id: route_by_operational_state
    type: io.kestra.plugin.core.flow.Switch
    description: Routes execution based on the evaluated health status of real-time
      consuming segments.
    value: "{{ outputs.evaluate_segment_health.vars.status }}"
    cases:
      HEALTHY:
        - id: log_nominal_state
          type: io.kestra.plugin.core.log.Log
          description: Logs confirmation that all real-time partitions are actively
            consuming within operational bounds.
          message: "Apache Pinot table {{ inputs.table_name }} is healthy. Ingested rows:
            {{ outputs.probe_broker_metrics.row.total_rows }}, latest event
            timestamp: {{ outputs.probe_broker_metrics.row.max_event_timestamp
            }}."

      NO_CONSUMING_SEGMENTS:
        - id: log_no_segments_warning
          type: io.kestra.plugin.core.log.Log
          description: Emits advisory log when no consuming segments are returned by the
            controller for the specified table.
          message: "No active consuming segments found for table {{ inputs.table_name }}.
            Verify table existence and streaming ingestion source."

      STALLED_COMMIT:
        - id: alert_stalled_commit
          type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
          description: Notifies operations team of wedged segment commits blocking
            partition ingestion.
          url: "{{ secret('SLACK_WEBHOOK_URL') }}"
          payload: |
            {
              "channel": "#data-platform-alerts",
              "text": "🚨 *Apache Pinot Ingestion Alert: Stalled Consuming Segment Detected*",
              "attachments": [
                {
                  "color": "#E01E5A",
                  "fields": [
                    {"title": "Table", "value": "`{{ inputs.table_name }}`", "short": true},
                    {"title": "Stall Threshold", "value": "{{ inputs.max_commit_stall_seconds }}s", "short": true},
                    {"title": "Broker Active Rows", "value": "{{ outputs.probe_broker_metrics.row.total_rows }}", "short": true},
                    {"title": "Auto-Remediation", "value": "{{ inputs.auto_remediate and not inputs.dry_run ? 'ENABLED (Executing pause/resume)' : 'DISABLED (Dry run)' }}", "short": true}
                  ]
                }
              ]
            }

        - id: check_remediation_gate
          type: io.kestra.plugin.core.flow.If
          description: Evaluates whether automated recovery execution is authorized by
            operational policy flags.
          condition: "{{ inputs.auto_remediate and not inputs.dry_run }}"
          then:
            - id: pause_table_consumption
              type: io.kestra.plugin.core.http.Request
              description: Instructs the Pinot Controller to pause real-time consumption and
                release stuck commit locks.
              uri: "{{ inputs.controller_url }}/tables/{{ inputs.table_name }}/pause"
              method: POST
              headers:
                Authorization: "Bearer {{ secret('PINOT_CONTROLLER_AUTH_TOKEN') }}"

            - id: resume_table_consumption
              type: io.kestra.plugin.core.http.Request
              description: Instructs the Pinot Controller to resume real-time consumption,
                spawning fresh consumer tasks.
              uri: "{{ inputs.controller_url }}/tables/{{ inputs.table_name }}/resume"
              method: POST
              headers:
                Authorization: "Bearer {{ secret('PINOT_CONTROLLER_AUTH_TOKEN') }}"

            - id: verify_broker_ingestion
              type: io.kestra.plugin.jdbc.pinot.Query
              description: Re-queries the broker following recovery to verify that partition
                consumption has resumed.
              url: "{{ inputs.pinot_jdbc_url }}"
              username: "{{ secret('PINOT_BROKER_USER') }}"
              password: "{{ secret('PINOT_BROKER_PASSWORD') }}"
              sql: "SELECT count(*) AS verified_total_rows, max(event_timestamp) AS
                verified_max_timestamp FROM {{ inputs.table_name }};"
              fetchType: FETCH_ONE

            - id: log_remediation_success
              type: io.kestra.plugin.core.log.Log
              description: Logs successful recycling of stalled Pinot real-time consumer tasks
                and post-recovery row counts.
              message: "Successfully executed pause-resume cycle on {{ inputs.table_name }}.
                Post-recovery rows: {{
                outputs.verify_broker_ingestion.row.verified_total_rows }}."
          else:
            - id: log_dry_run_remediation
              type: io.kestra.plugin.core.log.Log
              description: Records diagnostic plan when auto-remediation is disabled or
                operating in dry-run mode.
              message: "Stalled commit detected on {{ inputs.table_name }}, but auto_remediate
                is false or dry_run is true. Skipping controller mutation."

      CRITICAL_LAG:
        - id: alert_critical_lag
          type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
          description: Dispatches high-priority alert when consumer lag exceeds tolerance
            without a commit freeze.
          url: "{{ secret('SLACK_WEBHOOK_URL') }}"
          payload: |
            {
              "channel": "#data-platform-alerts",
              "text": "⚠️ *Apache Pinot Consumer Lag Warning: Partition Backlog Escalation*",
              "attachments": [
                {
                  "color": "#ECB22E",
                  "fields": [
                    {"title": "Table", "value": "`{{ inputs.table_name }}`", "short": true},
                    {"title": "Lag Threshold", "value": "{{ inputs.max_allowed_consumer_lag_records }} records", "short": true},
                    {"title": "Action", "value": "Inspect upstream Kafka/Pulsar topic partition throughput and Pinot server CPU/disk saturation.", "short": false}
                  ]
                }
              ]
            }

triggers:
  - id: hourly_schedule
    type: io.kestra.plugin.core.trigger.Schedule
    description: Evaluates Pinot real-time table consuming segment health once every hour.
    cron: "0 * * * *"
    disabled: true

  - id: webhook_trigger
    type: io.kestra.plugin.core.trigger.Webhook
    description: External webhook endpoint accepting incident alerts from Prometheus
      Alertmanager or Kafka monitoring agents.
    key: "{{ secret('PINOT_GUARD_WEBHOOK_KEY') }}"
    disabled: true

errors:
  - id: alert_flow_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Posts pipeline execution error alert if broker queries or
      controller API requests fail.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "channel": "#data-platform-alerts",
        "text": "❌ *Pinot Consuming Segment Guard Flow Execution Failed*",
        "attachments": [
          {
            "color": "#D32F2F",
            "fields": [
              {"title": "Execution ID", "value": "`{{ execution.id }}`", "short": true},
              {"title": "Failed Task", "value": "`{{ task.id }}`", "short": true},
              {"title": "Flow", "value": "`{{ flow.id }}`", "short": true}
            ]
          }
        ]
      }