440 lines
18 KiB
Python
440 lines
18 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from typing import Any
|
||
|
||
from core.config import (
|
||
DEPLOY_REGION,
|
||
GATEWAY_CLB_LISTENER_IDS,
|
||
GATEWAY_CLB_LOAD_BALANCER_ID,
|
||
GATEWAY_CLB_POLL_SECONDS,
|
||
GATEWAY_CLB_TASK_TIMEOUT_SECONDS,
|
||
TENCENT_SECRET_ID,
|
||
TENCENT_SECRET_KEY,
|
||
utc_now,
|
||
)
|
||
|
||
|
||
def gateway_clb_configuration_error() -> str:
|
||
# gateway 发布一旦要走摘流，就必须拿到完整的腾讯云和 CLB 配置。
|
||
missing: list[str] = []
|
||
if not TENCENT_SECRET_ID:
|
||
missing.append("TENCENT_SECRET_ID")
|
||
if not TENCENT_SECRET_KEY:
|
||
missing.append("TENCENT_SECRET_KEY")
|
||
if not DEPLOY_REGION:
|
||
missing.append("DEPLOY_REGION")
|
||
if not GATEWAY_CLB_LOAD_BALANCER_ID:
|
||
missing.append("GATEWAY_CLB_LOAD_BALANCER_ID")
|
||
if not GATEWAY_CLB_LISTENER_IDS:
|
||
missing.append("GATEWAY_CLB_LISTENER_IDS")
|
||
|
||
if missing:
|
||
return f"gateway CLB protection requires .env: {', '.join(missing)}"
|
||
return ""
|
||
|
||
|
||
def gateway_clb_settings() -> dict[str, Any]:
|
||
# 配置检查单独收口，发布页预警和真正执行都复用同一套口径。
|
||
error = gateway_clb_configuration_error()
|
||
if error:
|
||
raise RuntimeError(error)
|
||
|
||
return {
|
||
"loadBalancerId": GATEWAY_CLB_LOAD_BALANCER_ID,
|
||
"listenerIds": list(GATEWAY_CLB_LISTENER_IDS),
|
||
}
|
||
|
||
|
||
def build_clb_client() -> tuple[Any | None, str]:
|
||
# 先检查最基本的环境变量，缺任何一项都不继续请求腾讯云。
|
||
error = gateway_clb_configuration_error()
|
||
if error:
|
||
return None, error
|
||
|
||
try:
|
||
from tencentcloud.common import credential
|
||
from tencentcloud.common.profile.client_profile import ClientProfile
|
||
from tencentcloud.common.profile.http_profile import HttpProfile
|
||
from tencentcloud.clb.v20180317 import clb_client
|
||
except Exception as exc: # noqa: BLE001
|
||
return None, f"tencentcloud sdk unavailable: {exc}"
|
||
|
||
cred = credential.Credential(TENCENT_SECRET_ID, TENCENT_SECRET_KEY)
|
||
client = clb_client.ClbClient(
|
||
cred,
|
||
DEPLOY_REGION,
|
||
ClientProfile(httpProfile=HttpProfile(endpoint="clb.tencentcloudapi.com")),
|
||
)
|
||
return client, ""
|
||
|
||
|
||
def invoke_clb(action: Any, label: str, retries: int = 3, delay_seconds: int = 2) -> Any:
|
||
# SDK 异常这里集中做简单重试，避免上层每一步都手写同样的退避逻辑。
|
||
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
|
||
|
||
last_error: Exception | None = None
|
||
for attempt in range(1, retries + 1):
|
||
try:
|
||
return action()
|
||
except TencentCloudSDKException as exc:
|
||
last_error = exc
|
||
if attempt == retries:
|
||
break
|
||
time.sleep(delay_seconds * attempt)
|
||
|
||
if last_error is not None:
|
||
raise last_error
|
||
raise RuntimeError(f"{label} failed without error")
|
||
|
||
|
||
def call_clb(method_name: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||
# CLB 调用统一走 json 请求体，后面加字段时不用反复和 SDK 模型属性打交道。
|
||
client, error = build_clb_client()
|
||
if client is None:
|
||
raise RuntimeError(error)
|
||
|
||
from tencentcloud.clb.v20180317 import models as clb_models
|
||
|
||
request_class = getattr(clb_models, f"{method_name}Request")
|
||
request = request_class()
|
||
request.from_json_string(json.dumps(payload))
|
||
response = invoke_clb(lambda: getattr(client, method_name)(request), method_name)
|
||
return dict(json.loads(response.to_json_string()).get("Response") or {})
|
||
|
||
|
||
def wait_clb_task(task_id: str, *, timeout_seconds: int = GATEWAY_CLB_TASK_TIMEOUT_SECONDS) -> dict[str, Any]:
|
||
# 所有异步修改都轮询到终态，避免“接口返回成功但后端任务失败”的假成功。
|
||
if not task_id:
|
||
raise RuntimeError("missing CLB task id")
|
||
|
||
deadline = time.time() + max(float(timeout_seconds), 1.0)
|
||
last_payload: dict[str, Any] = {}
|
||
while time.time() < deadline:
|
||
last_payload = call_clb("DescribeTaskStatus", {"TaskId": task_id})
|
||
status = int(last_payload.get("Status") if last_payload.get("Status") is not None else 2)
|
||
if status == 0:
|
||
return last_payload
|
||
if status == 1:
|
||
message = str(last_payload.get("Message") or "CLB task failed").strip()
|
||
raise RuntimeError(f"CLB task failed: {message}")
|
||
time.sleep(max(float(GATEWAY_CLB_POLL_SECONDS), 0.1))
|
||
|
||
raise RuntimeError(f"CLB task timed out: {task_id}")
|
||
|
||
|
||
def target_matches_service(target: dict[str, Any], service: dict[str, Any]) -> bool:
|
||
# 匹配优先认实例 ID，其次再用私网 IP 和端口兜底，避免误把另一台 gateway 当成当前节点。
|
||
target_instance_id = str(target.get("InstanceId") or target.get("TargetId") or "").strip()
|
||
service_instance_id = str(service.get("instanceId") or "").strip()
|
||
if service_instance_id and target_instance_id and target_instance_id != service_instance_id:
|
||
return False
|
||
|
||
target_port = int(target.get("Port") or 0)
|
||
service_port = int(service.get("port") or 0)
|
||
if service_port and target_port and target_port != service_port:
|
||
return False
|
||
|
||
service_ip = str(service.get("ip") or "").strip()
|
||
target_ips = [str(item or "").strip() for item in (target.get("PrivateIpAddresses") or [target.get("IP")]) if str(item or "").strip()]
|
||
if service_ip and target_ips and service_ip not in target_ips:
|
||
return False
|
||
|
||
return bool(target_instance_id or target_ips)
|
||
|
||
|
||
def binding_key(binding: dict[str, Any]) -> str:
|
||
# 同一实例可能同时挂在多个监听器/规则下，这里把绑定维度完整拼成稳定 key。
|
||
return "|".join(
|
||
[
|
||
str(binding.get("listenerId") or "").strip(),
|
||
str(binding.get("locationId") or "").strip(),
|
||
str(binding.get("instanceId") or "").strip(),
|
||
str(int(binding.get("port") or 0)),
|
||
]
|
||
)
|
||
|
||
|
||
def binding_route_label(binding: dict[str, Any]) -> str:
|
||
# 日志里优先显示规则 ID，没有规则 ID 时退回到监听器 ID。
|
||
parts = [str(binding.get("listenerId") or "").strip()]
|
||
location_id = str(binding.get("locationId") or "").strip()
|
||
if location_id:
|
||
parts.append(location_id)
|
||
domain = str(binding.get("domain") or "").strip()
|
||
url = str(binding.get("url") or "").strip()
|
||
if domain or url:
|
||
parts.append(f"{domain or '*'}{url or '/'}")
|
||
return " / ".join(part for part in parts if part)
|
||
|
||
|
||
def build_binding(listener: dict[str, Any], service: dict[str, Any], target: dict[str, Any], *, rule: dict[str, Any] | None) -> dict[str, Any]:
|
||
# 摘流和恢复都只需要最小必要字段，这里统一整理成稳定结构。
|
||
private_ips = [str(item or "").strip() for item in (target.get("PrivateIpAddresses") or []) if str(item or "").strip()]
|
||
binding = {
|
||
"listenerId": str(listener.get("ListenerId") or "").strip(),
|
||
"locationId": str((rule or {}).get("LocationId") or "").strip(),
|
||
"domain": str((rule or {}).get("Domain") or "").strip(),
|
||
"url": str((rule or {}).get("Url") or "").strip(),
|
||
"instanceId": str(target.get("InstanceId") or target.get("TargetId") or service.get("instanceId") or "").strip(),
|
||
"port": int(target.get("Port") or service.get("port") or 0),
|
||
"weight": int(target.get("Weight") or 0),
|
||
"privateIp": private_ips[0] if private_ips else str(service.get("ip") or "").strip(),
|
||
"routeLabel": "",
|
||
}
|
||
binding["routeLabel"] = binding_route_label(binding)
|
||
return binding
|
||
|
||
|
||
def collect_service_bindings(listeners: list[dict[str, Any]], service: dict[str, Any]) -> list[dict[str, Any]]:
|
||
# 自动扫描当前实例在配置监听器里的全部绑定，避免只摘一条规则导致仍有流量漏进来。
|
||
bindings_by_key: dict[str, dict[str, Any]] = {}
|
||
|
||
for listener in listeners:
|
||
for target in listener.get("Targets") or []:
|
||
if not target_matches_service(target, service):
|
||
continue
|
||
binding = build_binding(listener, service, target, rule=None)
|
||
bindings_by_key[binding_key(binding)] = binding
|
||
|
||
for rule in listener.get("Rules") or []:
|
||
for target in rule.get("Targets") or []:
|
||
if not target_matches_service(target, service):
|
||
continue
|
||
binding = build_binding(listener, service, target, rule=rule)
|
||
bindings_by_key[binding_key(binding)] = binding
|
||
|
||
return list(bindings_by_key.values())
|
||
|
||
|
||
def describe_gateway_targets() -> dict[str, Any]:
|
||
# 直接按配置监听器拉绑定列表，摘流只关心这几个 listener 的后端状态。
|
||
settings = gateway_clb_settings()
|
||
return call_clb(
|
||
"DescribeTargets",
|
||
{
|
||
"LoadBalancerId": settings["loadBalancerId"],
|
||
"ListenerIds": settings["listenerIds"],
|
||
},
|
||
)
|
||
|
||
|
||
def current_gateway_bindings(snapshot: dict[str, Any]) -> list[dict[str, Any]]:
|
||
# 每次修改后都重新读一遍 CLB 当前值，确认线上状态已经真正落地。
|
||
payload = describe_gateway_targets()
|
||
return collect_service_bindings(list(payload.get("Listeners") or []), snapshot["service"])
|
||
|
||
|
||
def render_weight_summary(bindings: list[dict[str, Any]]) -> str:
|
||
# 日志只保留规则和权重摘要，避免一条进度消息塞太多冗余字段。
|
||
parts = [f"{binding['routeLabel']} -> {int(binding.get('weight') or 0)}" for binding in bindings]
|
||
return ", ".join(parts)
|
||
|
||
|
||
def prepare_gateway_release_clb_snapshot(service: dict[str, Any]) -> dict[str, Any]:
|
||
# 进入单节点重启前，先把这个节点当前在 CLB 里的全部绑定关系记下来。
|
||
if str(service.get("service") or "").strip() != "gateway":
|
||
raise RuntimeError("gateway CLB protection only supports gateway service")
|
||
|
||
settings = gateway_clb_settings()
|
||
listeners = list(describe_gateway_targets().get("Listeners") or [])
|
||
bindings = collect_service_bindings(listeners, service)
|
||
if not bindings:
|
||
raise RuntimeError(
|
||
f"gateway target not found in configured CLB listeners: {service.get('host')} {service.get('instanceId') or service.get('ip')}"
|
||
)
|
||
|
||
return {
|
||
"service": {
|
||
"host": str(service.get("host") or "").strip(),
|
||
"service": str(service.get("service") or "").strip(),
|
||
"ip": str(service.get("ip") or "").strip(),
|
||
"instanceId": str(service.get("instanceId") or "").strip(),
|
||
"port": int(service.get("port") or 0),
|
||
},
|
||
"loadBalancerId": settings["loadBalancerId"],
|
||
"listenerIds": list(settings["listenerIds"]),
|
||
"bindings": bindings,
|
||
"bindingSummary": render_weight_summary(bindings),
|
||
"updatedAt": utc_now(),
|
||
}
|
||
|
||
|
||
def modify_gateway_binding_weight(binding: dict[str, Any], weight: int) -> dict[str, Any]:
|
||
# 单条绑定单独改权重，出错时能直接定位到具体监听器/规则。
|
||
payload: dict[str, Any] = {
|
||
"LoadBalancerId": GATEWAY_CLB_LOAD_BALANCER_ID,
|
||
"ListenerId": binding["listenerId"],
|
||
"Targets": [
|
||
{
|
||
"InstanceId": binding["instanceId"],
|
||
"Port": int(binding["port"]),
|
||
}
|
||
],
|
||
"Weight": int(weight),
|
||
}
|
||
if str(binding.get("locationId") or "").strip():
|
||
payload["LocationId"] = binding["locationId"]
|
||
elif str(binding.get("domain") or "").strip() or str(binding.get("url") or "").strip():
|
||
payload["Domain"] = str(binding.get("domain") or "").strip()
|
||
payload["Url"] = str(binding.get("url") or "").strip()
|
||
|
||
response = call_clb("ModifyTargetWeight", payload)
|
||
wait_clb_task(str(response.get("RequestId") or "").strip())
|
||
return {
|
||
"routeLabel": binding["routeLabel"],
|
||
"weight": int(weight),
|
||
"updatedAt": utc_now(),
|
||
}
|
||
|
||
|
||
def wait_gateway_binding_weights(snapshot: dict[str, Any], expected_weights: dict[str, int], *, timeout_seconds: int) -> list[dict[str, Any]]:
|
||
# 权重调整是异步任务，必须读回线上值确认全部命中目标权重后才能继续重启。
|
||
deadline = time.time() + max(float(timeout_seconds), 1.0)
|
||
last_bindings: list[dict[str, Any]] = []
|
||
while time.time() < deadline:
|
||
last_bindings = current_gateway_bindings(snapshot)
|
||
current_by_key = {binding_key(item): item for item in last_bindings}
|
||
if current_by_key and all(
|
||
key in current_by_key and int(current_by_key[key].get("weight") or 0) == int(expected_weights[key])
|
||
for key in expected_weights
|
||
):
|
||
return last_bindings
|
||
time.sleep(max(float(GATEWAY_CLB_POLL_SECONDS), 0.1))
|
||
|
||
if not last_bindings:
|
||
raise RuntimeError("gateway CLB bindings disappeared while waiting for target weight")
|
||
|
||
current_by_key = {binding_key(item): item for item in last_bindings}
|
||
raise RuntimeError(
|
||
"gateway CLB weight not ready: "
|
||
+ ", ".join(
|
||
f"{current_by_key.get(key, {}).get('routeLabel', key)}={int(current_by_key.get(key, {}).get('weight') or -1)}"
|
||
for key in expected_weights
|
||
)
|
||
)
|
||
|
||
|
||
def drain_gateway_release_target(snapshot: dict[str, Any]) -> dict[str, Any]:
|
||
# 摘流直接把当前节点在所有匹配规则下的权重调成 0，避免新请求继续打到待重启节点。
|
||
expected_weights: dict[str, int] = {}
|
||
operations: list[dict[str, Any]] = []
|
||
for binding in snapshot["bindings"]:
|
||
operations.append(modify_gateway_binding_weight(binding, 0))
|
||
expected_weights[binding_key(binding)] = 0
|
||
|
||
verified = wait_gateway_binding_weights(
|
||
snapshot,
|
||
expected_weights,
|
||
timeout_seconds=GATEWAY_CLB_TASK_TIMEOUT_SECONDS,
|
||
)
|
||
return {
|
||
"ok": True,
|
||
"updatedAt": utc_now(),
|
||
"operations": operations,
|
||
"bindings": verified,
|
||
"summary": render_weight_summary(verified),
|
||
}
|
||
|
||
|
||
def restore_gateway_release_target(snapshot: dict[str, Any]) -> dict[str, Any]:
|
||
# 服务恢复后把每条绑定恢复成摘流前的原权重，避免线上流量分配被改坏。
|
||
expected_weights: dict[str, int] = {}
|
||
operations: list[dict[str, Any]] = []
|
||
for binding in snapshot["bindings"]:
|
||
original_weight = int(binding.get("weight") or 0)
|
||
operations.append(modify_gateway_binding_weight(binding, original_weight))
|
||
expected_weights[binding_key(binding)] = original_weight
|
||
|
||
verified = wait_gateway_binding_weights(
|
||
snapshot,
|
||
expected_weights,
|
||
timeout_seconds=GATEWAY_CLB_TASK_TIMEOUT_SECONDS,
|
||
)
|
||
return {
|
||
"ok": True,
|
||
"updatedAt": utc_now(),
|
||
"operations": operations,
|
||
"bindings": verified,
|
||
"summary": render_weight_summary(verified),
|
||
}
|
||
|
||
|
||
def health_key(listener_id: str, location_id: str, target_id: str, port: int) -> str:
|
||
# 健康状态和绑定状态用同一套 key，便于直接按绑定反查对应健康条目。
|
||
return "|".join([listener_id.strip(), location_id.strip(), target_id.strip(), str(int(port or 0))])
|
||
|
||
|
||
def collect_gateway_health_states(load_balancers: list[dict[str, Any]], snapshot: dict[str, Any]) -> list[dict[str, Any]]:
|
||
# CLB 健康查询返回的是 loadBalancer -> listener -> rule -> target 结构，这里拍平后只留当前实例。
|
||
matched: dict[str, dict[str, Any]] = {}
|
||
for load_balancer in load_balancers:
|
||
for listener in load_balancer.get("Listeners") or []:
|
||
listener_id = str(listener.get("ListenerId") or "").strip()
|
||
for rule in listener.get("Rules") or []:
|
||
location_id = str(rule.get("LocationId") or "").strip()
|
||
for target in rule.get("Targets") or []:
|
||
if not target_matches_service(target, snapshot["service"]):
|
||
continue
|
||
state = {
|
||
"listenerId": listener_id,
|
||
"locationId": location_id,
|
||
"targetId": str(target.get("TargetId") or "").strip(),
|
||
"port": int(target.get("Port") or 0),
|
||
"healthStatus": bool(target.get("HealthStatus")),
|
||
"detail": str(target.get("HealthStatusDetail") or target.get("HealthStatusDetial") or "").strip(),
|
||
}
|
||
matched[health_key(state["listenerId"], state["locationId"], state["targetId"], state["port"])] = state
|
||
return list(matched.values())
|
||
|
||
|
||
def describe_gateway_target_health(snapshot: dict[str, Any]) -> list[dict[str, Any]]:
|
||
# 健康查询只拉当前负载均衡和监听器，避免把无关 listener 的状态混进来。
|
||
payload = call_clb(
|
||
"DescribeTargetHealth",
|
||
{
|
||
"LoadBalancerIds": [snapshot["loadBalancerId"]],
|
||
"ListenerIds": list(snapshot["listenerIds"]),
|
||
},
|
||
)
|
||
return collect_gateway_health_states(list(payload.get("LoadBalancers") or []), snapshot)
|
||
|
||
|
||
def wait_gateway_release_target_healthy(snapshot: dict[str, Any], *, timeout_seconds: int) -> list[dict[str, Any]]:
|
||
# 只有等 CLB 自己也判定节点健康，才允许恢复权重并继续下一台，避免下一轮摘流时只剩空节点。
|
||
expected_keys = {
|
||
health_key(
|
||
str(binding.get("listenerId") or "").strip(),
|
||
str(binding.get("locationId") or "").strip(),
|
||
str(binding.get("instanceId") or "").strip(),
|
||
int(binding.get("port") or 0),
|
||
)
|
||
for binding in snapshot["bindings"]
|
||
}
|
||
deadline = time.time() + max(float(timeout_seconds), 1.0)
|
||
last_states: list[dict[str, Any]] = []
|
||
while time.time() < deadline:
|
||
last_states = describe_gateway_target_health(snapshot)
|
||
current_by_key = {
|
||
health_key(
|
||
str(item.get("listenerId") or "").strip(),
|
||
str(item.get("locationId") or "").strip(),
|
||
str(item.get("targetId") or "").strip(),
|
||
int(item.get("port") or 0),
|
||
): item
|
||
for item in last_states
|
||
}
|
||
if current_by_key and all(
|
||
key in current_by_key and bool(current_by_key[key].get("healthStatus"))
|
||
for key in expected_keys
|
||
):
|
||
return last_states
|
||
time.sleep(max(float(GATEWAY_CLB_POLL_SECONDS), 0.1))
|
||
|
||
detail = ", ".join(
|
||
f"{item.get('listenerId')}/{item.get('locationId') or '-'}={item.get('detail') or item.get('healthStatus')}"
|
||
for item in last_states
|
||
)
|
||
raise RuntimeError(f"gateway CLB health not ready: {detail or 'no matching target health'}")
|