Add optional label-based rate limiting to trigger_workflow action

playbooks/robusta_playbooks/workflow_trigger.py

@@ -15,6 +15,21 @@
       - trigger_workflow:
           workflow_id: "b7f9d2e4-1234-4c56-9abc-0123456789ab"
           api_key: "{{ env.ROBUSTA_PLATFORM_API_KEY }}"
+
+Optionally, rate limit the action by alert label combination. Below, once a
+workflow is triggered for an alert, further alerts matching on BOTH
+``alertname`` and ``pod`` are skipped for ``rate_limit_seconds``
+(default: 900)::
+
+    customPlaybooks:
+    - triggers:
+      - on_prometheus_alert: {}
+      actions:
+      - trigger_workflow:
+          workflow_id: "b7f9d2e4-1234-4c56-9abc-0123456789ab"
+          api_key: "{{ env.ROBUSTA_PLATFORM_API_KEY }}"
+          rate_limit_labels: ["alertname", "pod"]
+          rate_limit_seconds: 3600
 """
 
 import json
@@ -23,7 +38,7 @@
 
 import requests
 from pydantic import SecretStr
-from robusta.api import ActionException, ActionParams, ErrorCodes, PrometheusKubernetesAlert, action
+from robusta.api import ActionException, ActionParams, ErrorCodes, PrometheusKubernetesAlert, RateLimiter, action
 
 
 class TriggerWorkflowParams(ActionParams):
@@ -43,6 +58,14 @@ class TriggerWorkflowParams(ActionParams):
         workflow definition. Set False to always use the workflow's
         configured cluster.
     :var timeout: (optional) (Default: 30) Request timeout in seconds.
+    :var rate_limit_labels: (optional) Alert labels to rate limit by, e.g.
+        ``["alertname", "pod"]``. When set, alerts whose values match on ALL
+        of these labels share one rate limit bucket: after the workflow is
+        triggered once, further alerts with the same label combination are
+        skipped (with a log) until ``rate_limit_seconds`` passes. By default
+        no rate limit is applied.
+    :var rate_limit_seconds: (optional) (Default: 900) The rate limit period,
+        in seconds. Only relevant when ``rate_limit_labels`` is set.
     """
 
     workflow_id: Union[str, List[str]]
@@ -52,6 +75,8 @@ class TriggerWorkflowParams(ActionParams):
     origin: str = "robusta-runner"
     route_to_alert_cluster: bool = True
     timeout: int = 30
+    rate_limit_labels: Optional[List[str]] = None
+    rate_limit_seconds: int = 900
 
 
 def build_workflow_trigger_payload(alert: PrometheusKubernetesAlert) -> dict:
@@ -68,6 +93,28 @@ def build_workflow_trigger_payload(alert: PrometheusKubernetesAlert) -> dict:
     }
 
 
+def _rate_limit_allows(
+    alert: PrometheusKubernetesAlert, params: TriggerWorkflowParams, workflow_ids: List[str]
+) -> bool:
+    """Mark this alert's rate limit bucket and return whether the action may run.
+
+    Alerts share a bucket when they match on ALL of ``params.rate_limit_labels``
+    (a label missing from the alert matches other alerts missing it too). The
+    workflow ids are part of the bucket key, so separate trigger_workflow
+    configurations don't rate limit each other.
+    """
+    label_values = ",".join(f"{label}={alert.alert.labels.get(label, '')}" for label in sorted(params.rate_limit_labels))
+    limiter_id = f"{','.join(sorted(workflow_ids))}|{label_values}"
+    if RateLimiter.mark_and_test("trigger_workflow", limiter_id, params.rate_limit_seconds):
+        return True
+
+    logging.info(
+        f"trigger_workflow: rate limited for alert {alert.alert_name} ({label_values}); "
+        f"skipping workflow(s) {workflow_ids} for {params.rate_limit_seconds} seconds since the last trigger"
+    )
+    return False
+
+
 @action
 def trigger_workflow(alert: PrometheusKubernetesAlert, params: TriggerWorkflowParams):
     """
@@ -80,6 +127,9 @@ def trigger_workflow(alert: PrometheusKubernetesAlert, params: TriggerWorkflowPa
     if not workflow_ids:
         raise ActionException(ErrorCodes.ACTION_UNEXPECTED_ERROR, "trigger_workflow: no workflow_id provided")
 
+    if params.rate_limit_labels and not _rate_limit_allows(alert, params, workflow_ids):
+        return
+
     context = alert.get_context()
     account_id = params.account_id or context.account_id
 

tests/test_workflow_trigger.py

@@ -6,7 +6,7 @@
 
 import pytest
 from pydantic import SecretStr
-from robusta.api import ActionException
+from robusta.api import ActionException, RateLimiter
 from robusta.core.model.events import ExecutionContext
 from robusta.integrations.prometheus.models import PrometheusAlert, PrometheusKubernetesAlert
 
@@ -40,17 +40,25 @@
 }
 
 
-def make_alert() -> PrometheusKubernetesAlert:
+def make_alert(labels: dict = None) -> PrometheusKubernetesAlert:
+    alert_payload = {**NODE_CORDONED_ALERT, "labels": labels or NODE_CORDONED_ALERT["labels"]}
     alert = PrometheusKubernetesAlert(
-        alert=PrometheusAlert(**NODE_CORDONED_ALERT),
-        alert_name=NODE_CORDONED_ALERT["labels"]["alertname"],
-        alert_severity=NODE_CORDONED_ALERT["labels"]["severity"],
+        alert=PrometheusAlert(**alert_payload),
+        alert_name=alert_payload["labels"]["alertname"],
+        alert_severity=alert_payload["labels"].get("severity", "warning"),
         named_sinks=[],
     )
     alert.set_context(ExecutionContext(account_id=ACCOUNT_ID, cluster_name=CLUSTER_NAME))
     return alert
 
 
+@pytest.fixture(autouse=True)
+def clean_rate_limiter():
+    RateLimiter.limiter_map.clear()
+    yield
+    RateLimiter.limiter_map.clear()
+
+
 class _CaptureServer:
     """Minimal HTTP server capturing webhook requests, responding 200."""
 
@@ -164,6 +172,59 @@ def test_trigger_workflow_cluster_routing_opt_out():
     assert "cluster" not in query  # the workflow's configured cluster applies
 
 
+def _rate_limited_params(url: str, **overrides) -> TriggerWorkflowParams:
+    defaults = dict(
+        workflow_id=WORKFLOW_ID,
+        api_key=SecretStr(API_KEY),
+        url=url,
+        rate_limit_labels=["alertname", "node"],
+    )
+    defaults.update(overrides)
+    return TriggerWorkflowParams(**defaults)
+
+
+def test_rate_limit_skips_repeat_alert_with_same_label_combination():
+    with _CaptureServer() as server:
+        trigger_workflow(make_alert(), _rate_limited_params(server.url))
+        trigger_workflow(make_alert(), _rate_limited_params(server.url))
+
+    assert len(server.requests) == 1  # second firing was rate limited and skipped
+
+
+def test_rate_limit_requires_all_labels_to_match():
+    other_node_labels = {**NODE_CORDONED_ALERT["labels"], "node": "ip-10-0-2-42.ec2.internal"}
+    with _CaptureServer() as server:
+        trigger_workflow(make_alert(), _rate_limited_params(server.url))
+        # same alertname but a different node: only one of the two labels matches, so no rate limit
+        trigger_workflow(make_alert(labels=other_node_labels), _rate_limited_params(server.url))
+
+    assert len(server.requests) == 2
+
+
+def test_rate_limit_allows_after_period_expires():
+    params = _rate_limited_params("placeholder", rate_limit_seconds=600)
+    with _CaptureServer() as server:
+        params.url = server.url
+        trigger_workflow(make_alert(), params)
+        # backdate the stored timestamp past the rate limit period
+        for key in RateLimiter.limiter_map:
+            RateLimiter.limiter_map[key] -= 601
+        trigger_workflow(make_alert(), params)
+
+    assert len(server.requests) == 2
+
+
+def test_no_rate_limit_by_default():
+    with _CaptureServer() as server:
+        for _ in range(3):
+            trigger_workflow(
+                make_alert(),
+                TriggerWorkflowParams(workflow_id=WORKFLOW_ID, api_key=SecretStr(API_KEY), url=server.url),
+            )
+
+    assert len(server.requests) == 3
+
+
 def test_trigger_workflow_raises_on_http_error():
     with _CaptureServer(status_code=401) as server:
         with pytest.raises(ActionException):