From 954051f6d3f41eaac7e7efab563c8f2875bcc94c Mon Sep 17 00:00:00 2001 From: Alexanderamiri Date: Sat, 31 Oct 2026 13:07:26 +0100 Subject: [PATCH] fix(budget): make budget suspension stick and stop the team's RDS too budget-enforcer now writes /javabin/budget-suspensions/{team}; the ecs-service module plans desired_count=0 while it exists, so the next apply no longer scales a suspended team back up. The enforcer also stops the team's RDS instances, reconciles every 6 hours (re-stops RDS after its 7-day auto-start, re-zeros services) and lifts the suspension when spend is back under 200% or on a manual {"action":"lift"} invoke, restoring the previous desired counts and starting the stopped databases. Also parse the plain-text AWS Budgets SNS message; the old handler only accepted JSON and would have skipped every real notification. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 6 +- docs/cost-management.md | 47 +- docs/lambda-functions.md | 17 +- docs/reusable-modules.md | 2 + .../lambda-src/budget_enforcer/handler.py | 426 ++++++++++++++---- .../lambda-src/tests/test_budget_enforcer.py | 242 ++++++++++ terraform/modules/ecs-service/budget.tf | 21 + terraform/modules/ecs-service/main.tf | 2 +- terraform/modules/ecs-service/variables.tf | 6 + terraform/platform/lambdas/main.tf | 64 +++ 10 files changed, 733 insertions(+), 100 deletions(-) create mode 100644 terraform/lambda-src/tests/test_budget_enforcer.py create mode 100644 terraform/modules/ecs-service/budget.tf diff --git a/CLAUDE.md b/CLAUDE.md index 6e8b747..0e50264 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -170,7 +170,7 @@ terraform/state/ | `team-provisioner` | Syncs Google Groups, GitHub teams, AWS Budgets from registry team YAML; Atlassian groups from `groups/atlassian.yaml` | | `securityhub-summary` | Weekly Security Hub findings summary (Monday 08:00 UTC) | | `password-set` | Self-service password-set for new hero accounts (ALB only: password.javazone.no) | -| `budget-enforcer` | Scales ECS services to zero when team exceeds 200% budget | +| `budget-enforcer` | Suspends a team over 200% of budget: SSM marker, ECS to zero, RDS stopped; reconcile every 6 h re-applies it and lifts it when spend is back under 200% | | `resource-tagger` | EventBridge-triggered, auto-tags created-by + commit on new resources | | `ci-broker` | Verifies the job's GitHub OIDC token (repo from signed claims), checks repo ∈ app.yaml team, vends short-lived team role credentials | | `app-auth-sync` | Creates/deletes Cognito app clients requested by app.yaml `auth:` and `routing.login` (SSM request → client + credentials + allowed groups) | @@ -243,7 +243,8 @@ Scheduled: EventBridge (Create/Run) ──► compliance-reporter (report to Slack, no auto-fix) Hourly ──► override-cleanup (delete stale SSM override tokens) Registry merge ──► team-provisioner (Google/GitHub/Budget/Cognito/Identity Center sync + member provisioning + access group sync) -AWS Budgets (200%) ──► budget-enforcer Lambda ──► ECS scale-to-zero + #javabin-cost-alerts +AWS Budgets (200%) ──► budget-enforcer Lambda ──► suspension marker + ECS scale-to-zero + RDS stop + #javabin-cost-alerts + Every 6 hours ──► budget-enforcer reconcile (keep suspensions in force, lift when spend < 200%) EventBridge (Create/Run) ──► resource-tagger Lambda ──► Tag created-by + commit ``` @@ -272,6 +273,7 @@ All parameters are in `eu-central-1`. Use `--profile javabin --region eu-central | `/javabin/apps/{team}/{app}/network` | String | app.yaml `calls:`/`provides:` request (written by app TF, read by app-peers-sync) | | `/javabin/apps/{team}/{app}/peer-{t}-{a}-{url,scopes}`, `m2m-client-{id,secret}` | String / SecureString | Written by app-peers-sync, read by ECS execution role | | `/javabin/apps/{team}/{app}/peers-status` | String | Written by app-peers-sync: state of the app's calls and consumers | +| `/javabin/budget-suspensions/{team}` | String (JSON) | Written/deleted by budget-enforcer; read by the ecs-service module (plans `desired_count=0` while it exists) | ## GCP Connection diff --git a/docs/cost-management.md b/docs/cost-management.md index e6b4921..160df8c 100644 --- a/docs/cost-management.md +++ b/docs/cost-management.md @@ -18,7 +18,7 @@ Lambda posts to **#javabin-cost-alerts**. | Account budget `javabin-account-monthly` | `terraform/platform/cost-analytics/main.tf` | Forecast for the month > 100%; actual > 100%; actual > 150% | | Cost Anomaly Detection `javabin-anomaly-monitor` | `terraform/platform/monitoring/main.tf` | Per-service anomaly with total impact >= $10 (IMMEDIATE) | | Team budgets `javabin-team-` | `team-provisioner` Lambda (from registry `budget_nok`, default 500 NOK) | Actual > 80% (Slack); actual > 200% triggers `budget-enforcer` | -| `budget-enforcer` | `terraform/lambda-src/budget_enforcer` | Scales the team's ECS services to `desired_count=0` (nothing is deleted) | +| `budget-enforcer` | `terraform/lambda-src/budget_enforcer` | Suspends the team: ECS services to `desired_count=0`, RDS instances stopped (nothing is deleted). See [Budget suspension](#budget-suspension) | | `daily-cost-check` / `cost-report` | `terraform/lambda-src/` | Daily spike check; weekly breakdown every Monday | The account budget limit is `account_budget_usd` in `terraform/platform/main.tf` @@ -35,6 +35,51 @@ A per-team anomaly monitor was considered and skipped: a tag-based `aws_ce_anomaly_monitor` needs every team value listed explicitly and is replaced whenever the list changes, and per-team budgets already cover it. +## Budget suspension + +When a team's actual spend passes 200% of its monthly budget, AWS Budgets +notifies `javabin-budget-enforcement` and `budget-enforcer`: + +1. Writes the marker `/javabin/budget-suspensions/{team}` (JSON: when, which + budget, each service's previous desired count, the RDS instances it stopped). +2. Sets every ECS service tagged `team={team}` in `javabin-platform` to + `desired_count=0`. +3. Stops every RDS instance tagged `team={team}` that is `available` (Aurora + members and read replicas are skipped; they cannot be stopped this way). +4. Posts to #javabin-cost-alerts. + +**The suspension sticks.** The `ecs-service` module reads the marker and plans +`desired_count = 0` while it exists, so the team's next `terraform apply` or +deploy does not scale the services back up (`ecs-deploy` only forces a new +deployment; it never changes the count). Every 6 hours a reconcile run re-zeros +services scaled up by hand and stops the databases again: RDS starts a stopped +instance by itself after 7 days, and reconcile catches that within 6 hours. +Terraform changes to a stopped RDS instance fail until it is started. + +**How a team gets unsuspended** (any of): + +- **Raise the budget.** Raise `budget_nok` in the team's registry YAML. Once + team-provisioner has updated the budget, the next reconcile sees spend under + 200% of the new limit and lifts the suspension (within 6 hours). +- **Wait for the next month.** Budgets reset on the 1st (UTC); the first + reconcile of the month lifts every suspension. +- **Lift it now.** A platform admin invokes the Lambda: + `aws lambda invoke --function-name javabin-budget-enforcer --payload '{"action":"lift","team":""}' --cli-binary-format raw-in-base64-out /dev/stdout`. + AWS Budgets notifies once per threshold per month, so a manual lift holds + until next month even if spend stays above 200%. + +Lifting deletes the marker, sets each service the enforcer scaled back to its +previous desired count (if it is still at 0), starts the RDS instances it +stopped, and posts to #javabin-cost-alerts. Services created during the +suspension come up on their next deploy. If the team's budget cannot be read, +reconcile keeps the suspension rather than lifting it. + +The marker is not tamper-proof: it is untagged so that every team's Terraform +can read it, which also means a team role could delete it. The developer +boundary is ~45 bytes under AWS's 6,144-character limit, so a Deny for the +`budget-suspensions` path does not fit yet. Deleting the marker stops +reconcile from enforcing the suspension and shows in CloudTrail. + ## Attribution: cost-allocation tags `team`, `service`, `environment`, `managed-by` and `repo` are activated as diff --git a/docs/lambda-functions.md b/docs/lambda-functions.md index 8f490d9..a3a9470 100644 --- a/docs/lambda-functions.md +++ b/docs/lambda-functions.md @@ -142,16 +142,21 @@ IdP-synced groups are refused. ## budget-enforcer -**Trigger:** SNS notification from AWS Budgets (200% threshold) -**Purpose:** Scales a team's ECS services to `desired_count=0` when spending exceeds 200% of their monthly budget. Does NOT destroy resources — services can be scaled back up after resolution. +**Triggers:** SNS notification from AWS Budgets (200% threshold, plain-text message); EventBridge every 6 hours with `{"action": "reconcile"}`; manual invoke with `{"action": "lift", "team": ""}` +**Purpose:** Suspends a team that spends more than 200% of its monthly budget, and keeps it suspended until the suspension is lifted. Nothing is destroyed. -**Flow:** Parse budget name (`javabin-team-{team}`) → list ECS services tagged with team → scale to zero → post Slack alert. +**Suspend (SNS):** parse `Budget Name:` (`javabin-team-{team}`) → write marker `/javabin/budget-suspensions/{team}` → scale the team's ECS services to `desired_count=0` (previous counts recorded in the marker) → stop the team's `available` RDS instances (IDs recorded) → Slack. The `ecs-service` module plans `desired_count=0` while the marker exists. -| SSM Parameter | Channel | -|---------------|---------| +**Reconcile (schedule):** for each marker, read the team budget. Spend still above 200% of the limit (or unreadable) → re-zero services and re-stop RDS (RDS restarts stopped instances after 7 days). Otherwise (new month, raised `budget_nok`) → lift. + +**Lift:** delete the marker → set the recorded services (still at 0) back to their previous count → start the recorded RDS instances → Slack. See [cost-management.md](cost-management.md#budget-suspension). + +| SSM Parameter | Use | +|---------------|-----| | `/javabin/slack/platform-cost-alerts-webhook` | #javabin-cost-alerts | +| `/javabin/budget-suspensions/{team}` | Suspension marker (written and deleted here) | -**Environment vars:** `ECS_CLUSTER` (default: `javabin-platform`) +**Environment vars:** `ECS_CLUSTER` (default: `javabin-platform`), `COST_WEBHOOK_PARAM`, `BUDGET_NAME_PREFIX`, `SUSPENSION_PATH`, `AWS_ACCOUNT_ID` ## resource-tagger diff --git a/docs/reusable-modules.md b/docs/reusable-modules.md index 9c36146..1a39086 100644 --- a/docs/reusable-modules.md +++ b/docs/reusable-modules.md @@ -70,6 +70,8 @@ Supports `environment` (map) and `secrets` (map of name => ARN) for container co `service_registry_arns` (list, default `[]`) registers the tasks in Cloud Map (from `app-network`). +**Budget suspension:** while the SSM marker `{budget_suspension_path}/{team}` exists (default path `/javabin/budget-suspensions`, written by the `budget-enforcer` Lambda), the service is planned at `desired_count = 0` whatever `desired_count` says, so an apply does not undo the suspension. See [cost-management.md](cost-management.md#budget-suspension). + **Outputs:** `service_name`, `task_definition_arn`, `log_group_name` ## service-bucket diff --git a/terraform/lambda-src/budget_enforcer/handler.py b/terraform/lambda-src/budget_enforcer/handler.py index f7731d2..a4d2abc 100644 --- a/terraform/lambda-src/budget_enforcer/handler.py +++ b/terraform/lambda-src/budget_enforcer/handler.py @@ -1,15 +1,28 @@ -"""Budget enforcer — scales ECS services to zero when a team exceeds 200% of budget. +"""Budget enforcer — suspends a team that exceeds 200% of its monthly budget. -Trigger: SNS notification from AWS Budgets (200% threshold). -Parses budget name to extract team, finds ECS services tagged with that team, -scales them to desired_count=0, and posts a Slack alert to the cost channel. +Triggers: + - SNS notification from AWS Budgets (200% threshold): suspend the team. + - EventBridge schedule, {"action": "reconcile"}: keep every suspension in + force, and lift the ones whose budget is back under 200% (new month or a + raised budget_nok in the registry). + - Manual invoke, {"action": "lift", "team": ""}: lift early. -Does NOT destroy services — just scales to zero for easy recovery. +A suspension is an SSM marker at {SUSPENSION_PATH}/{team}. While it exists: + - the team's ECS services stay at desired_count=0. The ecs-service module + reads the marker and plans desired_count=0 too, so a terraform apply does + not scale them back up; + - the team's RDS instances are stopped. RDS starts a stopped instance again + after 7 days, so reconcile stops it again. + +Lifting deletes the marker, sets the services back to the desired count they +had before, and starts the RDS instances the enforcer stopped. Nothing is +ever deleted. """ import json import logging import os +import re from datetime import datetime, timezone import boto3 @@ -20,15 +33,89 @@ logger.setLevel(logging.INFO) ecs = boto3.client("ecs") +rds = boto3.client("rds") ssm = boto3.client("ssm") +budgets = boto3.client("budgets") CLUSTER = os.environ.get("ECS_CLUSTER", "javabin-platform") COST_WEBHOOK_PARAM = os.environ["COST_WEBHOOK_PARAM"] BUDGET_NAME_PREFIX = os.environ.get("BUDGET_NAME_PREFIX", "javabin-team-") +SUSPENSION_PATH = os.environ.get("SUSPENSION_PATH", "/javabin/budget-suspensions").rstrip("/") +ACCOUNT_ID = os.environ.get("AWS_ACCOUNT_ID", "") + +# Suspension holds while actual spend is above this share of the budget limit. +ENFORCE_RATIO = 2.0 + +# AWS Budgets SNS notifications are plain text, e.g. "Budget Name: javabin-team-x". +_TEXT_FIELD = re.compile(r"^\s*(Budget Name|Budgeted Amount)\s*:\s*\$?\s*(\S+)\s*$", re.MULTILINE) # --------------------------------------------------------------------------- -# ECS helpers +# Budget notification parsing +# --------------------------------------------------------------------------- +def parse_budget_message(text): + """Return (budget_name, budget_limit) from an AWS Budgets SNS message. + + AWS Budgets sends plain text; JSON ({"budgetName", "budgetLimit"}) is + accepted too for manual tests. + """ + try: + message = json.loads(text) + except (json.JSONDecodeError, TypeError): + message = None + if isinstance(message, dict): + return message.get("budgetName", ""), (message.get("budgetLimit") or {}).get("amount") + + fields = {k: v for k, v in _TEXT_FIELD.findall(text or "")} + limit = fields.get("Budgeted Amount") + return fields.get("Budget Name", ""), limit.replace(",", "") if limit else None + + +# --------------------------------------------------------------------------- +# Suspension marker (SSM) +# --------------------------------------------------------------------------- +def marker_name(team): + return f"{SUSPENSION_PATH}/{team}" + + +def read_marker(team): + try: + resp = ssm.get_parameter(Name=marker_name(team)) + except ssm.exceptions.ParameterNotFound: + return None + try: + return json.loads(resp["Parameter"]["Value"]) + except (json.JSONDecodeError, KeyError): + return {} + + +def write_marker(team, state): + ssm.put_parameter( + Name=marker_name(team), + Value=json.dumps(state, sort_keys=True), + Type="String", + Overwrite=True, + Description=f"Budget suspension for team {team} (budget-enforcer)", + ) + + +def delete_marker(team): + try: + ssm.delete_parameter(Name=marker_name(team)) + except ssm.exceptions.ParameterNotFound: + pass + + +def suspended_teams(): + teams = [] + for page in ssm.get_paginator("get_parameters_by_path").paginate(Path=SUSPENSION_PATH): + for p in page.get("Parameters", []): + teams.append(p["Name"].rsplit("/", 1)[-1]) + return teams + + +# --------------------------------------------------------------------------- +# ECS # --------------------------------------------------------------------------- def list_services_for_team(cluster, team_name): """List ECS services in the cluster that are tagged with the given team.""" @@ -52,36 +139,216 @@ def list_services_for_team(cluster, team_name): return matched -def scale_service_to_zero(cluster, service_name): - """Scale an ECS service to desired_count=0.""" - ecs.update_service( - cluster=cluster, - service=service_name, - desiredCount=0, - ) - logger.info("Scaled %s to 0 in cluster %s", service_name, cluster) +def set_desired_count(cluster, service_name, count): + ecs.update_service(cluster=cluster, service=service_name, desiredCount=count) + logger.info("Set %s to desired_count=%d in cluster %s", service_name, count, cluster) + + +def scale_team_to_zero(team, previous): + """Scale the team's running services to 0. Records each service's previous + desired count in `previous` (first value wins). Returns the names scaled.""" + scaled = [] + for svc in list_services_for_team(CLUSTER, team): + name, count = svc["serviceName"], svc.get("desiredCount", 0) + if count == 0: + continue + try: + set_desired_count(CLUSTER, name, 0) + except Exception as e: + logger.error("Failed to scale service %s: %s", name, e) + continue + previous.setdefault(name, count) + scaled.append(name) + return scaled + + +def restore_team_services(team, previous): + """Set services the enforcer scaled back to their previous count, if still at 0.""" + restored = [] + for svc in list_services_for_team(CLUSTER, team): + name = svc["serviceName"] + count = previous.get(name) + if not count or svc.get("desiredCount", 0) != 0: + continue + try: + set_desired_count(CLUSTER, name, count) + restored.append(name) + except Exception as e: + logger.error("Failed to restore service %s: %s", name, e) + return restored + + +# --------------------------------------------------------------------------- +# RDS +# --------------------------------------------------------------------------- +def list_db_instances_for_team(team_name): + matched = [] + for page in rds.get_paginator("describe_db_instances").paginate(): + for db in page.get("DBInstances", []): + tags = {t["Key"]: t["Value"] for t in db.get("TagList", [])} + if tags.get("team") == team_name: + matched.append(db) + return matched + + +def _stoppable(db): + # Aurora members are stopped per cluster; replicas and their sources can't be stopped. + return not (db.get("DBClusterIdentifier") + or db.get("ReadReplicaSourceDBInstanceIdentifier") + or db.get("ReadReplicaDBInstanceIdentifiers")) + + +def stop_team_databases(team, stopped): + """Stop the team's available RDS instances. Adds their IDs to `stopped`.""" + newly = [] + for db in list_db_instances_for_team(team): + ident = db["DBInstanceIdentifier"] + if db.get("DBInstanceStatus") != "available": + continue + if not _stoppable(db): + logger.info("RDS instance %s cannot be stopped (cluster or replica), skipping", ident) + continue + try: + rds.stop_db_instance(DBInstanceIdentifier=ident) + except Exception as e: + logger.error("Failed to stop RDS instance %s: %s", ident, e) + continue + logger.info("Stopped RDS instance %s", ident) + if ident not in stopped: + stopped.append(ident) + newly.append(ident) + return newly + + +def start_team_databases(stopped): + started = [] + for ident in stopped: + try: + rds.start_db_instance(DBInstanceIdentifier=ident) + started.append(ident) + logger.info("Started RDS instance %s", ident) + except rds.exceptions.InvalidDBInstanceStateFault: + logger.info("RDS instance %s is not stopped, leaving it", ident) + except Exception as e: + logger.error("Failed to start RDS instance %s: %s", ident, e) + return started + + +# --------------------------------------------------------------------------- +# Budget state +# --------------------------------------------------------------------------- +def over_limit(team): + """True while the team's actual spend this month is above 200% of its limit. + + Unknown (no budget, API error) counts as still over, so a suspension is + never lifted by accident. + """ + try: + b = budgets.describe_budget(AccountId=ACCOUNT_ID, BudgetName=f"{BUDGET_NAME_PREFIX}{team}")["Budget"] + limit = float(b["BudgetLimit"]["Amount"]) + actual = float(b["CalculatedSpend"]["ActualSpend"]["Amount"]) + except Exception as e: + logger.error("Could not read budget for team %s: %s", team, e) + return True + return actual > limit * ENFORCE_RATIO + + +# --------------------------------------------------------------------------- +# Actions +# --------------------------------------------------------------------------- +def suspend(team, budget_name, budget_limit): + state = read_marker(team) or {} + state.setdefault("suspended_at", datetime.now(timezone.utc).isoformat(timespec="seconds")) + state["budget"] = budget_name + services = state.setdefault("services", {}) + stopped = state.setdefault("rds", []) + # Marker first: from here on a terraform apply plans desired_count=0. + write_marker(team, state) + + scaled = scale_team_to_zero(team, services) + dbs = stop_team_databases(team, stopped) + write_marker(team, state) + + notify(build_suspend_blocks(team, budget_name, scaled, dbs, budget_limit), + f"Budget enforcement: team {team} exceeded 200% — {len(scaled)} service(s) " + f"scaled to zero, {len(dbs)} database(s) stopped") + return {"team": team, "action": "suspended", "services": scaled, "rds": dbs} + + +def lift(team, reason): + state = read_marker(team) + if state is None: + return {"team": team, "action": "not_suspended"} + # Marker first: the restored count matches what terraform now plans. + delete_marker(team) + restored = restore_team_services(team, state.get("services") or {}) + started = start_team_databases(state.get("rds") or []) + notify(build_lift_blocks(team, reason, restored, started), + f"Budget suspension lifted for team {team} ({reason})") + return {"team": team, "action": "lifted", "services": restored, "rds": started} + + +def reapply(team): + """Keep a suspension in force: re-zero services and re-stop databases.""" + state = read_marker(team) or {} + services = state.setdefault("services", {}) + stopped = state.setdefault("rds", []) + scaled = scale_team_to_zero(team, services) + dbs = stop_team_databases(team, stopped) + if scaled or dbs: + write_marker(team, state) + logger.info("Re-applied suspension for %s: services=%s rds=%s", team, scaled, dbs) + return {"team": team, "action": "enforced", "services": scaled, "rds": dbs} + + +def reconcile(): + results = [] + for team in suspended_teams(): + try: + if over_limit(team): + results.append(reapply(team)) + else: + results.append(lift(team, "spend is back under 200% of budget")) + except Exception as e: + logger.error("Reconcile failed for team %s: %s", team, e) + results.append({"team": team, "action": "error", "error": str(e)[:200]}) + return results # --------------------------------------------------------------------------- # Slack notification # --------------------------------------------------------------------------- -def build_alert_blocks(team_name, budget_name, services_scaled, budget_limit): - """Build Slack Block Kit message for budget enforcement action.""" +def notify(blocks, fallback): + try: + webhook_url = get_webhook_url(ssm, COST_WEBHOOK_PARAM) + post_to_slack(webhook_url, {"blocks": blocks, "text": fallback}) + except Exception as e: + logger.error("Failed to send Slack notification: %s", e) + + +def _bullets(items, empty): + return "\n".join(f" - `{i}`" for i in items) if items else f" _{empty}_" + + +def _context(): ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + return {"type": "context", "elements": [{"type": "mrkdwn", "text": f"{ts} | Cluster: `{CLUSTER}`"}]} - service_list = "\n".join( - f" - `{svc}`" for svc in services_scaled - ) if services_scaled else " _(no matching services found)_" - nok_limit = float(budget_limit) * USD_TO_NOK if budget_limit else None +def build_suspend_blocks(team_name, budget_name, services_scaled, dbs_stopped, budget_limit): + """Build Slack Block Kit message for budget enforcement action.""" + try: + nok_limit = float(budget_limit) * USD_TO_NOK if budget_limit else None + except ValueError: + nok_limit = None limit_text = f"${budget_limit}" + (f" (~{nok_limit:.0f} NOK)" if nok_limit else "") - blocks = [ + return [ { "type": "header", "text": { "type": "plain_text", - "text": ":rotating_light: Budget Enforcement — Services Scaled to Zero", + "text": ":rotating_light: Budget Enforcement — Team Suspended", "emoji": True, }, }, @@ -92,105 +359,84 @@ def build_alert_blocks(team_name, budget_name, services_scaled, budget_limit): "text": ( f"Team *{team_name}* has exceeded *200%* of their monthly budget " f"(`{budget_name}`, limit: {limit_text}).\n\n" - f"The following ECS services have been scaled to `desired_count=0`:" + "ECS services scaled to `desired_count=0`:\n" + + _bullets(services_scaled, "(none running)") + + "\n\nRDS instances stopped:\n" + + _bullets(dbs_stopped, "(none running)") ), }, }, + {"type": "divider"}, { "type": "section", "text": { "type": "mrkdwn", - "text": service_list, + "text": ( + f":lock: The suspension holds until it is lifted: deploys keep " + f"`desired_count=0` while `{marker_name(team_name)}` exists.\n" + ":arrow_right: *To lift it*, raise `budget_nok` in the registry (lifted within " + "6 hours), wait for the next month, or ask a platform admin to invoke " + "`javabin-budget-enforcer` with `{\"action\": \"lift\", \"team\": \"" + f"{team_name}\"}}`. Services and databases come back on their own." + ), }, }, - {"type": "divider"}, + _context(), + ] + + +def build_lift_blocks(team_name, reason, restored, started): + return [ + { + "type": "header", + "text": {"type": "plain_text", "text": ":white_check_mark: Budget Suspension Lifted", "emoji": True}, + }, { "type": "section", "text": { "type": "mrkdwn", "text": ( - ":arrow_right: *To restore services*, increase the team budget in " - "the registry or manually scale services back up via the AWS console/CLI." + f"Team *{team_name}* is no longer suspended ({reason}).\n\n" + "ECS services restored:\n" + _bullets(restored, "(none)") + + "\n\nRDS instances started:\n" + _bullets(started, "(none)") + + "\n\nServices created during the suspension come up on the next deploy." ), }, }, - { - "type": "context", - "elements": [ - { - "type": "mrkdwn", - "text": f"Enforced at {ts} | Cluster: `{CLUSTER}`", - } - ], - }, + _context(), ] - return blocks # --------------------------------------------------------------------------- # Lambda entry point # --------------------------------------------------------------------------- def handler(event, context): - for record in event.get("Records", []): - sns_message = record.get("Sns", {}).get("Message", "{}") + action = event.get("action") + if action == "reconcile": + return {"statusCode": 200, "results": reconcile()} + if action == "lift": + team = event.get("team", "") + if not team: + return {"statusCode": 400, "body": "team is required"} + return {"statusCode": 200, "results": [lift(team, "lifted manually")]} - try: - message = json.loads(sns_message) - except json.JSONDecodeError: - logger.error("Failed to parse SNS message: %s", sns_message) + results = [] + for record in event.get("Records", []): + sns_message = record.get("Sns", {}).get("Message", "") + budget_name, budget_limit = parse_budget_message(sns_message) + if not budget_name: + logger.error("No budget name in SNS message: %s", sns_message[:500]) continue - - budget_name = message.get("budgetName", "") if not budget_name.startswith(BUDGET_NAME_PREFIX): - logger.info( - "Ignoring budget notification for %s (not a team budget)", budget_name - ) + logger.info("Ignoring budget notification for %s (not a team budget)", budget_name) continue team_name = budget_name[len(BUDGET_NAME_PREFIX) :] - if not team_name: + if not team_name or "/" in team_name: logger.error("Could not extract team name from budget: %s", budget_name) continue - budget_limit = message.get("budgetLimit", {}).get("amount") - - logger.info( - "Budget enforcement triggered for team %s (budget: %s)", - team_name, - budget_name, - ) - - # Find and scale down team's ECS services - services = list_services_for_team(CLUSTER, team_name) - services_scaled = [] - - for svc in services: - svc_name = svc["serviceName"] - current_count = svc.get("desiredCount", 0) - if current_count == 0: - logger.info( - "Service %s already at desired_count=0, skipping", svc_name - ) - continue - try: - scale_service_to_zero(CLUSTER, svc_name) - services_scaled.append(svc_name) - except Exception as e: - logger.error("Failed to scale service %s: %s", svc_name, e) - - # Post Slack notification regardless of whether services were found - blocks = build_alert_blocks( - team_name, budget_name, services_scaled, budget_limit - ) - fallback = ( - f"Budget enforcement: team {team_name} exceeded 200% — " - f"{len(services_scaled)} service(s) scaled to zero" - ) - - try: - webhook_url = get_webhook_url(ssm, COST_WEBHOOK_PARAM) - post_to_slack(webhook_url, {"blocks": blocks, "text": fallback}) - except Exception as e: - logger.error("Failed to send Slack notification: %s", e) + logger.info("Budget enforcement triggered for team %s (budget: %s)", team_name, budget_name) + results.append(suspend(team_name, budget_name, budget_limit)) - return {"statusCode": 200, "body": "OK"} + return {"statusCode": 200, "body": "OK", "results": results} diff --git a/terraform/lambda-src/tests/test_budget_enforcer.py b/terraform/lambda-src/tests/test_budget_enforcer.py new file mode 100644 index 0000000..0160ba5 --- /dev/null +++ b/terraform/lambda-src/tests/test_budget_enforcer.py @@ -0,0 +1,242 @@ +"""Tests for budget_enforcer — run: python3 -m unittest discover -s terraform/lambda-src/tests""" + +import importlib.util +import json +import os +import sys +import types +import unittest + +LAMBDA_SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..") + +PATH = "/javabin/budget-suspensions" + +BUDGET_TEXT = """AWS Budget Notification November 12, 2026 +AWS Account 111111111111 + +Dear AWS Customer, + +You requested that we alert you when the ACTUAL Cost associated with your javabin-team-cupcake budget is greater than $95.24 for the current month. The ACTUAL Cost associated with this budget is $96.10. + +Budget Name: javabin-team-cupcake +Budget Type: Cost +Budgeted Amount: $47.62 +Alert Type: ACTUAL +Alert Threshold: > $95.24 +ACTUAL Amount: $96.10 +""" + + +class ParameterNotFound(Exception): + pass + + +class InvalidDBInstanceStateFault(Exception): + pass + + +class Paginator: + def __init__(self, fn): + self.fn = fn + + def paginate(self, **kw): + return [self.fn(**kw)] + + +class FakeSSM: + exceptions = types.SimpleNamespace(ParameterNotFound=ParameterNotFound) + + def __init__(self): + self.params = {"/x/cost": "https://hooks.example/cost"} + + def get_parameter(self, Name, WithDecryption=False): + if Name not in self.params: + raise ParameterNotFound(Name) + return {"Parameter": {"Name": Name, "Value": self.params[Name]}} + + def put_parameter(self, Name, Value, **kw): + self.params[Name] = Value + + def delete_parameter(self, Name): + if Name not in self.params: + raise ParameterNotFound(Name) + del self.params[Name] + + def get_paginator(self, name): + assert name == "get_parameters_by_path" + return Paginator(lambda Path: {"Parameters": [ + {"Name": n} for n in self.params if n.startswith(Path + "/") and "/" not in n[len(Path) + 1:] + ]}) + + +class FakeECS: + def __init__(self, services): + # name -> (team, desired) + self.services = dict(services) + self.updates = [] + + def get_paginator(self, name): + return Paginator(lambda cluster: {"serviceArns": list(self.services)}) + + def describe_services(self, cluster, services, include): + return {"services": [ + {"serviceName": n, "desiredCount": self.services[n][1], + "tags": [{"key": "team", "value": self.services[n][0]}]} + for n in services + ]} + + def update_service(self, cluster, service, desiredCount): + team, _ = self.services[service] + self.services[service] = (team, desiredCount) + self.updates.append((service, desiredCount)) + + +class FakeRDS: + exceptions = types.SimpleNamespace(InvalidDBInstanceStateFault=InvalidDBInstanceStateFault) + + def __init__(self, dbs): + self.dbs = dbs # list of dicts + self.stopped, self.started = [], [] + + def get_paginator(self, name): + return Paginator(lambda: {"DBInstances": self.dbs}) + + def _db(self, ident): + return next(d for d in self.dbs if d["DBInstanceIdentifier"] == ident) + + def stop_db_instance(self, DBInstanceIdentifier): + self._db(DBInstanceIdentifier)["DBInstanceStatus"] = "stopping" + self.stopped.append(DBInstanceIdentifier) + + def start_db_instance(self, DBInstanceIdentifier): + db = self._db(DBInstanceIdentifier) + if db["DBInstanceStatus"] != "stopped": + raise InvalidDBInstanceStateFault(DBInstanceIdentifier) + db["DBInstanceStatus"] = "starting" + self.started.append(DBInstanceIdentifier) + + +class FakeBudgets: + def __init__(self, limit, actual): + self.limit, self.actual = limit, actual + + def describe_budget(self, AccountId, BudgetName): + if self.limit is None: + raise RuntimeError("no budget") + return {"Budget": {"BudgetLimit": {"Amount": str(self.limit)}, + "CalculatedSpend": {"ActualSpend": {"Amount": str(self.actual)}}}} + + +def db(ident, team, status="available", **extra): + return {"DBInstanceIdentifier": ident, "DBInstanceStatus": status, + "TagList": [{"Key": "team", "Value": team}], **extra} + + +def load(ecs, rds, ssm, budgets): + clients = {"ecs": ecs, "rds": rds, "ssm": ssm, "budgets": budgets} + boto3 = types.ModuleType("boto3") + boto3.client = lambda name, *a, **kw: clients[name] + sys.modules["boto3"] = boto3 + os.environ.update({"COST_WEBHOOK_PARAM": "/x/cost", "SUSPENSION_PATH": PATH, + "AWS_ACCOUNT_ID": "111111111111"}) + if LAMBDA_SRC not in sys.path: + sys.path.insert(0, LAMBDA_SRC) + for name in [m for m in sys.modules if m == "shared" or m.startswith("shared.")]: + del sys.modules[name] + spec = importlib.util.spec_from_file_location( + "budget_enforcer_handler", os.path.join(LAMBDA_SRC, "budget_enforcer", "handler.py")) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + mod.post_to_slack = lambda url, msg: mod._posted.append(msg) + mod._posted = [] + return mod + + +def sns(message): + return {"Records": [{"Sns": {"Message": message}}]} + + +class BudgetEnforcerTest(unittest.TestCase): + def setUp(self): + self.ecs = FakeECS({"cupcake-web": ("cupcake", 2), "cupcake-idle": ("cupcake", 0), + "other-web": ("other", 1)}) + self.rds = FakeRDS([db("cupcake-db", "cupcake"), db("other-db", "other"), + db("cupcake-aurora", "cupcake", DBClusterIdentifier="c")]) + self.ssm = FakeSSM() + self.budgets = FakeBudgets(47.62, 96.10) + self.mod = load(self.ecs, self.rds, self.ssm, self.budgets) + + def marker(self): + return json.loads(self.ssm.params[f"{PATH}/cupcake"]) + + def test_parses_plain_text_budget_notification(self): + self.assertEqual(self.mod.parse_budget_message(BUDGET_TEXT), ("javabin-team-cupcake", "47.62")) + + def test_parses_json_notification(self): + msg = json.dumps({"budgetName": "javabin-team-x", "budgetLimit": {"amount": "5"}}) + self.assertEqual(self.mod.parse_budget_message(msg), ("javabin-team-x", "5")) + + def test_suspend_writes_marker_scales_services_and_stops_team_rds(self): + self.mod.handler(sns(BUDGET_TEXT), None) + + self.assertEqual(self.ecs.services["cupcake-web"], ("cupcake", 0)) + self.assertEqual(self.ecs.services["other-web"], ("other", 1)) + self.assertEqual(self.rds.stopped, ["cupcake-db"]) + state = self.marker() + self.assertEqual(state["services"], {"cupcake-web": 2}) + self.assertEqual(state["rds"], ["cupcake-db"]) + self.assertEqual(len(self.mod._posted), 1) + + def test_ignores_non_team_budget(self): + self.mod.handler(sns(BUDGET_TEXT.replace("javabin-team-cupcake", "javabin-account-monthly")), None) + self.assertEqual(self.ecs.updates, []) + self.assertNotIn(f"{PATH}/cupcake", self.ssm.params) + + def test_reconcile_reapplies_while_over_limit(self): + self.mod.handler(sns(BUDGET_TEXT), None) + # Someone scales up by hand; RDS auto-starts after 7 days. + self.ecs.services["cupcake-web"] = ("cupcake", 3) + self.rds._db("cupcake-db")["DBInstanceStatus"] = "available" + + self.mod.handler({"action": "reconcile"}, None) + + self.assertEqual(self.ecs.services["cupcake-web"], ("cupcake", 0)) + self.assertEqual(self.rds.stopped, ["cupcake-db", "cupcake-db"]) + # The original count is what lift restores, not the manual one. + self.assertEqual(self.marker()["services"], {"cupcake-web": 2}) + + def test_reconcile_lifts_when_spend_back_under_limit(self): + self.mod.handler(sns(BUDGET_TEXT), None) + self.rds._db("cupcake-db")["DBInstanceStatus"] = "stopped" + self.budgets.limit = 100 # budget_nok raised in the registry + + out = self.mod.handler({"action": "reconcile"}, None) + + self.assertEqual(out["results"][0]["action"], "lifted") + self.assertNotIn(f"{PATH}/cupcake", self.ssm.params) + self.assertEqual(self.ecs.services["cupcake-web"], ("cupcake", 2)) + self.assertEqual(self.ecs.services["cupcake-idle"], ("cupcake", 0)) + self.assertEqual(self.rds.started, ["cupcake-db"]) + + def test_reconcile_keeps_suspension_if_budget_unreadable(self): + self.mod.handler(sns(BUDGET_TEXT), None) + self.budgets.limit = None + self.mod.handler({"action": "reconcile"}, None) + self.assertIn(f"{PATH}/cupcake", self.ssm.params) + + def test_manual_lift(self): + self.mod.handler(sns(BUDGET_TEXT), None) + out = self.mod.handler({"action": "lift", "team": "cupcake"}, None) + self.assertEqual(out["results"][0]["action"], "lifted") + self.assertEqual(self.ecs.services["cupcake-web"], ("cupcake", 2)) + # RDS still stopping: start is skipped, not an error. + self.assertEqual(self.rds.started, []) + + def test_lift_without_suspension_is_noop(self): + out = self.mod.handler({"action": "lift", "team": "cupcake"}, None) + self.assertEqual(out["results"][0]["action"], "not_suspended") + self.assertEqual(self.ecs.updates, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/terraform/modules/ecs-service/budget.tf b/terraform/modules/ecs-service/budget.tf new file mode 100644 index 0000000..55b66f1 --- /dev/null +++ b/terraform/modules/ecs-service/budget.tf @@ -0,0 +1,21 @@ +################################################################################ +# Budget suspension +# +# The budget-enforcer Lambda scales a team's services to 0 when it passes 200% +# of its monthly budget and writes {budget_suspension_path}/{team}. Planning +# desired_count=0 while that marker exists keeps the next apply from scaling +# the service back up. The Lambda deletes the marker and restores the count +# when the suspension is lifted (docs/cost-management.md). +# +# By path, not a single-parameter lookup: a missing marker returns an empty +# list instead of failing the plan. +################################################################################ + +data "aws_ssm_parameters_by_path" "budget_suspension" { + path = var.budget_suspension_path + with_decryption = false +} + +locals { + budget_suspended = contains(data.aws_ssm_parameters_by_path.budget_suspension.names, "${trimsuffix(var.budget_suspension_path, "/")}/${var.team}") +} diff --git a/terraform/modules/ecs-service/main.tf b/terraform/modules/ecs-service/main.tf index e5b4416..947630f 100644 --- a/terraform/modules/ecs-service/main.tf +++ b/terraform/modules/ecs-service/main.tf @@ -82,7 +82,7 @@ resource "aws_ecs_service" "this" { name = "${var.team}-${var.name}" cluster = var.cluster_id task_definition = aws_ecs_task_definition.this.arn - desired_count = var.desired_count + desired_count = local.budget_suspended ? 0 : var.desired_count launch_type = "FARGATE" network_configuration { diff --git a/terraform/modules/ecs-service/variables.tf b/terraform/modules/ecs-service/variables.tf index bc7d82e..6d94c8f 100644 --- a/terraform/modules/ecs-service/variables.tf +++ b/terraform/modules/ecs-service/variables.tf @@ -113,3 +113,9 @@ variable "readonly_root_filesystem" { type = bool default = true } + +variable "budget_suspension_path" { + description = "SSM path of budget suspension markers ({path}/{team}, written by the budget-enforcer Lambda). While the team's marker exists the service is planned at desired_count=0." + type = string + default = "/javabin/budget-suspensions" +} diff --git a/terraform/platform/lambdas/main.tf b/terraform/platform/lambdas/main.tf index 1606a02..2f576f9 100644 --- a/terraform/platform/lambdas/main.tf +++ b/terraform/platform/lambdas/main.tf @@ -1106,6 +1106,44 @@ resource "aws_iam_role_policy" "budget_enforcer" { Action = "ecs:UpdateService" Resource = "arn:aws:ecs:${var.region}:${var.aws_account_id}:service/${var.project}-platform/*" }, + { + # Suspension markers: read by the ecs-service module in app Terraform + Sid = "SuspensionMarkers" + Effect = "Allow" + Action = [ + "ssm:GetParameter", + "ssm:GetParametersByPath", + "ssm:PutParameter", + "ssm:DeleteParameter", + ] + Resource = [ + "arn:aws:ssm:${var.region}:${var.aws_account_id}:parameter/${var.project}/budget-suspensions", + "arn:aws:ssm:${var.region}:${var.aws_account_id}:parameter/${var.project}/budget-suspensions/*", + ] + }, + { + Sid = "RDSDescribe" + Effect = "Allow" + Action = "rds:DescribeDBInstances" + Resource = "arn:aws:rds:${var.region}:${var.aws_account_id}:db:*" + }, + { + # Only team-tagged instances; the handler matches the team tag + Sid = "RDSStopStart" + Effect = "Allow" + Action = ["rds:StopDBInstance", "rds:StartDBInstance"] + Resource = "arn:aws:rds:${var.region}:${var.aws_account_id}:db:*" + Condition = { + Null = { "aws:ResourceTag/team" = "false" } + } + }, + { + # Lift a suspension once spend is back under 200% + Sid = "BudgetsRead" + Effect = "Allow" + Action = "budgets:ViewBudget" + Resource = "arn:aws:budgets::${var.aws_account_id}:budget/${var.project}-team-*" + }, ] }) } @@ -1130,6 +1168,8 @@ resource "aws_lambda_function" "budget_enforcer" { ECS_CLUSTER = "${var.project}-platform" COST_WEBHOOK_PARAM = "/javabin/slack/platform-cost-alerts-webhook" BUDGET_NAME_PREFIX = "${var.project}-team-" + SUSPENSION_PATH = "/${var.project}/budget-suspensions" + AWS_ACCOUNT_ID = var.aws_account_id } } } @@ -1148,6 +1188,30 @@ resource "aws_lambda_permission" "budget_enforcer_sns" { source_arn = aws_sns_topic.budget_enforcement.arn } +# Keeps suspensions in force (re-stops RDS, which restarts itself after 7 +# days; re-zeros services scaled up by hand) and lifts them once the team's +# spend is back under 200% (new month or raised budget_nok). +resource "aws_cloudwatch_event_rule" "budget_enforcer_reconcile" { + name = "${var.project}-budget-enforcer-reconcile" + description = "Reconcile team budget suspensions every 6 hours" + schedule_expression = "rate(6 hours)" +} + +resource "aws_cloudwatch_event_target" "budget_enforcer_reconcile" { + rule = aws_cloudwatch_event_rule.budget_enforcer_reconcile.name + target_id = "reconcile-budget-enforcer" + arn = aws_lambda_function.budget_enforcer.arn + input = jsonencode({ action = "reconcile" }) +} + +resource "aws_lambda_permission" "budget_enforcer_reconcile" { + statement_id = "AllowEventBridgeReconcile" + action = "lambda:InvokeFunction" + function_name = aws_lambda_function.budget_enforcer.function_name + principal = "events.amazonaws.com" + source_arn = aws_cloudwatch_event_rule.budget_enforcer_reconcile.arn +} + ################################################################################ # SNS Subscriptions — slack-alert subscribes to both topics ################################################################################