Coverage for cli / capacity / traffic_dial.py: 100.00%
173 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Manual Global Accelerator traffic-dial controls.
3Backs ``gco capacity traffic-dial show|set|clear``. The global stack publishes
4each region's endpoint-group ARN to SSM (``/{project}/endpoint-group-{region}-arn``
5in the global region), which this module uses for discovery so it works with
6any configured accelerator name. Runtime state shares one SSM tree with the
7scheduled controller (``lambda/traffic-dial-controller``):
9- ``/{project}/traffic-dial/state`` — the controller's last run summary.
10- ``/{project}/traffic-dial/override-{region}`` — a manual override recorded
11 by ``set``; the controller never touches an overridden region until
12 ``clear`` removes the parameter.
14Both are runtime-written, so CloudFormation never deletes them;
15``purge_runtime_parameters`` removes the whole tree and is invoked by a
16fully successful ``gco stacks destroy-all`` so a stale override can never
17pin a region in this account's next deployment.
19``set`` applies the dial via ``UpdateEndpointGroup`` carrying *only*
20``TrafficDialPercentage``: the API patches omitted fields, and omitting
21``EndpointConfigurations`` preserves the registered ALB endpoint.
22"""
24from __future__ import annotations
26import json
27import logging
28import re
29from dataclasses import dataclass, field
30from typing import Any
32import boto3
33from botocore.exceptions import ClientError
35from cli.config import GCOConfig, get_config
37logger = logging.getLogger(__name__)
39#: The Global Accelerator control plane is homed in us-west-2 in the
40#: commercial partition (same convention as the GCO Lambdas).
41GA_CONTROL_PLANE_REGION = "us-west-2"
44class TrafficDialError(Exception):
45 """Raised when a traffic-dial operation cannot be performed."""
48@dataclass
49class RegionDialStatus:
50 """One region's dial, endpoint health, and controller/override state."""
52 region: str
53 traffic_dial: int
54 endpoint_health: str
55 override: str | None = None
56 controller_reason: str | None = None
57 healthy_percent: float | None = None
58 endpoint_group_arn: str = ""
59 warnings: list[str] = field(default_factory=list)
62class TrafficDialManager:
63 """Reads and mutates per-region Global Accelerator traffic dials."""
65 def __init__(self, config: GCOConfig | None = None):
66 self.config = config or get_config()
67 self._session = boto3.Session()
69 def _ssm_client(self) -> Any:
70 return self._session.client("ssm", region_name=self.config.global_region)
72 def _ga_client(self) -> Any:
73 return self._session.client("globalaccelerator", region_name=GA_CONTROL_PLANE_REGION)
75 def _endpoint_group_parameter_pattern(self) -> re.Pattern[str]:
76 project = re.escape(self.config.project_name)
77 return re.compile(rf"^/{project}/endpoint-group-(?P<region>[a-z0-9-]+)-arn$")
79 def _override_parameter_name(self, region: str) -> str:
80 return f"/{self.config.project_name}/traffic-dial/override-{region}"
82 def _state_parameter_name(self) -> str:
83 return f"/{self.config.project_name}/traffic-dial/state"
85 def discover_endpoint_groups(self) -> dict[str, str]:
86 """Return ``{region: endpoint_group_arn}`` from the SSM registry."""
87 ssm = self._ssm_client()
88 pattern = self._endpoint_group_parameter_pattern()
89 groups: dict[str, str] = {}
90 token: str | None = None
91 while True:
92 kwargs: dict[str, Any] = {
93 "Path": f"/{self.config.project_name}",
94 "Recursive": False,
95 }
96 if token:
97 kwargs["NextToken"] = token
98 response = ssm.get_parameters_by_path(**kwargs)
99 for parameter in response.get("Parameters", []):
100 match = pattern.match(str(parameter.get("Name", "")))
101 if match:
102 groups[match.group("region")] = str(parameter.get("Value", ""))
103 token = response.get("NextToken")
104 if not token:
105 break
106 if not groups:
107 raise TrafficDialError(
108 "No Global Accelerator endpoint groups found in the SSM registry "
109 f"(searched /{self.config.project_name}/endpoint-group-*-arn in "
110 f"{self.config.global_region}). Traffic dialing requires the "
111 "commercial-partition Global Accelerator topology."
112 )
113 return groups
115 def read_overrides(self) -> dict[str, str]:
116 """Return ``{region: value}`` for every manual override parameter."""
117 ssm = self._ssm_client()
118 prefix = f"/{self.config.project_name}/traffic-dial/"
119 marker = f"{prefix}override-"
120 overrides: dict[str, str] = {}
121 token: str | None = None
122 while True:
123 kwargs: dict[str, Any] = {"Path": prefix, "Recursive": False}
124 if token:
125 kwargs["NextToken"] = token
126 response = ssm.get_parameters_by_path(**kwargs)
127 for parameter in response.get("Parameters", []):
128 name = str(parameter.get("Name", ""))
129 if name.startswith(marker):
130 overrides[name.removeprefix(marker)] = str(parameter.get("Value", ""))
131 token = response.get("NextToken")
132 if not token:
133 break
134 return overrides
136 def read_controller_state(self) -> dict[str, Any] | None:
137 """Return the controller's last run summary, or None when absent."""
138 ssm = self._ssm_client()
139 try:
140 response = ssm.get_parameter(Name=self._state_parameter_name())
141 except ClientError as exc:
142 if exc.response.get("Error", {}).get("Code") == "ParameterNotFound":
143 return None
144 raise
145 try:
146 state = json.loads(str(response["Parameter"]["Value"]))
147 except KeyError, ValueError:
148 logger.warning("Traffic-dial state parameter holds invalid JSON")
149 return None
150 return state if isinstance(state, dict) else None
152 @staticmethod
153 def _summarize_endpoint_health(group: dict[str, Any]) -> str:
154 descriptions = group.get("EndpointDescriptions", [])
155 if not descriptions:
156 return "no endpoints"
157 healthy = sum(1 for endpoint in descriptions if endpoint.get("HealthState") == "HEALTHY")
158 return f"{healthy}/{len(descriptions)} healthy"
160 def get_status(self) -> list[RegionDialStatus]:
161 """Describe every region's dial, endpoint health, and override state."""
162 groups = self.discover_endpoint_groups()
163 overrides = self.read_overrides()
164 state = self.read_controller_state() or {}
165 decisions = {
166 str(decision.get("region")): decision
167 for decision in state.get("decisions", [])
168 if isinstance(decision, dict)
169 }
171 ga = self._ga_client()
172 statuses: list[RegionDialStatus] = []
173 for region in sorted(groups):
174 arn = groups[region]
175 try:
176 group = ga.describe_endpoint_group(EndpointGroupArn=arn).get("EndpointGroup", {})
177 except ClientError as exc:
178 raise TrafficDialError(
179 f"Failed to describe the {region} endpoint group: {exc}"
180 ) from exc
181 decision = decisions.get(region, {})
182 healthy_percent = decision.get("healthy_percent")
183 statuses.append(
184 RegionDialStatus(
185 region=region,
186 traffic_dial=int(round(float(group.get("TrafficDialPercentage", 100.0)))),
187 endpoint_health=self._summarize_endpoint_health(group),
188 override=overrides.get(region),
189 controller_reason=decision.get("reason"),
190 healthy_percent=(
191 float(healthy_percent) if healthy_percent is not None else None
192 ),
193 endpoint_group_arn=arn,
194 )
195 )
196 return statuses
198 def set_dial(self, region: str, percentage: int) -> RegionDialStatus:
199 """Apply a manual dial and record the override the controller honors."""
200 if not isinstance(percentage, int) or isinstance(percentage, bool):
201 raise TrafficDialError(f"Percentage must be an integer, got {percentage!r}")
202 if not 0 <= percentage <= 100:
203 raise TrafficDialError(f"Percentage must be between 0 and 100, got {percentage}")
205 groups = self.discover_endpoint_groups()
206 if region not in groups:
207 raise TrafficDialError(
208 f"No endpoint group registered for region '{region}'. "
209 f"Known regions: {', '.join(sorted(groups))}"
210 )
212 warnings: list[str] = []
213 if percentage < 100:
214 ga = self._ga_client()
215 others_below_100 = True
216 for other_region, other_arn in groups.items():
217 if other_region == region:
218 continue
219 other = ga.describe_endpoint_group(EndpointGroupArn=other_arn).get(
220 "EndpointGroup", {}
221 )
222 if float(other.get("TrafficDialPercentage", 100.0)) >= 100.0:
223 others_below_100 = False
224 break
225 if others_below_100:
226 warnings.append(
227 "Every other region is already dialed below 100; this leaves the "
228 "listener with no fully dialed region to absorb redirected "
229 "traffic — a configuration whose resulting distribution Global "
230 "Accelerator does not document. The scheduled controller never "
231 "creates this state; proceeding because a manual override is "
232 "explicit operator intent."
233 )
235 ga = self._ga_client()
236 try:
237 # Only the dial: UpdateEndpointGroup patches omitted fields, and
238 # omitting EndpointConfigurations preserves the registered ALB.
239 updated = ga.update_endpoint_group(
240 EndpointGroupArn=groups[region],
241 TrafficDialPercentage=float(percentage),
242 ).get("EndpointGroup", {})
243 except ClientError as exc:
244 raise TrafficDialError(f"Failed to update the {region} traffic dial: {exc}") from exc
246 ssm = self._ssm_client()
247 ssm.put_parameter(
248 Name=self._override_parameter_name(region),
249 Value=str(percentage),
250 Type="String",
251 Overwrite=True,
252 Description=(
253 f"Manual traffic-dial override for {region}; the scheduled "
254 "controller skips this region until the override is cleared."
255 ),
256 )
258 return RegionDialStatus(
259 region=region,
260 traffic_dial=int(round(float(updated.get("TrafficDialPercentage", percentage)))),
261 endpoint_health=self._summarize_endpoint_health(updated),
262 override=str(percentage),
263 endpoint_group_arn=groups[region],
264 warnings=warnings,
265 )
267 def purge_runtime_parameters(self) -> list[str]:
268 """Delete every runtime parameter under ``/{project}/traffic-dial``.
270 The controller Lambda writes ``state`` and ``gco capacity
271 traffic-dial set`` writes ``override-{region}`` at runtime, so
272 CloudFormation never owns or deletes them. Left behind after a full
273 teardown, a stale override is the hazard: the controller honors
274 overrides indefinitely, so a later redeployment would silently pin
275 that region until an operator noticed. Returns the deleted names.
276 """
277 ssm = self._ssm_client()
278 prefix = f"/{self.config.project_name}/traffic-dial"
279 names: list[str] = []
280 token: str | None = None
281 while True:
282 kwargs: dict[str, Any] = {"Path": prefix, "Recursive": True}
283 if token:
284 kwargs["NextToken"] = token
285 response = ssm.get_parameters_by_path(**kwargs)
286 names.extend(
287 str(parameter.get("Name", ""))
288 for parameter in response.get("Parameters", [])
289 if parameter.get("Name")
290 )
291 token = response.get("NextToken")
292 if not token:
293 break
294 deleted: list[str] = []
295 # DeleteParameters accepts at most ten names per call.
296 for start in range(0, len(names), 10):
297 response = ssm.delete_parameters(Names=names[start : start + 10])
298 deleted.extend(str(name) for name in response.get("DeletedParameters", []))
299 return sorted(deleted)
301 def clear_override(self, region: str) -> bool:
302 """Remove a manual override; returns whether one existed.
304 The dial itself is left unchanged: with the controller disabled or in
305 monitor mode it keeps the last manual value, and in enforce mode the
306 controller re-converges it from the region's health signal on its
307 next cycle.
308 """
309 ssm = self._ssm_client()
310 try:
311 ssm.delete_parameter(Name=self._override_parameter_name(region))
312 except ClientError as exc:
313 if exc.response.get("Error", {}).get("Code") == "ParameterNotFound":
314 return False
315 raise
316 return True
319def get_traffic_dial_manager(config: GCOConfig | None = None) -> TrafficDialManager:
320 """Get a configured traffic-dial manager instance."""
321 return TrafficDialManager(config)