From 64ffa63b7cc5f42bda881e651a032c14212b1815 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 18:08:31 +0000 Subject: [PATCH 01/17] docs: design AWS locality-aware backend selection --- ...026-08-13-aws-locality-awareness-design.md | 638 ++++++++++++++++++ 1 file changed, 638 insertions(+) create mode 100644 docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md diff --git a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md new file mode 100644 index 0000000000..d5db3a75f0 --- /dev/null +++ b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md @@ -0,0 +1,638 @@ +# AWS Locality-Aware MySQL Backend Selection + +**Status:** Approved design + +**Date:** 2026-08-13 + +## Summary + +ProxySQL 4.0 will optionally prefer MySQL backends that are in the same AWS +Region or Availability Zone as the ProxySQL process. The feature changes only +the temporary weights used by a server-selection attempt. It never modifies +`mysql_servers.weight`, `runtime_mysql_servers.weight`, the saved +configuration, or ProxySQL Cluster checksums. + +The MySQL module owns the feature's configuration and traffic policy. The +general AWS plugin provides asynchronous, normalized AWS metadata and makes no +traffic-routing decisions. + +The first version supports RDS and Aurora endpoints. It does not discover +arbitrary MySQL servers on EC2. + +## Goals + +- Preserve configured server weights while giving operators a bounded local + Region and local AZ preference. +- Keep all AWS, IMDS, DNS, and credential-provider work out of connection + selection. +- Keep locality configuration in the MySQL module because MySQL Hostgroup + Manager consumes it. +- Discover every ProxySQL process's location independently so ProxySQL Cluster + cannot propagate one process's Region or AZ to another. +- Degrade to ordinary configured weights whenever metadata is unavailable, + expired, invalid, or unsupported. +- Make every classification and active multiplier observable without changing + `runtime_mysql_servers`. +- Reuse the general AWS plugin and its statically linked vendored AWS SDK + runtime alongside the IAM database-authentication capability. + +## Non-goals + +- Mutating configured or runtime server weights. +- Synchronizing discovered metadata through ProxySQL Cluster. +- Moving, closing, or rebalancing existing backend connections. +- Discovering arbitrary EC2-hosted MySQL servers. +- Resolving custom CNAMEs to infer an RDS or Aurora target. +- Scanning every AWS Region to locate an endpoint. +- Assuming roles into other AWS accounts. +- Managing RDS/Aurora topology or hostgroup membership. +- Supporting RDS Proxy endpoints in the first version. +- Providing PostgreSQL locality awareness in the first version. +- Making real AWS infrastructure mandatory for ordinary CI. + +## User-facing configuration + +### Global master switch + +The MySQL module adds one dynamic global variable: + +```text +mysql-aws_locality_awareness = false +``` + +It is available in the ProxySQL 4.0 build and defaults to `false`. + +Changing it follows the normal MySQL-variable lifecycle: + +```sql +SET mysql-aws_locality_awareness = true; +LOAD MYSQL VARIABLES TO RUNTIME; +SAVE MYSQL VARIABLES TO DISK; +``` + +There are deliberately no configured Region, AZ, or AWS-account variables. +Such variables could be synchronized to ProxySQL processes in other locations +and would therefore be unsafe. + +When the switch is disabled: + +- selection uses configured weights exactly as it does today; +- no new locality metadata refreshes are scheduled; +- in-flight locality requests are cancelled or ignored by generation; +- cached rows may remain visible in the diagnostic table with status + `disabled`, but cannot affect selection. + +### Per-hostgroup policy + +The existing `mysql_hostgroup_attributes.hostgroup_settings` JSON is the +configuration extension point: + +```json +{ + "aws": { + "locality_awareness": { + "same_region_multiplier": 2.0, + "same_az_multiplier": 4.0, + "refresh_interval_seconds": 300, + "stale_ttl_seconds": 1800 + } + } +} +``` + +Presence of a valid `aws.locality_awareness` object enables locality awareness +for that hostgroup. There is no additional per-hostgroup `enabled` field. + +Both multipliers are required and must be finite JSON numbers satisfying: + +```text +1.0 <= same_region_multiplier <= same_az_multiplier <= 10.0 +``` + +The timing fields are optional. Their defaults and accepted bounds are: + +```text +refresh_interval_seconds = 300 +stale_ttl_seconds = 1800 + +30 <= refresh_interval_seconds <= 86400 +refresh_interval_seconds <= stale_ttl_seconds <= 604800 +``` + +An invalid locality object disables locality bias for that hostgroup after the +load. Diagnostics identify the rejected field and hostgroup but never log the +complete JSON document. + +The existing `aws_iam_region` key remains an IAM-authentication setting. It is +not required by, or treated as authoritative for, locality discovery. + +## Selection semantics + +Locality produces a temporary effective weight for an eligible server: + +```text +remote or unknown configured_weight +same Region, different AZ int(configured_weight * same_region_multiplier) +same AZ int(configured_weight * same_az_multiplier) +``` + +The Region and AZ tiers are mutually exclusive. A same-AZ server receives only +`same_az_multiplier`; the two multipliers are never multiplied together. + +Conversion to an integer truncates toward zero. Weight zero remains zero. +Arithmetic uses a wide intermediate and saturates safely before entering the +64-bit weighted-selection accumulator. The current MySQL weight bounds make +saturation unlikely, but the operation must still be defined for every input. + +Examples with configured weights `10`, `20`, and `30`, Region multiplier +`2.0`, and AZ multiplier `4.0`: + +- first server in the same AZ: effective weight `40`; +- second server in the same Region but another AZ: effective weight `40`; +- third server in another Region: effective weight `30`. + +The values `10`, `20`, and `30` remain stored and reported by +`mysql_servers` and `runtime_mysql_servers`. + +### Classification rules + +The selector classifies a backend from one immutable metadata snapshot: + +- `same_az`: local and backend Regions match, local and backend AZ names match, + and both sides have the same confirmed AWS account ID; +- `same_region`: Regions match, but AZ is different, unavailable, inapplicable, + or cannot be trusted because account identity is unavailable or different; +- `remote`: both Regions are known and differ; +- `unknown`: either Region required for comparison is unknown. + +AZ names can map to different physical zones in different AWS accounts. The +same-AZ multiplier is therefore never applied without a same-account check. +Same-Region preference does not require matching accounts. + +### Existing eligibility remains authoritative + +Locality changes only the weighted lottery among candidates that have already +passed the existing rules, including: + +- ONLINE status and shun recovery; +- `max_connections` capacity; +- latency bounds; +- GTID requirements; +- replication-lag and Aurora-lag requirements; +- session-tracking capability backoff; +- the existing Aurora writer/replica filtering. + +Locality never makes a backend healthy, eligible, or available. + +### Global and thread-local pool paths + +The global path in `MyHGC::get_random_MySrvC()` computes each final +candidate's effective weight and performs its existing weighted selection with +a 64-bit accumulator. + +The per-thread local connection cache must also honor locality. Otherwise, a +remote cached connection could repeatedly bypass Hostgroup Manager's weighted +lottery. For locality-enabled hostgroups only, the local-cache path: + +1. finds connections that pass all existing compatibility, GTID, lag, health, + and session-state checks; +2. groups them by parent server, so a server with more idle connections does + not gain more selection probability; +3. selects a parent server using the same effective server weight helper; +4. returns a compatible cached connection belonging to that parent. + +When locality is inactive, the current local-cache first-match fast path stays +unchanged. + +One snapshot is retained for an entire selection attempt, preventing a refresh +from mixing classifications within one lottery. Metadata changes affect only +future selections. Existing connections are not migrated or closed. + +## Ownership and architecture + +### Chosen approach + +Core owns locality state and consumes an asynchronous AWS metadata provider. + +This was selected over two alternatives: + +1. A synchronous cached lookup into the plugin from every selection would add + plugin ABI calls and lifecycle/synchronization risk to a hot path. +2. A plugin-owned server-selection hook would move MySQL traffic policy into + the capability provider and violate the intended ownership boundary. + +### MySQL core responsibilities + +A core `MySQLAwsLocalityManager` owns: + +- parsed hostgroup policies; +- registered backend endpoint identities; +- configuration generations; +- refresh scheduling and request coalescing; +- normalized results received from the plugin; +- last-attempt and last-success times; +- per-policy fresh/stale/expired evaluation; +- immutable snapshots used by selection; +- diagnostic-table rows and redacted failure state. + +Core types contain only ProxySQL-owned strings, enums, timestamps, request +IDs, and result structures. They expose no AWS SDK types. + +The manager starts work only when the master switch is enabled and at least one +hostgroup has a valid locality policy. `LOAD MYSQL SERVERS TO RUNTIME` rebuilds +the endpoint registration set, advances its generation, and schedules the +necessary asynchronous refreshes. `LOAD MYSQL VARIABLES TO RUNTIME` activates +or bypasses the manager according to the master switch. + +### AWS plugin responsibilities + +The general `aws` plugin owns: + +- AWS SDK initialization and shutdown; +- the default AWS credential-provider chain; +- IMDSv2 access; +- regional RDS clients; +- paginated RDS API calls; +- bounded retries, timeouts, and background execution; +- cancellation and clean shutdown; +- normalization into the core-defined result contract. + +It makes no multiplier, eligibility, hostgroup, or traffic decision. + +The plugin extends its advertised capabilities beyond `aws_iam`, for example +with local-instance metadata and RDS-topology capabilities. The plugin reuses +the same SDK runtime already used by IAM authentication. + +### Generic asynchronous provider ABI + +ProxySQL's plugin services gain a generic AWS metadata-provider installation +contract. The first request kinds are: + +- discover the local ProxySQL process's AWS location; +- describe the RDS/Aurora endpoints in one candidate Region. + +Requests carry opaque IDs, deadlines, endpoint sets, and core configuration +generations. Results contain normalized endpoint type, Region, AZ where +applicable, account identity for comparison, timestamps, and a redacted status +category. + +The provider uses the same lease/drain principle as the IAM token source: + +- a plugin module cannot unload while requests or callbacks retain leases; +- shutdown stops accepting work, cancels queued work, and drains active work; +- callbacks target weak/core-owned completion sinks; +- callbacks run without holding plugin or Hostgroup Manager locks; +- core rejects completions from an obsolete configuration generation. + +No selection path invokes this ABI. + +### Immutable snapshot publication + +Core publishes immutable per-hostgroup locality snapshots. A snapshot maps the +current stable backend identity `(hostgroup_id, normalized hostname, port)` to +its classification inputs and metadata timestamps. Publication is atomic; a +selection retains one snapshot for its duration and performs no network calls +or mutable-cache locking. + +The optional feature may pay for immutable-map lookups. With the global switch +off or no hostgroup policy, the existing hot path bypasses those lookups. + +## Local ProxySQL location discovery + +Every ProxySQL process discovers its own location independently. Discovery is +node-local state and is never persisted or cluster-synchronized. + +Discovery order: + +1. Retrieve the EC2 instance identity document through IMDSv2. It provides + Region, Availability Zone, and account ID. +2. If IMDSv2 is unavailable, use process environment fallback: + - Region: `AWS_REGION`, then `AWS_DEFAULT_REGION`; + - AZ: `AWS_AVAILABILITY_ZONE`; + - account assertion: optional `AWS_ACCOUNT_ID`. +3. Leave any unavailable field unknown. + +An environment AZ is usable only with an environment Region. The same-AZ tier +also requires `AWS_ACCOUNT_ID`; without it, same-Region preference still +works. In Kubernetes, operators can inject the node topology AZ and Region as +pod environment values without adding synchronized ProxySQL settings. + +The EC2 instance identity document and fields are documented by AWS at: + + + +Local metadata follows the same refresh and stale policy as backend metadata. +If local Region expires, all locality classifications are neutral. If local +Region remains usable but AZ/account becomes unavailable, same-Region +classification remains possible while same-AZ does not. + +## Backend endpoint discovery + +### Candidate recognition + +Core recognizes official RDS/Aurora endpoint DNS forms only to extract a +candidate AWS Region and partition. This is routing for the API request, not +authoritative metadata. + +The implementation normalizes endpoint hostnames by lowercasing ASCII and +removing one trailing DNS dot. It does not resolve DNS or follow CNAMEs. +Supported official suffixes include the standard/GovCloud AWS suffix and the +China partition suffix. An unrecognized endpoint remains neutral. + +### Authoritative API matching + +Requests are coalesced by candidate Region. The plugin performs paginated: + +- `rds:DescribeDBInstances`; +- `rds:DescribeDBClusters`; +- `rds:DescribeDBClusterEndpoints`. + +Core accepts a result only when the normalized configured hostname exactly +matches an endpoint returned by the AWS APIs. Where the response supplies a +port, a configured port mismatch is rejected. A custom cluster endpoint that +does not expose a distinct port is matched by its exact authoritative +hostname. + +The endpoint mappings are: + +- RDS DB instance endpoint: `instance`, with Region, instance AZ, and account; +- Aurora DB instance endpoint: `instance`, with Region, instance AZ, and + account; +- Aurora or Multi-AZ cluster writer endpoint: `cluster`, with Region and + account but no stable endpoint AZ; +- Aurora reader endpoint: `reader`, with Region and account but no stable + endpoint AZ; +- Aurora custom endpoint: `custom`, with Region and account but no stable + endpoint AZ; +- unmatched or unsupported endpoint: `unknown`. + +Cluster, reader, and custom endpoints can route to instances in multiple AZs. +They can receive the same-Region multiplier but never the same-AZ multiplier. + +`DescribeDBInstances` exposes both an endpoint address and Availability Zone: + + + +`DescribeDBClusters` exposes cluster, reader, and member information: + + + +`DescribeDBClusterEndpoints` exposes custom and managed cluster endpoints: + + + +Custom CNAMEs, RDS Proxy endpoints, arbitrary hosts, and AWS-looking hostnames +that do not appear in an authoritative response stay `unknown`. + +## Refresh, sharing, and stale data + +Successful metadata records include a monotonic success time and a wall-clock +time for diagnostics. A failed refresh records an attempt time and redacted +error but does not immediately discard the last successful value. + +Each hostgroup evaluates freshness using its own policy: + +- `fresh`: age is no greater than `refresh_interval_seconds`; +- `stale`: a refresh is due or has failed, but age is no greater than + `stale_ttl_seconds`; the last successful metadata remains active; +- `expired`: age exceeds `stale_ttl_seconds`; metadata becomes unknown and the + configured weight is used; +- `error`: no usable successful metadata exists for the endpoint; +- `pending`: discovery has not completed yet; +- `disabled`: the global switch is off. + +If several hostgroups reference the same endpoint with different intervals, +the endpoint is refreshed at the shortest active interval. The shared result +retains its success timestamp; each hostgroup independently determines whether +that result is fresh, stale, or expired under its own TTL. + +Regional API scans are coalesced so one in-flight scan serves all registered +endpoints in that Region. Repeated load operations cancel or supersede older +generations. Late completions cannot attach to removed servers, removed +policies, or a newer generation. + +A successful scan that does not contain a configured endpoint records +`endpoint_not_found`. A prior match may remain active only through its bounded +stale TTL, after which the endpoint becomes neutral. + +## Failure behavior and security + +Every failure is fail-neutral, not fail-closed for database traffic: + +- missing AWS plugin; +- plugin unload or shutdown; +- missing credentials; +- IMDS disabled or unreachable; +- Kubernetes without injected location; +- access denied; +- throttling; +- timeout; +- malformed or unsupported endpoint; +- endpoint not found; +- callback cancellation; +- metadata expiration. + +In all cases, the backend remains subject to its ordinary eligibility and +configured weight. + +Logs are rate-limited by stable endpoint/Region/error-category keys. They never +include credentials, authorization headers, IMDS tokens, raw AWS errors, +account IDs, or the complete hostgroup JSON. Supported fixed categories +include: + +```text +access_denied +throttled +provider_unavailable +imds_unavailable +endpoint_not_found +timeout +cancelled +invalid_response +``` + +The plugin uses the normal AWS SDK credential provider chain. ProxySQL adds no +access-key or secret-key settings. Expected deployments include EC2 instance +profiles, EKS IRSA or Pod Identity, and externally provided standard AWS +credential sources. + +The read-only RDS policy required for backend discovery is: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "rds:DescribeDBInstances", + "rds:DescribeDBClusters", + "rds:DescribeDBClusterEndpoints" + ], + "Resource": "*" + } + ] +} +``` + +IMDS and environment discovery require no AWS API permission. + +## Observability + +The MySQL module exposes a read-only `stats_mysql_aws_locality` table with one +row for each backend in a hostgroup that currently has a valid locality policy: + +```text +hostgroup_id +hostname +port +endpoint_type +configured_weight +effective_weight +local_region +local_az +backend_region +backend_az +account_match +locality +active_multiplier +metadata_status +last_success_timestamp +last_attempt_timestamp +last_error_category +``` + +Definitions: + +- `endpoint_type`: `instance`, `cluster`, `reader`, `custom`, or `unknown`; +- `account_match`: `yes`, `no`, `unknown`, or `not_applicable`; +- `locality`: `same_az`, `same_region`, `remote`, or `unknown`; +- `active_multiplier`: the multiplier currently affecting selection, otherwise + `1.0`; +- `effective_weight`: a diagnostic calculation only; it is never written back + to a server table; +- timestamps: Unix epoch seconds, or zero if the event has never occurred; +- `metadata_status`: one of the lifecycle states defined above. + +Account IDs are never exposed. When the master switch is disabled, rows may +retain cached location text for diagnosis, but `active_multiplier` is `1.0`, +`effective_weight` equals `configured_weight`, and status is `disabled`. + +## Runtime sequence + +1. ProxySQL loads the AWS plugin and installs its generic metadata provider. +2. `LOAD MYSQL VARIABLES TO RUNTIME` enables the global feature. +3. `LOAD MYSQL SERVERS TO RUNTIME` parses locality policies, advances the core + registration generation, and schedules local and regional discovery. +4. The AWS plugin performs IMDS and RDS work on bounded background workers. +5. Completions return normalized metadata to the core manager. +6. Core validates request ID and generation, updates timestamps/error state, + and atomically publishes immutable hostgroup snapshots. +7. Global and local-cache selection attempts retain one snapshot and calculate + temporary effective weights for eligible server parents. +8. Periodic refreshes repeat at the shortest interval required by registered + hostgroups. Stale and expiration decisions remain per hostgroup. +9. Disabling the global variable immediately bypasses the snapshot and stops + scheduling new work. + +Until step 6 first succeeds, selection behaves exactly as it did before the +feature. + +## Verification strategy + +### Configuration and arithmetic unit tests + +- valid policy with defaults and explicit timing values; +- missing fields, wrong JSON types, NaN/infinity-equivalent rejection, + multiplier bounds and ordering; +- timing minimum, maximum, and `refresh <= stale` relationship; +- invalid reload removes prior locality influence; +- `1.0` and `10.0` multiplier boundaries; +- integer truncation, zero weight, non-cumulative tiers, and saturation; +- proof that configured/runtime table weights and checksums do not change. + +### Classification and discovery unit tests + +- IMDSv2 success and every failure phase; +- environment fallback precedence and partial values; +- missing account, matching account, and cross-account AZ-name collision; +- RDS instance and Aurora instance endpoints; +- cluster writer, reader, and custom endpoints; +- Multi-AZ/failover metadata refresh; +- remote Region and unknown local Region; +- custom CNAME, arbitrary host, unsupported RDS Proxy, and false AWS-looking + endpoint; +- exact normalized API endpoint match and port validation; +- paginated regional responses and duplicated endpoints across hostgroups; +- fixed/redacted errors without secrets or account IDs. + +### Cache and lifecycle unit tests + +- pending to fresh, fresh to stale, stale to expired, and recovery transitions + under a fake clock; +- different refresh/TTL policies sharing one endpoint; +- regional request coalescing and bounded queues; +- configuration reload, server removal, late completion, cancellation, and + generation rejection; +- provider replacement, plugin stop, core shutdown, and callback lifetime; +- enable, disable, and re-enable behavior; +- missing plugin and provider-unavailable neutral fallback; +- TSan coverage for publication, callbacks, reload, and shutdown. + +### Selection tests + +- deterministic effective-weight selection for the global Hostgroup Manager + path; +- deterministic parent-server weighting in the thread-local connection cache; +- proof that multiple cached connections do not amplify a server's weight; +- configured-weight relative ratios within each locality tier; +- no multiplier for unknown or expired metadata; +- cluster/reader/custom endpoints receive only same-Region preference; +- existing health, status, latency, lag, GTID, capacity, and backoff filters win + before locality; +- no locality-specific allocation, lock, plugin call, DNS, or network operation + in the hot path; +- current fast path remains in use when the feature is inactive. + +### Integration and regression tests + +- a fake asynchronous AWS provider drives the real MySQL Hostgroup Manager and + produces deterministic distributions; +- policy reload affects future selections only; +- `mysql_servers`, `runtime_mysql_servers`, saved configuration, and cluster + checksums remain byte-for-byte unchanged by discovered metadata; +- exact `stats_mysql_aws_locality` rows for fresh, stale, expired, error, and + disabled states; +- existing IAM database-authentication behavior continues through the shared + AWS plugin runtime; +- ASan, TSan, plugin lifecycle, static-linkage, SDK-free daemon, and existing + IAM/pool selection regressions remain green; +- optional externally provisioned AWS integration verifies one RDS/Aurora + instance endpoint and one cluster or reader endpoint without becoming a + normal-CI requirement. + +## Acceptance criteria + +The feature is complete when all of the following are true: + +1. The global switch defaults off and inactive builds preserve the existing + selection fast paths. +2. No locality operation mutates configured/runtime weights or cluster-visible + state. +3. Valid hostgroups apply bounded, non-cumulative Region/AZ multipliers only at + selection time. +4. The global pool and thread-local cache use identical server-level effective + weight semantics. +5. Instance endpoints can receive same-AZ preference; cluster, reader, and + custom endpoints cannot. +6. Same-AZ preference requires a confirmed same-account identity. +7. AWS/IMDS work is asynchronous, bounded, cancellable, and absent from the + hot path. +8. Refresh failure retains last-known metadata only through the configured + stale TTL, then returns to configured weighting. +9. Missing capability, credentials, permissions, or metadata never prevents a + database connection solely because locality awareness is enabled. +10. The diagnostic table explains every active or neutral decision without + exposing account IDs or sensitive AWS data. +11. Sanitizer, lifecycle, linkage, existing IAM, and selection regression gates + pass. From 211d2ba15b7cb66a419fe1c53db03a77d68bd37d Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 18:52:43 +0000 Subject: [PATCH 02/17] docs: define AWS locality stats lifecycle --- ...026-08-13-aws-locality-awareness-design.md | 59 ++++++++++++++++--- 1 file changed, 50 insertions(+), 9 deletions(-) diff --git a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md index d5db3a75f0..2d41459a1c 100644 --- a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md +++ b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md @@ -79,8 +79,8 @@ When the switch is disabled: - selection uses configured weights exactly as it does today; - no new locality metadata refreshes are scheduled; - in-flight locality requests are cancelled or ignored by generation; -- cached rows may remain visible in the diagnostic table with status - `disabled`, but cannot affect selection. +- when the AWS plugin is loaded, cached rows remain visible in its diagnostic + table with status `disabled`, but cannot affect selection. ### Per-hostgroup policy @@ -233,7 +233,7 @@ A core `MySQLAwsLocalityManager` owns: - last-attempt and last-success times; - per-policy fresh/stale/expired evaluation; - immutable snapshots used by selection; -- diagnostic-table rows and redacted failure state. +- diagnostic snapshot data and redacted failure state. Core types contain only ProxySQL-owned strings, enums, timestamps, request IDs, and result structures. They expose no AWS SDK types. @@ -263,6 +263,11 @@ The plugin extends its advertised capabilities beyond `aws_iam`, for example with local-instance metadata and RDS-topology capabilities. The plugin reuses the same SDK runtime already used by IAM authentication. +The plugin also registers the `stats_mysql_aws_locality` schema and its +query-time refresh callback. The MySQL module remains the source of the rows; +the plugin registration only makes the AWS-specific diagnostic surface exist +when the AWS capability is actually present. + ### Generic asynchronous provider ABI ProxySQL's plugin services gain a generic AWS metadata-provider installation @@ -479,8 +484,35 @@ IMDS and environment discovery require no AWS API permission. ## Observability -The MySQL module exposes a read-only `stats_mysql_aws_locality` table with one -row for each backend in a hostgroup that currently has a valid locality policy: +The AWS plugin registers `stats_mysql_aws_locality` as a read-only table in the +stats database. Its existence follows the plugin lifecycle: + +- when the AWS plugin loads successfully, its schema-registration phase adds + the table before the Admin databases are materialized; +- when the AWS plugin is not configured or does not load successfully, the + table is not created, and querying it returns the normal SQLite + `no such table` error; +- ProxySQL does not currently support hot unloading configured plugins. If hot + unload is introduced, the unload contract must unregister and drop this + table rather than leave an empty or stale table behind. + +The table is a query-time projection of the MySQL locality manager's current +immutable in-memory snapshot. Before a query that references the table is +executed, Admin invokes the registered refresh callback. The callback replaces +the prior SQLite rows in one transaction from one retained manager snapshot, +so a result never mixes locality generations. This is the same materialized- +on-query model used by other runtime and stats views; the SQLite rows are not +the authoritative locality state. + +Refreshing the table never performs an IMDS or AWS API request and never waits +for metadata discovery. Network refresh remains bounded asynchronous plugin +work; a table query reports the most recently published state, including +`pending`, `stale`, `expired`, or `error` as applicable. + +The projection is non-persistent: it is not saved to disk, loaded to runtime, +included in ProxySQL Cluster checksums, or accepted as configuration. Writes +to it are unsupported. Each refresh emits one row for each backend in a +hostgroup that currently has a valid locality policy: ```text hostgroup_id @@ -514,9 +546,12 @@ Definitions: - timestamps: Unix epoch seconds, or zero if the event has never occurred; - `metadata_status`: one of the lifecycle states defined above. -Account IDs are never exposed. When the master switch is disabled, rows may -retain cached location text for diagnosis, but `active_multiplier` is `1.0`, -`effective_weight` equals `configured_weight`, and status is `disabled`. +Account IDs are never exposed. When the AWS plugin is loaded but the master +switch is disabled, the table remains present and its rows retain cached +location text for diagnosis. Every row has `active_multiplier` equal to `1.0`, +`effective_weight` equal to `configured_weight`, and `metadata_status` equal +to `disabled`. If no hostgroup currently has a valid locality policy, the +table exists but is empty. ## Runtime sequence @@ -601,6 +636,11 @@ feature. - policy reload affects future selections only; - `mysql_servers`, `runtime_mysql_servers`, saved configuration, and cluster checksums remain byte-for-byte unchanged by discovered metadata; +- `stats_mysql_aws_locality` is absent without the AWS plugin and is registered + only when that plugin loads successfully; +- each table query projects one consistent in-memory snapshot without issuing + an IMDS or AWS API request, and projected rows are never persisted or + clustered; - exact `stats_mysql_aws_locality` rows for fresh, stale, expired, error, and disabled states; - existing IAM database-authentication behavior continues through the shared @@ -632,7 +672,8 @@ The feature is complete when all of the following are true: stale TTL, then returns to configured weighting. 9. Missing capability, credentials, permissions, or metadata never prevents a database connection solely because locality awareness is enabled. -10. The diagnostic table explains every active or neutral decision without +10. The plugin-conditional, query-refreshed diagnostic table explains every + active or neutral decision without performing network discovery or exposing account IDs or sensitive AWS data. 11. Sanitizer, lifecycle, linkage, existing IAM, and selection regression gates pass. From 7b4666e183062c3e6c57fb2358b3df3a85b17c2f Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 19:43:19 +0000 Subject: [PATCH 03/17] docs: plan AWS locality awareness implementation --- .../2026-08-13-aws-locality-awareness.md | 481 ++++++++++++++++++ 1 file changed, 481 insertions(+) create mode 100644 docs/superpowers/plans/2026-08-13-aws-locality-awareness.md diff --git a/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md b/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md new file mode 100644 index 0000000000..86e42cfcd4 --- /dev/null +++ b/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md @@ -0,0 +1,481 @@ +# AWS Locality-Aware Backend Selection Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Prefer eligible RDS/Aurora backends in the ProxySQL process's AWS Region or Availability Zone using temporary selection weights, without changing configured/runtime server weights. + +**Architecture:** MySQL core parses policy, owns refresh/cache state, publishes immutable selection snapshots, and projects diagnostics. The general `aws` plugin owns the shared AWS SDK runtime, IMDSv2 and RDS calls, and an asynchronous metadata provider installed through the plugin ABI. Both global and thread-local connection selection retain one core snapshot and use one pure effective-weight helper; no hot path calls the plugin or performs network work. + +**Tech Stack:** C++17, ProxySQL plugin ABI, AWS SDK for C++ 1.11.869 (`core` and `rds`), vendored libcurl, nlohmann JSON, SQLite stats runtime views, TAP unit/integration tests, ASan and TSan. + +## Global Constraints + +- Build tier is selected only with `PROXYSQL40=1`; there is no AWS-specific build flag. +- Invoke builds with `make -j`; never hard-code `-j` inside a Makefile recipe. +- The AWS SDK remains statically linked only into `plugins/aws/ProxySQL_Aws_Plugin.so`; the ProxySQL daemon and `libproxysql.a` remain free of AWS SDK symbols and DSOs. +- The AWS SDK release stays pinned to the existing vendored 1.11.869 LFS tarball and uses ProxySQL's vendored dependencies. +- Locality never changes `mysql_servers.weight`, `runtime_mysql_servers.weight`, saved configuration, or ProxySQL Cluster checksums. +- `mysql-aws_locality_awareness` is the only global control; Region, AZ, and account identity are node-local discoveries, never synchronized variables. +- Policy lives at `mysql_hostgroup_attributes.hostgroup_settings.aws.locality_awareness`. +- Multipliers are finite JSON numbers satisfying `1.0 <= same_region_multiplier <= same_az_multiplier <= 10.0`. +- Default refresh is 300 seconds; default stale TTL is 1800 seconds; accepted bounds are `30 <= refresh <= 86400` and `refresh <= stale_ttl <= 604800`. +- Every failure is fail-neutral and all logged/provider failure data is fixed-category and redacted. +- `stats_mysql_aws_locality` exists only when the AWS plugin loads successfully and querying it never performs network discovery. + +--- + +## File Structure + +- `include/Aws_Locality_Types.h`: SDK-free provider request/result, policy, endpoint identity, classification, diagnostics, and interfaces. +- `include/Aws_Locality_Manager.h`: provider lease/registry plus the core manager public API. +- `lib/Aws_Locality_Manager.cpp`: parsing helpers, endpoint recognition, classification, arithmetic, registry lifetime, scheduler/cache, immutable publication, and diagnostic snapshots. +- `include/ProxySQL_Plugin.h`, `lib/ProxySQL_PluginManager.cpp`: ABI-6 provider installation, core snapshot projection callback, and service wiring. +- `include/MySQL_HostGroups_Manager.h`, `lib/MySQL_HostGroups_Manager.cpp`: manager ownership, hostgroup policy storage, reload registration, and diagnostics projection. +- `include/MySQL_Thread.h`, `lib/MySQL_Thread.cpp`, `lib/Admin_FlushVariables.cpp`: `mysql-aws_locality_awareness` lifecycle and manager enable/disable notification. +- `lib/MyHGC.cpp`: global Hostgroup Manager effective weighting. +- `lib/MySQL_Thread.cpp`: locality-aware parent-server selection for the thread-local connection cache. +- `plugins/aws/src/aws_plugin.cpp`: shared SDK runtime and capability installation. +- `plugins/aws/src/aws_locality_provider.h`, `plugins/aws/src/aws_locality_provider.cpp`: bounded provider, IMDSv2/environment discovery, paginated RDS discovery, normalization, and cancellation. +- `plugins/aws/Makefile`, `lib/Makefile`: new compilation units and dependencies. +- `test/tap/tests/unit/aws_locality_policy_unit-t.cpp`: policy, DNS recognition, classification, and arithmetic. +- `test/tap/tests/unit/aws_locality_manager_unit-t.cpp`: cache, refresh, generation, stale, cancellation, and concurrency. +- `test/tap/tests/unit/aws_locality_selection_unit-t.cpp`: global and local-cache selection behavior. +- `test/tap/tests/unit/aws_locality_plugin_unit-t.cpp`: fake discovery backend exercising provider queue/normalization and environment fallback. +- `test/tap/tests/unit/aws_locality_stats_unit-t.cpp`: plugin-conditional table lifecycle and query-time projection. +- `test/tap/tests/unit/Makefile`, `test/tap/groups/groups.json`: targets and CI groups. +- `doc/aws-locality-awareness.md`, `README.md`: operator configuration, permissions, behavior, and diagnostics. + +--- + +### Task 1: SDK-Free Policy, Classification, and Weight Arithmetic + +**Files:** +- Create: `include/Aws_Locality_Types.h` +- Create: `include/Aws_Locality_Manager.h` +- Create: `lib/Aws_Locality_Manager.cpp` +- Modify: `lib/Makefile` +- Create: `test/tap/tests/unit/aws_locality_policy_unit-t.cpp` +- Modify: `test/tap/tests/unit/Makefile` +- Modify: `test/tap/groups/groups.json` + +**Interfaces:** +- Produces `AwsLocalityPolicy parse_aws_locality_policy(const nlohmann::json&, uint32_t hostgroup_id, AwsLocalityPolicyError&)`. +- Produces `AwsEndpointCandidate recognize_rds_endpoint(uint32_t, std::string_view, uint16_t)`. +- Produces `AwsLocalityClass classify_aws_locality(const AwsLocalLocation&, const AwsBackendLocation&)`. +- Produces `uint64_t aws_locality_effective_weight(int64_t configured_weight, double multiplier)`. + +- [ ] **Step 1: Write the failing policy and arithmetic test** + + Use literal cases for defaults, explicit timing, missing required multipliers, wrong JSON types, NaN-like invalid values, `1.0`/`10.0` bounds, ordering, truncation, zero, and `uint64_t` saturation. Include official instance/cluster/reader/custom endpoint names in standard, GovCloud, and China partitions, plus custom CNAME/RDS Proxy/arbitrary-host negatives. + + ```cpp + AwsLocalityPolicyError error; + const auto policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.5,"same_az_multiplier":4.75})"), 10, error); + ok(policy.valid && policy.refresh_interval_seconds == 300 && + policy.stale_ttl_seconds == 1800, "valid policy uses timing defaults"); + ok(aws_locality_effective_weight(3, 2.5) == 7, + "effective weight truncates toward zero"); + ``` + +- [ ] **Step 2: Run the focused target and verify RED** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_policy_unit-t` + + Expected: compilation fails only because the new locality types/functions are absent. + +- [ ] **Step 3: Implement the minimal pure model** + + Define SDK-free enums and structs with owned strings: + + ```cpp + enum class AwsEndpointType : uint8_t { unknown, instance, cluster, reader, custom }; + enum class AwsLocalityClass : uint8_t { unknown, remote, same_region, same_az }; + enum class AwsLocalityMetadataStatus : uint8_t { disabled, pending, fresh, stale, expired, error }; + + struct AwsLocalityPolicy { + bool valid{false}; + double same_region_multiplier{1.0}; + double same_az_multiplier{1.0}; + uint32_t refresh_interval_seconds{300}; + uint32_t stale_ttl_seconds{1800}; + }; + ``` + + Normalize hostnames by ASCII-lowercasing and removing one trailing dot. Recognition extracts only candidate Region/partition and never claims authoritative endpoint type. Use `long double` for multiplication, truncate toward zero, preserve zero, and saturate at `uint64_t::max()`. + +- [ ] **Step 4: Run focused GREEN and existing parser/selection regressions** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_policy_unit-t server_selection_unit-t aws_iam_policy_unit-t` + + Expected: all TAP plans pass with no warnings or SDK references in the focused locality binary. + +- [ ] **Step 5: Commit** + + ```bash + git add include/Aws_Locality_Types.h include/Aws_Locality_Manager.h \ + lib/Aws_Locality_Manager.cpp lib/Makefile \ + test/tap/tests/unit/aws_locality_policy_unit-t.cpp \ + test/tap/tests/unit/Makefile test/tap/groups/groups.json + git commit -m "feat(mysql): parse AWS locality policies" + ``` + +--- + +### Task 2: Provider Registry, Refresh Manager, and Immutable Snapshots + +**Files:** +- Modify: `include/Aws_Locality_Types.h` +- Modify: `include/Aws_Locality_Manager.h` +- Modify: `lib/Aws_Locality_Manager.cpp` +- Create: `test/tap/tests/unit/aws_locality_manager_unit-t.cpp` +- Modify: `test/tap/tests/unit/Makefile` +- Modify: `test/tap/groups/groups.json` + +**Interfaces:** +- Consumes Task 1 policy/candidate/classification types. +- Produces `AwsMetadataProvider`, `AwsMetadataCompletionSink`, `AwsMetadataProviderLease`, `install_global_aws_metadata_provider()`, `acquire_global_aws_metadata_provider()`, and `shutdown_global_aws_metadata_provider()`. +- Produces `MySQLAwsLocalityManager::{configure,set_enabled,snapshot,diagnostic_rows,shutdown}`. + +- [ ] **Step 1: Write a fake-provider manager test** + + Exercise a real manager against a deterministic provider that retains request IDs/generations and posts complete normalized results. Cover local discovery followed by coalesced regional requests, duplicate endpoints across hostgroups, shortest refresh interval, pending/fresh/stale/expired/error states, endpoint-not-found, provider absence, enable/disable/re-enable, invalid reload removal, late-generation rejection, cancel, provider replacement, and shutdown drain. + + ```cpp + class FakeAwsMetadataProvider final : public AwsMetadataProvider { + public: + AwsMetadataRequestHandle request(const AwsMetadataRequest& request, + std::weak_ptr sink) override; + void cancel(AwsMetadataRequestHandle handle) override; + }; + ``` + + Use an injected steady/wall clock; no sleeps for freshness assertions. Verify monotonic age calculations and wall-clock diagnostic timestamps independently. Add a deterministic callback-vs-shutdown test and a 100-repeat multi-producer publication test. + +- [ ] **Step 2: Run and verify RED** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_manager_unit-t` + + Expected: compile failure on the absent provider/manager APIs. + +- [ ] **Step 3: Implement the provider registry and lease** + + Mirror the proven IAM lease/drain contract, but keep a separate generic metadata registry. Installation transfers provider ownership plus an optional retained module handle. Shutdown disables new leases, waits for active leases, calls provider shutdown/destructor, and only then `dlclose()`s the retained module reference. + +- [ ] **Step 4: Implement manager scheduling and publication** + + `configure()` receives a copied vector of hostgroup policy/backend identities and advances a generation. A lazy scheduler thread exists only while enabled with at least one valid policy. It requests local identity and coalesces endpoint scans by Region, never holding the manager mutex across provider calls. Completions update mutable cache state under the manager mutex, then build and atomically publish `std::shared_ptr` objects. Selection snapshots contain no mutable locks or plugin pointers. + +- [ ] **Step 5: Run focused GREEN and TSan** + + Run: + + ```bash + PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_manager_unit-t + PROXYSQL40=1 NOJEMALLOC=1 WITHTSAN=1 make -C test/tap/tests/unit -j aws_locality_manager_unit-t + TSAN_OPTIONS=halt_on_error=1 test/tap/tests/unit/aws_locality_manager_unit-t + ``` + + Expected: full TAP plan passes and TSan reports no race. + +- [ ] **Step 6: Commit** + + ```bash + git add include/Aws_Locality_Types.h include/Aws_Locality_Manager.h \ + lib/Aws_Locality_Manager.cpp test/tap/tests/unit/aws_locality_manager_unit-t.cpp \ + test/tap/tests/unit/Makefile test/tap/groups/groups.json + git commit -m "feat(mysql): manage asynchronous AWS locality metadata" + ``` + +--- + +### Task 3: MySQL Variable and Hostgroup Reload Integration + +**Files:** +- Modify: `include/Base_HostGroups_Manager.h` +- Modify: `include/MySQL_HostGroups_Manager.h` +- Modify: `include/MySQL_Thread.h` +- Modify: `lib/BaseHGC.cpp` +- Modify: `lib/MySQL_HostGroups_Manager.cpp` +- Modify: `lib/MySQL_Thread.cpp` +- Modify: `lib/Admin_FlushVariables.cpp` +- Create: `test/tap/tests/unit/aws_locality_config_unit-t.cpp` +- Modify: `test/tap/tests/unit/Makefile` +- Modify: `test/tap/groups/groups.json` + +**Interfaces:** +- Consumes `MySQLAwsLocalityManager::configure()` and `set_enabled()`. +- Produces `MyHGC::attributes.aws_locality_policy` and `MySQL_HostGroups_Manager::refresh_aws_locality_configuration()`. +- Produces the dynamic boolean variable `mysql-aws_locality_awareness`, default `false`, only in `PROXYSQL40`. + +- [ ] **Step 1: Write failing configuration lifecycle tests** + + Initialize real hostgroup attributes from JSON and assert valid policy installation, exact defaults/bounds, malformed field rejection without full JSON logging, and removal of the prior policy after invalid reload. Exercise `MySQL_Threads_Handler::{set_variable,get_variable,commit}` for false/true parsing and notify a real manager after `LOAD` semantics. Compare `runtime_mysql_servers` and checksum input before/after metadata completions. + +- [ ] **Step 2: Run and verify RED** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_config_unit-t` + + Expected: compile failure on missing variable/policy fields and refresh method. + +- [ ] **Step 3: Add the variable and policy storage** + + Register `aws_locality_awareness` in the existing bool variable table, default it false, copy it to worker variables, and expose it as `mysql-aws_locality_awareness`. Compile all feature behavior under `PROXYSQL40`. Add an owned policy value to `MyHGC` and reset it on every attributes reload before parsing. + +- [ ] **Step 4: Wire load boundaries** + + At the end of `MySQL_HostGroups_Manager::commit()`, while server/hostgroup state is stable, copy valid policies and backend identities and call `configure()` after releasing the HGM lock. After MySQL variable commit and lock release, call `set_enabled()` once with the master value. Never include discovered data in generated tables/checksums. + +- [ ] **Step 5: Run focused GREEN and existing config regressions** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_config_unit-t aws_iam_connection_config_unit-t hostgroups_unit-t cluster_sync_unit-t` + +- [ ] **Step 6: Commit** + + ```bash + git add include/Base_HostGroups_Manager.h include/MySQL_HostGroups_Manager.h \ + include/MySQL_Thread.h lib/BaseHGC.cpp lib/MySQL_HostGroups_Manager.cpp \ + lib/MySQL_Thread.cpp lib/Admin_FlushVariables.cpp \ + test/tap/tests/unit/aws_locality_config_unit-t.cpp \ + test/tap/tests/unit/Makefile test/tap/groups/groups.json + git commit -m "feat(mysql): load AWS locality configuration" + ``` + +--- + +### Task 4: Global and Thread-Local Weighted Selection + +**Files:** +- Modify: `include/Aws_Locality_Manager.h` +- Modify: `include/MySQL_HostGroups_Manager.h` +- Modify: `lib/MyHGC.cpp` +- Modify: `lib/MySQL_Thread.cpp` +- Create: `test/tap/tests/unit/aws_locality_selection_unit-t.cpp` +- Modify: `test/tap/tests/unit/Makefile` +- Modify: `test/tap/groups/groups.json` + +**Interfaces:** +- Consumes immutable `AwsLocalitySnapshot` and `effective_weight(hostgroup,host,port,configured_weight)`. +- Produces identical server-level weighted selection semantics in `MyHGC::get_random_MySrvC()` and `MySQL_Thread::get_MyConn_local()`. + +- [ ] **Step 1: Write failing global/local selection tests** + + Use real `MyHGC`, `MySrvC`, and cached `MySQL_Connection` fixtures. Publish literal same-AZ/same-Region/remote/unknown metadata for weights 10/20/30 and multipliers 4.0/2.0; assert deterministic effective weights 40/40/30. Verify same-AZ requires matching account, cluster/reader/custom receive only Region bias, stale remains active, expired is neutral, and configured weights never change. + + For local cache, place multiple connections on one remote parent and one connection on one local parent. Run seeded selections and prove probability follows parent weights, not connection count. Include health, GTID, lag, auth compatibility, session state, and backoff exclusions before locality. + +- [ ] **Step 2: Run and verify RED** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_selection_unit-t` + + Expected: locality distribution assertions fail while legacy controls pass. + +- [ ] **Step 3: Integrate global selection** + + Preserve the entire existing eligibility scan. Retain one snapshot before candidate evaluation only when the thread-local master flag and hostgroup policy are active. Store a parallel `uint64_t` effective-weight array, use a saturating 64-bit sum, and run the existing lottery over those values. The inactive path retains current branches and configured-weight arithmetic. + +- [ ] **Step 4: Integrate local-cache selection** + + Keep the existing first-match implementation unchanged when locality is inactive. When active, scan compatible eligible connections, group candidate indices by `MySrvC*`, calculate one effective weight per parent, choose a parent with the same helper/lottery, then remove one best compatible connection belonging to that parent. Do not allocate/call plugins on the inactive path; reuse bounded stack storage before falling back to a vector for unusually large candidate sets. + +- [ ] **Step 5: Run focused GREEN and pool regressions** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_selection_unit-t server_selection_unit-t aws_iam_pool_unit-t connection_pool_unit-t` + +- [ ] **Step 6: Commit** + + ```bash + git add include/Aws_Locality_Manager.h include/MySQL_HostGroups_Manager.h \ + lib/MyHGC.cpp lib/MySQL_Thread.cpp \ + test/tap/tests/unit/aws_locality_selection_unit-t.cpp \ + test/tap/tests/unit/Makefile test/tap/groups/groups.json + git commit -m "feat(mysql): apply AWS locality during backend selection" + ``` + +--- + +### Task 5: AWS Plugin Metadata Provider + +**Files:** +- Create: `plugins/aws/src/aws_locality_provider.h` +- Create: `plugins/aws/src/aws_locality_provider.cpp` +- Modify: `plugins/aws/src/aws_plugin.cpp` +- Modify: `plugins/aws/Makefile` +- Modify: `include/ProxySQL_Plugin.h` +- Modify: `lib/ProxySQL_PluginManager.cpp` +- Create: `test/tap/tests/unit/aws_locality_plugin_unit-t.cpp` +- Modify: `test/tap/tests/unit/Makefile` +- Modify: `test/tap/groups/groups.json` + +**Interfaces:** +- Consumes Task 2 `AwsMetadataProvider` and ownership callbacks. +- Produces ABI-6 services `install_aws_metadata_provider` and the `aws_locality` advertised plugin capability. +- Produces `AwsSdkMetadataProvider`, with an injectable `AwsLocalityDiscoveryBackend` for deterministic tests. + +- [ ] **Step 1: Write the failing provider test** + + Drive the real bounded queue/provider with a fake discovery backend. Assert full request/result fields, two-worker bound, queue rejection, cancellation, deadline rejection before/after work, no callback after shutdown, redacted categories, and generation/opaque ID preservation. Test environment precedence (`AWS_REGION`, `AWS_DEFAULT_REGION`, `AWS_AVAILABILITY_ZONE`, `AWS_ACCOUNT_ID`) and partial fallback using scoped environment restoration. + + Build literal normalized response fixtures for RDS instances, Aurora instances, cluster writer/reader/custom endpoints, missing ports, duplicate pages, Multi-AZ/failover changes, and endpoint-not-found. Never assert on fake call existence alone; assert emitted normalized results. + +- [ ] **Step 2: Run and verify RED** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_plugin_unit-t` + + Expected: compilation fails on missing provider/backend types. + +- [ ] **Step 3: Extend the ABI and shared SDK lifetime** + + Increment the plugin ABI maximum/current version to 6 and append metadata-provider installation to `ProxySQL_PluginServices`. Wire it only during plugin init. Refactor the plugin so IAM signer/token source and locality provider each retain a `std::shared_ptr`; `Aws::InitAPI` occurs once and `Aws::ShutdownAPI` occurs only after both core-owned capabilities drain. + +- [ ] **Step 4: Implement local discovery** + + Use IMDSv2 token `PUT /latest/api/token` with a bounded TTL header, then `GET /latest/dynamic/instance-identity/document` with the token. Use ProxySQL's vendored libcurl already linked into the plugin, enforce link-local target/timeouts/response bounds, and parse only `region`, `availabilityZone`, and `accountId`. On IMDS failure, apply the exact environment fallback order. Never log the document, token, account ID, raw curl error, or environment values. + +- [ ] **Step 5: Implement paginated RDS discovery** + + Maintain regional `Aws::RDS::RDSClient` instances under a client-map mutex. Issue paginated `DescribeDBInstances`, `DescribeDBClusters`, and `DescribeDBClusterEndpoints` requests until marker exhaustion or deadline/cancellation. Normalize only authoritative endpoint addresses and supplied ports. Map SDK errors to fixed categories (`access_denied`, `throttled`, `timeout`, `invalid_response`, `provider_unavailable`) and discard raw messages. Rate-limit logs by stable Region/endpoint/category keys without including raw responses or account identity. + +- [ ] **Step 6: Run focused GREEN, plugin loader, and secret scans** + + Run: + + ```bash + PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_plugin_unit-t aws_plugin_load_unit-t + PROXYSQL40=1 make -C plugins/aws -j + nm -C plugins/aws/ProxySQL_Aws_Plugin.so > /tmp/proxysql-aws-plugin-nm.txt + ldd plugins/aws/ProxySQL_Aws_Plugin.so > /tmp/proxysql-aws-plugin-ldd.txt + ``` + + Assert the plugin has locality symbols, has no AWS/CRT shared-library dependencies, and test/log output contains none of the fixture credentials/account IDs/tokens. + +- [ ] **Step 7: Commit** + + ```bash + git add plugins/aws/src/aws_locality_provider.h plugins/aws/src/aws_locality_provider.cpp \ + plugins/aws/src/aws_plugin.cpp plugins/aws/Makefile include/ProxySQL_Plugin.h \ + lib/ProxySQL_PluginManager.cpp test/tap/tests/unit/aws_locality_plugin_unit-t.cpp \ + test/tap/tests/unit/Makefile test/tap/groups/groups.json + git commit -m "feat(aws): discover RDS locality metadata" + ``` + +--- + +### Task 6: Plugin-Conditional Stats Table + +**Files:** +- Modify: `include/ProxySQL_Plugin.h` +- Modify: `lib/ProxySQL_PluginManager.cpp` +- Modify: `include/MySQL_HostGroups_Manager.h` +- Modify: `lib/MySQL_HostGroups_Manager.cpp` +- Modify: `plugins/aws/src/aws_plugin.cpp` +- Create: `test/tap/tests/unit/aws_locality_stats_unit-t.cpp` +- Modify: `test/tap/tests/unit/Makefile` +- Modify: `test/tap/groups/groups.json` + +**Interfaces:** +- Produces the plugin-owned `stats_mysql_aws_locality` schema and runtime-view registration. +- Produces ABI service `refresh_mysql_aws_locality_stats(SQLite3DB*)`, callable by the plugin's static refresh callback. +- Consumes `MySQLAwsLocalityManager::diagnostic_rows()`. + +- [ ] **Step 1: Write failing table lifecycle/projection tests** + + Bootstrap Admin with no AWS plugin and assert `SELECT * FROM stats_mysql_aws_locality` returns `no such table`. Bootstrap through the real AWS plugin schema-registration phase and assert the exact 17-column schema exists. Publish manager rows for pending/fresh/stale/expired/error/disabled, query through `ProxySQL_Admin` twice across a generation swap, and assert each result is a complete single-generation snapshot. + + Use a provider request counter to prove querying the table issues zero metadata requests. Assert writes are rejected and no disk/config table/checksum contains the name. + +- [ ] **Step 2: Run and verify RED** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_stats_unit-t` + + Expected: plugin-loaded table query fails because the schema/callback is absent. + +- [ ] **Step 3: Register schema and refresh callback** + + In `register_schemas`, register only a stats-db table: + + ```sql + CREATE TABLE stats_mysql_aws_locality ( + hostgroup_id INT NOT NULL, hostname VARCHAR NOT NULL, port INT NOT NULL, + endpoint_type VARCHAR NOT NULL, configured_weight INT NOT NULL, + effective_weight INT NOT NULL, local_region VARCHAR NOT NULL, + local_az VARCHAR NOT NULL, backend_region VARCHAR NOT NULL, + backend_az VARCHAR NOT NULL, account_match VARCHAR NOT NULL, + locality VARCHAR NOT NULL, active_multiplier REAL NOT NULL, + metadata_status VARCHAR NOT NULL, last_success_timestamp INT NOT NULL, + last_attempt_timestamp INT NOT NULL, last_error_category VARCHAR NOT NULL, + PRIMARY KEY(hostgroup_id, hostname, port)) + ``` + + The plugin callback invokes the core service. Core retains one diagnostics snapshot, executes `BEGIN; DELETE; INSERT...; COMMIT`, and never calls the provider. With the master off, preserve cached text but force multiplier 1.0/effective configured/status disabled. No valid policies means zero rows. + +- [ ] **Step 4: Run focused GREEN and lifecycle regressions** + + Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_stats_unit-t aws_plugin_load_unit-t plugin_runtime_views_unit-t test_aws_iam_metrics-t` + +- [ ] **Step 5: Commit** + + ```bash + git add include/ProxySQL_Plugin.h lib/ProxySQL_PluginManager.cpp \ + include/MySQL_HostGroups_Manager.h lib/MySQL_HostGroups_Manager.cpp \ + plugins/aws/src/aws_plugin.cpp test/tap/tests/unit/aws_locality_stats_unit-t.cpp \ + test/tap/tests/unit/Makefile test/tap/groups/groups.json + git commit -m "feat(stats): expose AWS locality decisions" + ``` + +--- + +### Task 7: Operator Documentation and Final Verification + +**Files:** +- Create: `doc/aws-locality-awareness.md` +- Modify: `README.md` +- Modify: `.github/workflows/CI-aws.yml` +- Modify: `docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md` only if implementation review exposes an approved contract correction. + +**Interfaces:** +- Consumes all preceding production/test behavior. +- Produces operator documentation and CI gates; no new runtime API. + +- [ ] **Step 1: Write operator documentation** + + Document the variable, JSON example, numeric/timing bounds, integer truncation, non-cumulative tiers, instance-vs-cluster AZ behavior, account requirement, environment fallback, EC2/EKS credential delivery, exact read-only RDS IAM policy, stale/fail-neutral behavior, and plugin-conditional stats table. State explicitly that displayed/runtime configured weights never change. + +- [ ] **Step 2: Extend established container CI** + + Add the five locality tests to the existing AWS workflow's established ProxySQL build container. Preserve `actions/checkout` LFS hydration, pass `PROXYSQL40=1 make -j` from workflow invocations, and do not install compiler/development dependencies directly onto the GitHub runner. + +- [ ] **Step 3: Run the complete normal regression gate** + + Run: + + ```bash + PROXYSQL40=1 make -j clean + PROXYSQL40=1 make -j + PROXYSQL40=1 make -C test/tap/tests/unit -j + ``` + + Execute every generated unit binary and record TAP totals. Run all existing AWS IAM, connection-pool, hostgroup, cluster/checksum, plugin lifecycle/runtime-view, controlled TLS, and metrics targets explicitly. + +- [ ] **Step 4: Run sanitizer gates** + + Build and run policy, manager, config, selection, plugin, stats, and affected IAM/pool tests under ASan+LSan. Run manager, selection, provider, stats, and plugin lifecycle concurrency tests under TSan. Restore normal artifacts afterward using `PROXYSQL40=1 make -j clean && PROXYSQL40=1 make -j`. + +- [ ] **Step 5: Run linkage/security/final-diff gates** + + Capture `nm` and `ldd` output to files before grepping. Prove daemon/archive contain no `Aws::` symbols and daemon/plugin have no AWS/CRT DSOs; prove plugin contains expected static SDK/locality symbols. Scan the combined diff/test output for credentials, tokens, account IDs, raw AWS errors, and unredacted fixture markers. Run `git diff --check` and validate the vendored archive remains unmodified. + + If an externally provisioned AWS integration runner, credentials, and RDS endpoints are configured, run one instance-endpoint and one cluster/reader-endpoint locality test. Otherwise record the optional gate as `NOT RUN`; never substitute fake-provider coverage and label it real AWS verification. + +- [ ] **Step 6: Request independent review and fix all Critical/Important findings** + + Provide the reviewer the approved design, this plan, base SHA, head SHA, exact verification evidence, and explicit non-goals. For every valid finding, write a focused failing test before the production correction, then rerun affected and full gates. + +- [ ] **Step 7: Commit documentation/CI and prepare handoff** + + ```bash + git add doc/aws-locality-awareness.md README.md .github/workflows/CI-aws.yml + git commit -m "docs: document AWS locality awareness" + ``` + + Do not push or open/retarget a PR until the user requests publication. From 7d1220a1ccdb8e1d1880c957d759e1032a946b9a Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 19:55:30 +0000 Subject: [PATCH 04/17] feat(mysql): parse AWS locality policies --- include/Aws_Locality_Manager.h | 28 ++ include/Aws_Locality_Types.h | 66 +++++ lib/Aws_Locality_Manager.cpp | 240 ++++++++++++++++++ lib/Makefile | 2 +- test/tap/groups/groups.json | 1 + test/tap/tests/unit/Makefile | 8 +- .../tests/unit/aws_locality_policy_unit-t.cpp | 180 +++++++++++++ 7 files changed, 523 insertions(+), 2 deletions(-) create mode 100644 include/Aws_Locality_Manager.h create mode 100644 include/Aws_Locality_Types.h create mode 100644 lib/Aws_Locality_Manager.cpp create mode 100644 test/tap/tests/unit/aws_locality_policy_unit-t.cpp diff --git a/include/Aws_Locality_Manager.h b/include/Aws_Locality_Manager.h new file mode 100644 index 0000000000..cf9938964f --- /dev/null +++ b/include/Aws_Locality_Manager.h @@ -0,0 +1,28 @@ +#ifndef AWS_LOCALITY_MANAGER_H +#define AWS_LOCALITY_MANAGER_H + +#include "Aws_Locality_Types.h" +#include "json_fwd.hpp" + +#include +#include + +AwsLocalityPolicy parse_aws_locality_policy( + const nlohmann::json& policy_json, + uint32_t hostgroup_id, + AwsLocalityPolicyError& error); + +AwsEndpointCandidate recognize_rds_endpoint( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port); + +AwsLocalityClass classify_aws_locality( + const AwsLocalLocation& local, + const AwsBackendLocation& backend); + +uint64_t aws_locality_effective_weight( + int64_t configured_weight, + double multiplier); + +#endif // AWS_LOCALITY_MANAGER_H diff --git a/include/Aws_Locality_Types.h b/include/Aws_Locality_Types.h new file mode 100644 index 0000000000..fa2fc37501 --- /dev/null +++ b/include/Aws_Locality_Types.h @@ -0,0 +1,66 @@ +#ifndef AWS_LOCALITY_TYPES_H +#define AWS_LOCALITY_TYPES_H + +#include +#include + +enum class AwsEndpointType : uint8_t { + unknown, + instance, + cluster, + reader, + custom, +}; + +enum class AwsLocalityClass : uint8_t { + unknown, + remote, + same_region, + same_az, +}; + +enum class AwsLocalityMetadataStatus : uint8_t { + disabled, + pending, + fresh, + stale, + expired, + error, +}; + +struct AwsLocalityPolicy { + bool valid { false }; + double same_region_multiplier { 1.0 }; + double same_az_multiplier { 1.0 }; + uint32_t refresh_interval_seconds { 300 }; + uint32_t stale_ttl_seconds { 1800 }; +}; + +struct AwsLocalityPolicyError { + uint32_t hostgroup_id { 0 }; + std::string field; +}; + +struct AwsEndpointCandidate { + bool recognized { false }; + uint32_t hostgroup_id { 0 }; + std::string hostname; + uint16_t port { 0 }; + std::string region; + std::string partition; +}; + +struct AwsLocalLocation { + std::string region; + std::string availability_zone; + std::string account_id; +}; + +struct AwsBackendLocation { + AwsEndpointType endpoint_type { AwsEndpointType::unknown }; + std::string region; + std::string availability_zone; + std::string account_id; +}; + +#endif // AWS_LOCALITY_TYPES_H diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp new file mode 100644 index 0000000000..75db6e3c7f --- /dev/null +++ b/lib/Aws_Locality_Manager.cpp @@ -0,0 +1,240 @@ +#include "Aws_Locality_Manager.h" + +#include "json.hpp" + +#include +#include +#include +#include +#include + +using nlohmann::json; + +namespace { + +AwsLocalityPolicy invalid_policy( + uint32_t hostgroup_id, + const char* field, + AwsLocalityPolicyError& error) { + error.hostgroup_id = hostgroup_id; + error.field = field; + return {}; +} + +bool read_multiplier(const json& object, const char* field, double& value) { + const auto it = object.find(field); + if (it == object.end() || !it->is_number()) { + return false; + } + + value = it->get(); + return std::isfinite(value) && value >= 1.0 && value <= 10.0; +} + +bool read_seconds( + const json& object, + const char* field, + uint32_t default_value, + uint32_t minimum, + uint32_t maximum, + uint32_t& value) { + const auto it = object.find(field); + if (it == object.end()) { + value = default_value; + return true; + } + if (!it->is_number_unsigned() && !it->is_number_integer()) { + return false; + } + + const int64_t parsed = it->get(); + if (parsed < static_cast(minimum) || + parsed > static_cast(maximum)) { + return false; + } + value = static_cast(parsed); + return true; +} + +std::string normalized_hostname(std::string_view input) { + if (input.empty()) { + return {}; + } + + std::string hostname(input); + if (hostname.back() == '.') { + hostname.pop_back(); + } + if (hostname.empty() || hostname.back() == '.') { + return {}; + } + + for (char& character : hostname) { + const unsigned char value = static_cast(character); + if (!(std::isalnum(value) || character == '-' || character == '.')) { + return {}; + } + character = static_cast(std::tolower(value)); + } + return hostname; +} + +bool ends_with(const std::string& value, const std::string& suffix) { + return value.size() >= suffix.size() && + value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +bool valid_region(const std::string& region) { + if (region.empty() || region.front() == '-' || region.back() == '-') { + return false; + } + + unsigned int hyphens = 0; + for (const char character : region) { + if (character == '-') { + ++hyphens; + } else if (!std::islower(static_cast(character)) && + !std::isdigit(static_cast(character))) { + return false; + } + } + return hyphens >= 2 && + std::isdigit(static_cast(region.back())); +} + +bool contains_proxy_label(const std::string& prefix) { + size_t begin = 0; + while (begin < prefix.size()) { + const size_t end = prefix.find('.', begin); + const size_t length = end == std::string::npos + ? prefix.size() - begin : end - begin; + if (length >= 6 && prefix.compare(begin, 6, "proxy-") == 0) { + return true; + } + if (end == std::string::npos) { + break; + } + begin = end + 1; + } + return false; +} + +} // namespace + +AwsLocalityPolicy parse_aws_locality_policy( + const json& policy_json, + uint32_t hostgroup_id, + AwsLocalityPolicyError& error) { + error = {}; + if (!policy_json.is_object()) { + return invalid_policy(hostgroup_id, "locality_awareness", error); + } + + AwsLocalityPolicy policy; + if (!read_multiplier(policy_json, "same_region_multiplier", + policy.same_region_multiplier)) { + return invalid_policy(hostgroup_id, "same_region_multiplier", error); + } + if (!read_multiplier(policy_json, "same_az_multiplier", + policy.same_az_multiplier) || + policy.same_az_multiplier < policy.same_region_multiplier) { + return invalid_policy(hostgroup_id, "same_az_multiplier", error); + } + if (!read_seconds(policy_json, "refresh_interval_seconds", 300, 30, 86400, + policy.refresh_interval_seconds)) { + return invalid_policy(hostgroup_id, "refresh_interval_seconds", error); + } + if (!read_seconds(policy_json, "stale_ttl_seconds", 1800, + policy.refresh_interval_seconds, 604800, + policy.stale_ttl_seconds)) { + return invalid_policy(hostgroup_id, "stale_ttl_seconds", error); + } + + policy.valid = true; + return policy; +} + +AwsEndpointCandidate recognize_rds_endpoint( + uint32_t hostgroup_id, + std::string_view hostname_input, + uint16_t port) { + AwsEndpointCandidate result; + result.hostgroup_id = hostgroup_id; + result.port = port; + result.hostname = normalized_hostname(hostname_input); + if (result.hostname.empty()) { + return result; + } + + const std::string china_suffix = ".rds.amazonaws.com.cn"; + const std::string standard_suffix = ".rds.amazonaws.com"; + const std::string* suffix = nullptr; + if (ends_with(result.hostname, china_suffix)) { + suffix = &china_suffix; + } else if (ends_with(result.hostname, standard_suffix)) { + suffix = &standard_suffix; + } else { + return result; + } + + const std::string before_suffix = result.hostname.substr( + 0, result.hostname.size() - suffix->size()); + const size_t region_separator = before_suffix.rfind('.'); + if (region_separator == std::string::npos || region_separator == 0 || + region_separator + 1 == before_suffix.size()) { + return result; + } + + const std::string endpoint_prefix = before_suffix.substr(0, region_separator); + result.region = before_suffix.substr(region_separator + 1); + if (contains_proxy_label(endpoint_prefix) || !valid_region(result.region)) { + result.region.clear(); + return result; + } + + if (result.region.compare(0, 3, "cn-") == 0) { + result.partition = "aws-cn"; + } else if (result.region.compare(0, 7, "us-gov-") == 0) { + result.partition = "aws-us-gov"; + } else { + result.partition = "aws"; + } + result.recognized = true; + return result; +} + +AwsLocalityClass classify_aws_locality( + const AwsLocalLocation& local, + const AwsBackendLocation& backend) { + if (local.region.empty() || backend.region.empty()) { + return AwsLocalityClass::unknown; + } + if (local.region != backend.region) { + return AwsLocalityClass::remote; + } + if (backend.endpoint_type == AwsEndpointType::instance && + !local.availability_zone.empty() && + local.availability_zone == backend.availability_zone && + !local.account_id.empty() && + local.account_id == backend.account_id) { + return AwsLocalityClass::same_az; + } + return AwsLocalityClass::same_region; +} + +uint64_t aws_locality_effective_weight( + int64_t configured_weight, + double multiplier) { + if (configured_weight <= 0 || !std::isfinite(multiplier) || multiplier <= 0.0) { + return 0; + } + + const long double product = static_cast(configured_weight) * + static_cast(multiplier); + const long double maximum = static_cast( + std::numeric_limits::max()); + if (product >= maximum) { + return std::numeric_limits::max(); + } + return static_cast(product); +} diff --git a/lib/Makefile b/lib/Makefile index 3a9e46ba22..f7871d82e7 100644 --- a/lib/Makefile +++ b/lib/Makefile @@ -93,7 +93,7 @@ MYCXXFLAGS := $(STDCPP) $(MYCFLAGS) $(PSQLCH) $(PSQL40) $(PSQL31) $(PSQLFFTO) $( default: libproxysql.a .PHONY: default -_OBJ_CXX := ProxySQL_GloVars.oo network.oo debug.oo configfile.oo Query_Cache.oo SpookyV2.oo MySQL_Authentication.oo MySQL_Backend_Auth.oo Aws_Iam_Provider.oo MySQL_Passthrough_Auth_Cache.oo gen_utils.oo sqlite3db.oo mysql_connection.oo MySQL_HostGroups_Manager.oo mysql_data_stream.oo MySQL_Thread.oo MySQL_Session.oo MySQL_Protocol.oo mysql_backend.oo Query_Processor.oo MySQL_Query_Processor.oo PgSQL_Query_Processor.oo ProxySQL_Admin.oo ProxySQL_Config.oo ProxySQL_Restapi.oo MySQL_Monitor.oo MySQL_Logger.oo log_utils.oo thread.oo MySQL_PreparedStatement.oo ProxySQL_Cluster.oo ClickHouse_Authentication.oo ClickHouse_Server.oo ProxySQL_Statistics.oo Chart_bundle_js.oo ProxySQL_HTTP_Server.oo ProxySQL_RESTAPI_Server.oo font-awesome.min.css.oo main-bundle.min.css.oo MySQL_Variables.oo MySQL_User_Variables.oo c_tokenizer.oo proxysql_utils.oo proxysql_coredump.oo proxysql_sslkeylog.oo \ +_OBJ_CXX := ProxySQL_GloVars.oo network.oo debug.oo configfile.oo Query_Cache.oo SpookyV2.oo MySQL_Authentication.oo MySQL_Backend_Auth.oo Aws_Iam_Provider.oo Aws_Locality_Manager.oo MySQL_Passthrough_Auth_Cache.oo gen_utils.oo sqlite3db.oo mysql_connection.oo MySQL_HostGroups_Manager.oo mysql_data_stream.oo MySQL_Thread.oo MySQL_Session.oo MySQL_Protocol.oo mysql_backend.oo Query_Processor.oo MySQL_Query_Processor.oo PgSQL_Query_Processor.oo ProxySQL_Admin.oo ProxySQL_Config.oo ProxySQL_Restapi.oo MySQL_Monitor.oo MySQL_Logger.oo log_utils.oo thread.oo MySQL_PreparedStatement.oo ProxySQL_Cluster.oo ClickHouse_Authentication.oo ClickHouse_Server.oo ProxySQL_Statistics.oo Chart_bundle_js.oo ProxySQL_HTTP_Server.oo ProxySQL_RESTAPI_Server.oo font-awesome.min.css.oo main-bundle.min.css.oo MySQL_Variables.oo MySQL_User_Variables.oo c_tokenizer.oo proxysql_utils.oo proxysql_coredump.oo proxysql_sslkeylog.oo \ sha256crypt.oo \ ProxySQL_PluginManager.oo \ BaseSrvList.oo BaseHGC.oo Base_HostGroups_Manager.oo \ diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index 1ac357b547..e73c69801b 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -21,6 +21,7 @@ "aws_iam_pool_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_iam_session_state_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], "aws_iam_provider_boundary_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], + "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], "basic-t" : [ "legacy-g1","mariadb10-galera-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql84-gr-g1","mysql90-g1","mysql90-gr-g1","mysql93-g1","mysql93-gr-g1","mysql95-g1","mysql95-gr-g1" ], "c_tokenizer_unit-t" : [ "unit-tests-g1" ], diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index a1877a0ca1..1d2fd33afd 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -419,7 +419,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ @@ -935,6 +935,12 @@ aws_iam_completion_queue_unit-t: aws_iam_completion_queue_unit-t.cpp \ $(CXX) $< $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o $(IDIRS) $(OPT) \ -lpthread -lcrypto -o $@ +aws_locality_policy_unit-t: aws_locality_policy_unit-t.cpp \ + $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ + $(IDIRS) $(OPT) -lpthread -o $@ + test_aws_iam_metrics-t: ../test_aws_iam_metrics-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) $(CXX) $< $(TEST_HELPERS_OBJ) -I$(TEST_HELPERS_DIR) \ $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ diff --git a/test/tap/tests/unit/aws_locality_policy_unit-t.cpp b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp new file mode 100644 index 0000000000..41d64359e0 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp @@ -0,0 +1,180 @@ +#include "tap.h" + +#include "json.hpp" +#include "Aws_Locality_Manager.h" + +#include +#include +#include + +using nlohmann::json; + +namespace { + +void test_policy_validation() { + AwsLocalityPolicyError error; + AwsLocalityPolicy policy = parse_aws_locality_policy( + json::parse(R"({"same_region_multiplier":2.5,"same_az_multiplier":4.75})"), + 10, error); + ok(policy.valid && policy.same_region_multiplier == 2.5 && + policy.same_az_multiplier == 4.75 && + policy.refresh_interval_seconds == 300 && + policy.stale_ttl_seconds == 1800 && error.field.empty(), + "valid policy uses explicit multipliers and timing defaults"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":1.0,"same_az_multiplier":10.0,"refresh_interval_seconds":30,"stale_ttl_seconds":604800})"), + 11, error); + ok(policy.valid && policy.refresh_interval_seconds == 30 && + policy.stale_ttl_seconds == 604800, + "inclusive multiplier and timing boundaries are accepted"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_az_multiplier":4.0})"), 12, error); + ok(!policy.valid && error.field == "same_region_multiplier", + "missing same-Region multiplier rejects the complete locality policy"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":"4"})"), 13, error); + ok(!policy.valid && error.field == "same_az_multiplier", + "a numeric-looking string is not accepted as a multiplier"); + + json non_finite = { + {"same_region_multiplier", std::numeric_limits::quiet_NaN()}, + {"same_az_multiplier", 4.0} + }; + policy = parse_aws_locality_policy(non_finite, 14, error); + ok(!policy.valid && error.field == "same_region_multiplier", + "non-finite multiplier is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":5.0,"same_az_multiplier":4.0})"), 15, error); + ok(!policy.valid && error.field == "same_az_multiplier", + "same-AZ multiplier cannot be lower than same-Region multiplier"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":0.99,"same_az_multiplier":4.0})"), 16, error); + ok(!policy.valid && error.field == "same_region_multiplier", + "multiplier below 1.0 is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":10.01})"), 17, error); + ok(!policy.valid && error.field == "same_az_multiplier", + "multiplier above 10.0 is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":4.0,"refresh_interval_seconds":29})"), + 18, error); + ok(!policy.valid && error.field == "refresh_interval_seconds", + "refresh interval below 30 seconds is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":4.0,"refresh_interval_seconds":500,"stale_ttl_seconds":499})"), + 19, error); + ok(!policy.valid && error.field == "stale_ttl_seconds", + "stale TTL shorter than refresh interval is rejected"); + + policy = parse_aws_locality_policy(json::array(), 20, error); + ok(!policy.valid && error.field == "locality_awareness", + "non-object locality policy is rejected"); +} + +void test_endpoint_recognition() { + AwsEndpointCandidate candidate = recognize_rds_endpoint( + 10, "DB-1.ABCDEF.us-east-1.rds.amazonaws.com.", 3306); + ok(candidate.recognized && + candidate.hostname == "db-1.abcdef.us-east-1.rds.amazonaws.com" && + candidate.region == "us-east-1" && candidate.partition == "aws" && + candidate.hostgroup_id == 10 && candidate.port == 3306, + "standard RDS endpoint is normalized and routed to its Region"); + + candidate = recognize_rds_endpoint( + 11, "db.cluster-abcdefghijkl.us-gov-west-1.rds.amazonaws.com", 3306); + ok(candidate.recognized && candidate.region == "us-gov-west-1" && + candidate.partition == "aws-us-gov", + "GovCloud RDS endpoint is recognized"); + + candidate = recognize_rds_endpoint( + 12, "db.cluster-ro-abcdefghijkl.cn-north-1.rds.amazonaws.com.cn", 3306); + ok(candidate.recognized && candidate.region == "cn-north-1" && + candidate.partition == "aws-cn", + "China RDS endpoint is recognized"); + + candidate = recognize_rds_endpoint( + 13, "db.cluster-custom-abcdefghijkl.eu-west-1.rds.amazonaws.com", 3306); + ok(candidate.recognized && candidate.region == "eu-west-1", + "Aurora custom endpoint is a valid discovery candidate"); + + ok(!recognize_rds_endpoint(14, "db.proxy-abcdefghijkl.us-east-1.rds.amazonaws.com", 3306).recognized, + "RDS Proxy endpoint is excluded from the first version"); + ok(!recognize_rds_endpoint(15, "db.internal.example.com", 3306).recognized, + "custom CNAME is not inferred as RDS"); + ok(!recognize_rds_endpoint(16, "db.us-east-1.rds.amazonaws.com.evil.example", 3306).recognized, + "AWS-looking hostname with a false suffix is rejected"); + ok(!recognize_rds_endpoint(17, "db.INVALID_REGION.rds.amazonaws.com", 3306).recognized, + "malformed candidate Region is rejected"); +} + +void test_classification() { + const AwsLocalLocation local {"us-east-1", "us-east-1a", "111122223333"}; + + AwsBackendLocation backend { + AwsEndpointType::instance, "us-east-1", "us-east-1a", "111122223333" + }; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_az, + "instance in same account and AZ receives same-AZ classification"); + + backend.account_id = "444455556666"; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "same AZ name in a different account is only same-Region"); + + backend.account_id.clear(); + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "unknown backend account cannot receive same-AZ classification"); + + backend = {AwsEndpointType::cluster, "us-east-1", "us-east-1a", "111122223333"}; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "cluster endpoint never receives same-AZ classification"); + + backend = {AwsEndpointType::reader, "us-east-1", "", "111122223333"}; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "reader endpoint can receive same-Region classification"); + + backend = {AwsEndpointType::instance, "eu-west-1", "eu-west-1a", "111122223333"}; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::remote, + "different known Region is remote"); + + backend.region.clear(); + ok(classify_aws_locality(local, backend) == AwsLocalityClass::unknown, + "missing backend Region is neutral"); + + AwsLocalLocation unknown_local {"", "", ""}; + backend.region = "us-east-1"; + ok(classify_aws_locality(unknown_local, backend) == AwsLocalityClass::unknown, + "missing local Region is neutral"); +} + +void test_effective_weight() { + ok(aws_locality_effective_weight(3, 2.5) == 7, + "effective weight truncates toward zero"); + ok(aws_locality_effective_weight(0, 10.0) == 0, + "configured zero weight remains zero"); + ok(aws_locality_effective_weight(-1, 4.0) == 0, + "unexpected negative configured weight is safely neutralized"); + ok(aws_locality_effective_weight(std::numeric_limits::max(), 10.0) == + std::numeric_limits::max(), + "effective weight saturates before the 64-bit accumulator"); + ok(aws_locality_effective_weight(10, 1.0) == 10, + "neutral multiplier preserves configured weight"); +} + +} // namespace + +int main() { + plan(32); + test_policy_validation(); + test_endpoint_recognition(); + test_classification(); + test_effective_weight(); + return exit_status(); +} From 567419bfebdaeb7dc1f7793b6ece0ccf14584103 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 20:12:40 +0000 Subject: [PATCH 05/17] feat(mysql): manage asynchronous AWS locality metadata --- include/Aws_Locality_Manager.h | 100 +++ include/Aws_Locality_Types.h | 88 ++ lib/Aws_Locality_Manager.cpp | 807 ++++++++++++++++++ test/tap/groups/groups.json | 1 + test/tap/tests/unit/Makefile | 8 +- .../unit/aws_locality_manager_unit-t.cpp | 512 +++++++++++ 6 files changed, 1515 insertions(+), 1 deletion(-) create mode 100644 test/tap/tests/unit/aws_locality_manager_unit-t.cpp diff --git a/include/Aws_Locality_Manager.h b/include/Aws_Locality_Manager.h index cf9938964f..a3fc63ffda 100644 --- a/include/Aws_Locality_Manager.h +++ b/include/Aws_Locality_Manager.h @@ -4,8 +4,14 @@ #include "Aws_Locality_Types.h" #include "json_fwd.hpp" +#include #include +#include +#include +#include #include +#include +#include AwsLocalityPolicy parse_aws_locality_policy( const nlohmann::json& policy_json, @@ -25,4 +31,98 @@ uint64_t aws_locality_effective_weight( int64_t configured_weight, double multiplier); +using AwsMetadataProviderDestroyFn = void (*)(AwsMetadataProvider*); + +class AwsMetadataProviderLease { +public: + AwsMetadataProviderLease() = default; + ~AwsMetadataProviderLease(); + AwsMetadataProviderLease(AwsMetadataProviderLease&& other) noexcept; + AwsMetadataProviderLease& operator=(AwsMetadataProviderLease&& other) noexcept; + + AwsMetadataProvider* get() const { return provider_; } + AwsMetadataProvider* operator->() const { return provider_; } + explicit operator bool() const { return provider_ != nullptr; } + + AwsMetadataProviderLease(const AwsMetadataProviderLease&) = delete; + AwsMetadataProviderLease& operator=(const AwsMetadataProviderLease&) = delete; + +private: + explicit AwsMetadataProviderLease(AwsMetadataProvider* provider) + : provider_(provider) {} + void release(); + AwsMetadataProvider* provider_ { nullptr }; + + friend AwsMetadataProviderLease acquire_global_aws_metadata_provider(); +}; + +bool install_global_aws_metadata_provider( + AwsMetadataProvider* provider, + AwsMetadataProviderDestroyFn destroy, + void* module_handle); +AwsMetadataProviderLease acquire_global_aws_metadata_provider(); +void shutdown_global_aws_metadata_provider(); + +struct AwsLocalitySnapshotEntry { + uint32_t hostgroup_id { 0 }; + std::string hostname; + uint16_t port { 0 }; + AwsEndpointType endpoint_type { AwsEndpointType::unknown }; + int64_t configured_weight { 0 }; + AwsLocalLocation local; + AwsBackendLocation backend; + AwsLocalityClass locality { AwsLocalityClass::unknown }; + double multiplier { 1.0 }; + AwsLocalityMetadataStatus status { AwsLocalityMetadataStatus::disabled }; + int64_t last_success_timestamp { 0 }; + int64_t last_attempt_timestamp { 0 }; + std::string failure_category; +}; + +struct AwsLocalitySnapshot { + uint64_t generation { 0 }; + bool enabled { false }; + std::unordered_map entries; + + const AwsLocalitySnapshotEntry* find( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port) const; + uint64_t effective_weight( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port, + int64_t configured_weight) const; +}; + +struct AwsLocalityManagerConfig { + using SteadyClock = std::function; + using WallClock = std::function; + + AwsLocalityManagerConfig(); + SteadyClock steady_clock; + WallClock wall_clock; + std::chrono::milliseconds request_timeout { std::chrono::seconds(5) }; + std::function before_completion; +}; + +class MySQLAwsLocalityManager { +public: + explicit MySQLAwsLocalityManager(AwsLocalityManagerConfig config = {}); + ~MySQLAwsLocalityManager(); + MySQLAwsLocalityManager(const MySQLAwsLocalityManager&) = delete; + MySQLAwsLocalityManager& operator=(const MySQLAwsLocalityManager&) = delete; + + void configure(std::vector hostgroups); + void set_enabled(bool enabled); + void request_refresh(); + std::shared_ptr snapshot() const; + std::vector diagnostic_rows() const; + void shutdown(); + +private: + class Impl; + std::unique_ptr impl_; +}; + #endif // AWS_LOCALITY_MANAGER_H diff --git a/include/Aws_Locality_Types.h b/include/Aws_Locality_Types.h index fa2fc37501..720d7db96c 100644 --- a/include/Aws_Locality_Types.h +++ b/include/Aws_Locality_Types.h @@ -1,8 +1,12 @@ #ifndef AWS_LOCALITY_TYPES_H #define AWS_LOCALITY_TYPES_H +#include #include +#include #include +#include +#include enum class AwsEndpointType : uint8_t { unknown, @@ -63,4 +67,88 @@ struct AwsBackendLocation { std::string account_id; }; +enum class AwsMetadataRequestKind : uint8_t { + local_location, + rds_region, +}; + +enum class AwsMetadataStatus : uint8_t { + ok, + provider_unavailable, + access_denied, + throttled, + imds_unavailable, + timeout, + cancelled, + invalid_response, + shutdown, +}; + +struct AwsMetadataRequestHandle { + uint64_t value { 0 }; +}; + +struct AwsMetadataRequest { + AwsMetadataRequestKind kind { AwsMetadataRequestKind::local_location }; + uint64_t opaque_id { 0 }; + uint64_t generation { 0 }; + std::string region; + std::string partition; + std::vector endpoints; + std::chrono::steady_clock::time_point deadline {}; +}; + +struct AwsMetadataEndpoint { + std::string hostname; + uint16_t port { 0 }; + AwsEndpointType endpoint_type { AwsEndpointType::unknown }; + std::string region; + std::string availability_zone; + std::string account_id; +}; + +struct AwsMetadataResult { + AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; + AwsLocalLocation local; + std::vector endpoints; + std::string failure_category; +}; + +struct AwsMetadataCompletion { + uint64_t opaque_id { 0 }; + uint64_t generation { 0 }; + AwsMetadataResult result; +}; + +class AwsMetadataCompletionSink { +public: + virtual void post(AwsMetadataCompletion&& completion) = 0; + virtual ~AwsMetadataCompletionSink() = default; +}; + +class AwsMetadataProvider { +public: + virtual AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) = 0; + virtual void cancel(AwsMetadataRequestHandle handle) = 0; + virtual void shutdown() = 0; + virtual ~AwsMetadataProvider() = default; +}; + +struct AwsLocalityBackendConfig { + AwsEndpointCandidate endpoint; + int64_t configured_weight { 0 }; + + AwsLocalityBackendConfig() = default; + AwsLocalityBackendConfig(AwsEndpointCandidate value, int64_t weight = 0) + : endpoint(std::move(value)), configured_weight(weight) {} +}; + +struct AwsLocalityHostgroupConfig { + uint32_t hostgroup_id { 0 }; + AwsLocalityPolicy policy; + std::vector backends; +}; + #endif // AWS_LOCALITY_TYPES_H diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp index 75db6e3c7f..d27cf8505e 100644 --- a/lib/Aws_Locality_Manager.cpp +++ b/lib/Aws_Locality_Manager.cpp @@ -3,10 +3,20 @@ #include "json.hpp" #include +#include #include +#include #include +#include +#include #include +#include +#include #include +#include +#include +#include +#include using nlohmann::json; @@ -238,3 +248,800 @@ uint64_t aws_locality_effective_weight( } return static_cast(product); } + +namespace { + +std::mutex metadata_provider_mutex; +std::condition_variable metadata_provider_cv; +AwsMetadataProvider* leased_metadata_provider = nullptr; +AwsMetadataProviderDestroyFn metadata_provider_destroy = nullptr; +void* metadata_provider_module = nullptr; +size_t metadata_provider_leases = 0; +bool metadata_provider_accepting = false; + +std::string snapshot_key( + uint32_t hostgroup_id, + std::string_view hostname_input, + uint16_t port) { + const std::string hostname = normalized_hostname(hostname_input); + return std::to_string(hostgroup_id) + "\n" + hostname + "\n" + + std::to_string(port); +} + +std::string endpoint_key(std::string_view hostname_input, uint16_t port) { + return normalized_hostname(hostname_input) + "\n" + std::to_string(port); +} + +const char* failure_category(AwsMetadataStatus status) { + switch (status) { + case AwsMetadataStatus::ok: return ""; + case AwsMetadataStatus::provider_unavailable: return "provider_unavailable"; + case AwsMetadataStatus::access_denied: return "access_denied"; + case AwsMetadataStatus::throttled: return "throttled"; + case AwsMetadataStatus::imds_unavailable: return "imds_unavailable"; + case AwsMetadataStatus::timeout: return "timeout"; + case AwsMetadataStatus::cancelled: return "cancelled"; + case AwsMetadataStatus::invalid_response: return "invalid_response"; + case AwsMetadataStatus::shutdown: return "cancelled"; + } + return "invalid_response"; +} + +int64_t wall_seconds(std::chrono::system_clock::time_point value) { + return std::chrono::duration_cast( + value.time_since_epoch()).count(); +} + +} // namespace + +void AwsMetadataProviderLease::release() { + if (provider_ == nullptr) { + return; + } + { + std::lock_guard lock(metadata_provider_mutex); + if (metadata_provider_leases != 0) { + --metadata_provider_leases; + } + } + provider_ = nullptr; + metadata_provider_cv.notify_all(); +} + +AwsMetadataProviderLease::~AwsMetadataProviderLease() { + release(); +} + +AwsMetadataProviderLease::AwsMetadataProviderLease( + AwsMetadataProviderLease&& other) noexcept + : provider_(other.provider_) { + other.provider_ = nullptr; +} + +AwsMetadataProviderLease& AwsMetadataProviderLease::operator=( + AwsMetadataProviderLease&& other) noexcept { + if (this != &other) { + release(); + provider_ = other.provider_; + other.provider_ = nullptr; + } + return *this; +} + +bool install_global_aws_metadata_provider( + AwsMetadataProvider* provider, + AwsMetadataProviderDestroyFn destroy, + void* module_handle) { + if (provider == nullptr || destroy == nullptr) { + return false; + } + + std::lock_guard lock(metadata_provider_mutex); + if (metadata_provider_accepting || leased_metadata_provider != nullptr || + metadata_provider_destroy != nullptr) { + return false; + } + leased_metadata_provider = provider; + metadata_provider_destroy = destroy; + metadata_provider_module = module_handle; + metadata_provider_accepting = true; + return true; +} + +AwsMetadataProviderLease acquire_global_aws_metadata_provider() { + std::lock_guard lock(metadata_provider_mutex); + if (!metadata_provider_accepting || leased_metadata_provider == nullptr) { + return {}; + } + ++metadata_provider_leases; + return AwsMetadataProviderLease(leased_metadata_provider); +} + +void shutdown_global_aws_metadata_provider() { + AwsMetadataProvider* provider = nullptr; + AwsMetadataProviderDestroyFn destroy = nullptr; + void* module_handle = nullptr; + { + std::unique_lock lock(metadata_provider_mutex); + metadata_provider_accepting = false; + metadata_provider_cv.wait(lock, [] { + return metadata_provider_leases == 0; + }); + provider = leased_metadata_provider; + destroy = metadata_provider_destroy; + module_handle = metadata_provider_module; + leased_metadata_provider = nullptr; + metadata_provider_destroy = nullptr; + metadata_provider_module = nullptr; + } + + if (provider != nullptr) { + provider->shutdown(); + destroy(provider); + } + if (module_handle != nullptr) { + dlclose(module_handle); + } +} + +const AwsLocalitySnapshotEntry* AwsLocalitySnapshot::find( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port) const { + const auto it = entries.find(snapshot_key(hostgroup_id, hostname, port)); + return it == entries.end() ? nullptr : &it->second; +} + +uint64_t AwsLocalitySnapshot::effective_weight( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port, + int64_t configured_weight) const { + const auto* entry = find(hostgroup_id, hostname, port); + return aws_locality_effective_weight( + configured_weight, entry == nullptr ? 1.0 : entry->multiplier); +} + +AwsLocalityManagerConfig::AwsLocalityManagerConfig() + : steady_clock([] { return std::chrono::steady_clock::now(); }), + wall_clock([] { return std::chrono::system_clock::now(); }) {} + +class MySQLAwsLocalityManager::Impl { +public: + explicit Impl(AwsLocalityManagerConfig config) + : config_(std::move(config)), sink_(std::make_shared(this)) { + auto initial = std::make_shared(); + std::atomic_store_explicit( + &published_, std::shared_ptr(std::move(initial)), + std::memory_order_release); + } + + ~Impl() { + shutdown(); + } + + void configure(std::vector hostgroups) { + std::lock_guard lock(mutex_); + if (stopping_) { + return; + } + for (auto& hostgroup : hostgroups) { + for (auto& backend : hostgroup.backends) { + backend.endpoint.hostgroup_id = hostgroup.hostgroup_id; + } + } + hostgroups_.clear(); + for (auto& hostgroup : hostgroups) { + if (hostgroup.policy.valid) { + hostgroups_.push_back(std::move(hostgroup)); + } + } + ++generation_; + cancel_requested_ = true; + force_refresh_ = enabled_ && !hostgroups_.empty(); + publish_locked(); + cv_.notify_all(); + } + + void set_enabled(bool enabled) { + std::unique_lock lock(mutex_); + if (stopping_) { + return; + } + enabled_ = enabled; + if (enabled_ && !hostgroups_.empty()) { + ensure_worker_locked(); + force_refresh_ = true; + } else { + cancel_requested_ = true; + disable_acknowledged_ = !worker_.joinable(); + } + publish_locked(); + cv_.notify_all(); + if (!enabled_ && worker_.joinable()) { + cv_.wait(lock, [&] { + return disable_acknowledged_ || stopping_; + }); + } + } + + void request_refresh() { + std::lock_guard lock(mutex_); + if (stopping_ || !enabled_ || hostgroups_.empty()) { + return; + } + force_refresh_ = true; + publish_locked(); + cv_.notify_all(); + } + + std::shared_ptr snapshot() const { + return std::atomic_load_explicit(&published_, std::memory_order_acquire); + } + + std::vector diagnostic_rows() const { + const auto current = snapshot(); + std::vector rows; + rows.reserve(current->entries.size()); + for (const auto& item : current->entries) { + rows.push_back(item.second); + } + std::sort(rows.begin(), rows.end(), [](const auto& lhs, const auto& rhs) { + if (lhs.hostgroup_id != rhs.hostgroup_id) { + return lhs.hostgroup_id < rhs.hostgroup_id; + } + if (lhs.hostname != rhs.hostname) { + return lhs.hostname < rhs.hostname; + } + return lhs.port < rhs.port; + }); + return rows; + } + + void shutdown() { + std::thread worker; + { + std::lock_guard lock(mutex_); + if (shutdown_complete_) { + return; + } + stopping_ = true; + enabled_ = false; + cancel_requested_ = true; + publish_locked(); + cv_.notify_all(); + worker = std::move(worker_); + } + if (worker.joinable()) { + worker.join(); + } + sink_->detach(); + { + std::lock_guard lock(mutex_); + shutdown_complete_ = true; + publish_locked(); + } + } + +private: + struct EndpointRecord { + bool has_value { false }; + AwsBackendLocation value; + std::chrono::steady_clock::time_point success_steady {}; + int64_t success_wall { 0 }; + int64_t attempt_wall { 0 }; + std::string error; + }; + + struct LocalRecord { + bool has_value { false }; + AwsLocalLocation value; + std::chrono::steady_clock::time_point success_steady {}; + int64_t success_wall { 0 }; + int64_t attempt_wall { 0 }; + std::string error; + }; + + struct InFlight { + AwsMetadataRequest request; + AwsMetadataRequestHandle handle; + }; + + class CompletionSink final : public AwsMetadataCompletionSink { + public: + explicit CompletionSink(Impl* owner) : owner_(owner) {} + + void post(AwsMetadataCompletion&& completion) override { + Impl* owner = nullptr; + { + std::lock_guard lock(mutex_); + if (owner_ == nullptr) { + return; + } + owner = owner_; + ++active_; + } + owner->on_completion(std::move(completion)); + { + std::lock_guard lock(mutex_); + --active_; + cv_.notify_all(); + } + } + + void detach() { + std::unique_lock lock(mutex_); + owner_ = nullptr; + cv_.wait(lock, [&] { return active_ == 0; }); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + Impl* owner_ { nullptr }; + size_t active_ { 0 }; + }; + + void ensure_worker_locked() { + if (!worker_.joinable()) { + worker_ = std::thread([this] { worker_loop(); }); + } + } + + uint32_t minimum_refresh_seconds_locked() const { + uint32_t result = 86400; + for (const auto& hostgroup : hostgroups_) { + result = std::min(result, hostgroup.policy.refresh_interval_seconds); + } + return result; + } + + std::vector build_cycle_locked( + std::chrono::steady_clock::time_point now) { + std::vector requests; + AwsMetadataRequest local; + local.kind = AwsMetadataRequestKind::local_location; + local.opaque_id = next_opaque_id_++; + local.generation = generation_; + local.deadline = now + config_.request_timeout; + requests.push_back(local); + + std::map> regions; + for (const auto& hostgroup : hostgroups_) { + for (const auto& backend : hostgroup.backends) { + const auto& endpoint = backend.endpoint; + if (!endpoint.recognized) { + continue; + } + regions[endpoint.region][endpoint_key(endpoint.hostname, endpoint.port)] = endpoint; + } + } + for (const auto& region : regions) { + AwsMetadataRequest request; + request.kind = AwsMetadataRequestKind::rds_region; + request.opaque_id = next_opaque_id_++; + request.generation = generation_; + request.region = region.first; + request.deadline = now + config_.request_timeout; + for (const auto& item : region.second) { + request.endpoints.push_back(item.second); + if (request.partition.empty()) { + request.partition = item.second.partition; + } + } + requests.push_back(std::move(request)); + } + + const int64_t attempt = wall_seconds(config_.wall_clock()); + local_.attempt_wall = attempt; + ++local_in_flight_; + for (const auto& request : requests) { + in_flight_.emplace(request.opaque_id, InFlight {request, {}}); + if (request.kind == AwsMetadataRequestKind::rds_region) { + ++region_in_flight_[request.region]; + for (const auto& endpoint : request.endpoints) { + endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)].attempt_wall = attempt; + } + } + } + publish_locked(); + return requests; + } + + void cancel_all_requests(AwsMetadataProvider* provider) { + std::vector handles; + { + std::lock_guard lock(mutex_); + for (const auto& item : in_flight_) { + if (item.second.handle.value != 0) { + handles.push_back(item.second.handle); + } + } + in_flight_.clear(); + local_in_flight_ = 0; + region_in_flight_.clear(); + publish_locked(); + } + if (provider != nullptr) { + for (const auto handle : handles) { + provider->cancel(handle); + } + } + } + + void provider_unavailable(uint64_t generation) { + std::lock_guard lock(mutex_); + if (generation != generation_ || stopping_) { + return; + } + const int64_t attempt = wall_seconds(config_.wall_clock()); + local_.attempt_wall = attempt; + local_.error = "provider_unavailable"; + for (const auto& hostgroup : hostgroups_) { + for (const auto& backend : hostgroup.backends) { + auto& record = endpoint_cache_[endpoint_key( + backend.endpoint.hostname, backend.endpoint.port)]; + record.attempt_wall = attempt; + record.error = "provider_unavailable"; + } + } + in_flight_.clear(); + local_in_flight_ = 0; + region_in_flight_.clear(); + publish_locked(); + } + + void worker_loop() { + AwsMetadataProviderLease provider_lease; + bool cycle_started = false; + auto next_due = config_.steady_clock(); + + for (;;) { + std::vector requests; + bool cancel = false; + bool release_provider = false; + bool stop_now = false; + bool acknowledge_disable = false; + uint64_t dispatch_generation = 0; + { + std::unique_lock lock(mutex_); + while (!stopping_) { + if (cancel_requested_ || force_refresh_) { + break; + } + if (!enabled_ || hostgroups_.empty()) { + cv_.wait(lock); + continue; + } + const auto now = config_.steady_clock(); + if (!cycle_started || now >= next_due) { + break; + } + const auto delay = next_due - now; + cv_.wait_for(lock, delay); + } + + if (stopping_) { + cancel = true; + release_provider = true; + } else { + cancel = cancel_requested_; + cancel_requested_ = false; + if (cancel) { + // Cancel the previous generation before constructing new work. + // Otherwise the cancellation sweep can consume requests that + // were created in this same scheduler iteration. + if (!enabled_) { + release_provider = true; + acknowledge_disable = true; + } + } else if (!enabled_ || hostgroups_.empty()) { + release_provider = true; + acknowledge_disable = !enabled_; + force_refresh_ = false; + publish_locked(); + } else { + const auto now = config_.steady_clock(); + if (force_refresh_ || !cycle_started || now >= next_due) { + force_refresh_ = false; + requests = build_cycle_locked(now); + dispatch_generation = generation_; + cycle_started = true; + next_due = now + std::chrono::seconds( + minimum_refresh_seconds_locked()); + } + } + } + stop_now = stopping_; + } + + if (cancel) { + cancel_all_requests(provider_lease.get()); + } + if (release_provider) { + provider_lease = {}; + } + if (acknowledge_disable) { + std::lock_guard lock(mutex_); + disable_acknowledged_ = true; + cv_.notify_all(); + } + if (stop_now) { + break; + } + if (requests.empty()) { + continue; + } + + if (!provider_lease) { + provider_lease = acquire_global_aws_metadata_provider(); + } + if (!provider_lease) { + provider_unavailable(dispatch_generation); + continue; + } + + for (const auto& request : requests) { + AwsMetadataRequestHandle handle; + try { + handle = provider_lease->request(request, sink_); + } catch (...) { + AwsMetadataCompletion completion; + completion.opaque_id = request.opaque_id; + completion.generation = request.generation; + completion.result.status = AwsMetadataStatus::provider_unavailable; + on_completion(std::move(completion)); + continue; + } + std::lock_guard lock(mutex_); + const auto it = in_flight_.find(request.opaque_id); + if (it != in_flight_.end()) { + it->second.handle = handle; + } + } + } + } + + void on_completion(AwsMetadataCompletion&& completion) { + if (config_.before_completion) { + config_.before_completion(); + } + std::lock_guard lock(mutex_); + const auto pending = in_flight_.find(completion.opaque_id); + if (pending == in_flight_.end()) { + return; + } + const AwsMetadataRequest request = pending->second.request; + in_flight_.erase(pending); + if (request.kind == AwsMetadataRequestKind::local_location) { + if (local_in_flight_ != 0) { + --local_in_flight_; + } + } else { + auto region = region_in_flight_.find(request.region); + if (region != region_in_flight_.end() && region->second != 0 && + --region->second == 0) { + region_in_flight_.erase(region); + } + } + if (completion.generation != request.generation || + completion.generation != generation_ || stopping_) { + publish_locked(); + return; + } + + const auto now_steady = config_.steady_clock(); + const int64_t now_wall = wall_seconds(config_.wall_clock()); + if (request.kind == AwsMetadataRequestKind::local_location) { + local_.attempt_wall = now_wall; + if (completion.result.status == AwsMetadataStatus::ok && + !completion.result.local.region.empty()) { + local_.has_value = true; + local_.value = std::move(completion.result.local); + local_.success_steady = now_steady; + local_.success_wall = now_wall; + local_.error.clear(); + } else { + local_.error = completion.result.status == AwsMetadataStatus::ok + ? "invalid_response" : failure_category(completion.result.status); + } + } else { + apply_region_completion_locked(request, completion.result, + now_steady, now_wall); + } + publish_locked(); + } + + void apply_region_completion_locked( + const AwsMetadataRequest& request, + const AwsMetadataResult& result, + std::chrono::steady_clock::time_point now_steady, + int64_t now_wall) { + if (result.status != AwsMetadataStatus::ok) { + for (const auto& endpoint : request.endpoints) { + auto& record = endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)]; + record.attempt_wall = now_wall; + record.error = failure_category(result.status); + } + return; + } + + std::unordered_map returned; + for (const auto& endpoint : result.endpoints) { + const std::string hostname = normalized_hostname(endpoint.hostname); + if (!hostname.empty() && endpoint.region == request.region) { + returned[endpoint_key(hostname, endpoint.port)] = &endpoint; + if (endpoint.port == 0) { + returned[hostname + "\n0"] = &endpoint; + } + } + } + + for (const auto& endpoint : request.endpoints) { + auto& record = endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)]; + record.attempt_wall = now_wall; + const auto exact = returned.find(endpoint_key(endpoint.hostname, endpoint.port)); + const auto no_port = returned.find(normalized_hostname(endpoint.hostname) + "\n0"); + const AwsMetadataEndpoint* match = exact != returned.end() + ? exact->second : (no_port != returned.end() ? no_port->second : nullptr); + if (match == nullptr || match->endpoint_type == AwsEndpointType::unknown) { + record.error = "endpoint_not_found"; + continue; + } + record.has_value = true; + record.value = {match->endpoint_type, match->region, + match->availability_zone, match->account_id}; + record.success_steady = now_steady; + record.success_wall = now_wall; + record.error.clear(); + } + } + + AwsLocalitySnapshotEntry build_entry_locked( + const AwsLocalityHostgroupConfig& hostgroup, + const AwsLocalityBackendConfig& backend_config, + std::chrono::steady_clock::time_point now) const { + AwsLocalitySnapshotEntry entry; + entry.hostgroup_id = hostgroup.hostgroup_id; + entry.hostname = backend_config.endpoint.hostname; + entry.port = backend_config.endpoint.port; + entry.configured_weight = backend_config.configured_weight; + entry.local = local_.value; + + const auto endpoint = endpoint_cache_.find(endpoint_key( + backend_config.endpoint.hostname, backend_config.endpoint.port)); + if (endpoint != endpoint_cache_.end()) { + entry.backend = endpoint->second.value; + entry.endpoint_type = endpoint->second.value.endpoint_type; + entry.last_attempt_timestamp = std::max( + local_.attempt_wall, endpoint->second.attempt_wall); + if (local_.success_wall != 0 && endpoint->second.success_wall != 0) { + entry.last_success_timestamp = std::min( + local_.success_wall, endpoint->second.success_wall); + } + } else { + entry.last_attempt_timestamp = local_.attempt_wall; + } + + if (!enabled_) { + entry.status = AwsLocalityMetadataStatus::disabled; + return entry; + } + if (!backend_config.endpoint.recognized) { + entry.status = AwsLocalityMetadataStatus::error; + entry.failure_category = "invalid_response"; + return entry; + } + + const bool endpoint_pending = region_in_flight_.find( + backend_config.endpoint.region) != region_in_flight_.end(); + if (!local_.has_value || endpoint == endpoint_cache_.end() || + !endpoint->second.has_value) { + if ((!local_.has_value && local_in_flight_ != 0) || + (endpoint != endpoint_cache_.end() && !endpoint->second.has_value && + endpoint_pending)) { + entry.status = AwsLocalityMetadataStatus::pending; + } else { + entry.status = AwsLocalityMetadataStatus::error; + } + if (endpoint != endpoint_cache_.end() && !endpoint->second.error.empty()) { + entry.failure_category = endpoint->second.error; + } else if (!local_.error.empty()) { + entry.failure_category = local_.error; + } + return entry; + } + + const auto local_age = now >= local_.success_steady + ? now - local_.success_steady : std::chrono::steady_clock::duration::zero(); + const auto endpoint_age = now >= endpoint->second.success_steady + ? now - endpoint->second.success_steady : std::chrono::steady_clock::duration::zero(); + const auto age = std::max(local_age, endpoint_age); + const auto refresh = std::chrono::seconds(hostgroup.policy.refresh_interval_seconds); + const auto stale_ttl = std::chrono::seconds(hostgroup.policy.stale_ttl_seconds); + if (age > stale_ttl) { + entry.status = AwsLocalityMetadataStatus::expired; + return entry; + } + const bool refresh_failed = !local_.error.empty() || !endpoint->second.error.empty(); + entry.status = age > refresh || refresh_failed + ? AwsLocalityMetadataStatus::stale : AwsLocalityMetadataStatus::fresh; + entry.failure_category = !endpoint->second.error.empty() + ? endpoint->second.error : local_.error; + entry.locality = classify_aws_locality(local_.value, endpoint->second.value); + if (entry.locality == AwsLocalityClass::same_az) { + entry.multiplier = hostgroup.policy.same_az_multiplier; + } else if (entry.locality == AwsLocalityClass::same_region) { + entry.multiplier = hostgroup.policy.same_region_multiplier; + } + return entry; + } + + void publish_locked() { + auto next = std::make_shared(); + next->generation = generation_; + next->enabled = enabled_; + const auto now = config_.steady_clock(); + for (const auto& hostgroup : hostgroups_) { + for (const auto& backend : hostgroup.backends) { + auto entry = build_entry_locked(hostgroup, backend, now); + next->entries.emplace(snapshot_key(entry.hostgroup_id, + entry.hostname, entry.port), std::move(entry)); + } + } + std::atomic_store_explicit( + &published_, std::shared_ptr(std::move(next)), + std::memory_order_release); + } + + AwsLocalityManagerConfig config_; + mutable std::mutex mutex_; + std::condition_variable cv_; + std::vector hostgroups_; + uint64_t generation_ { 0 }; + uint64_t next_opaque_id_ { 1 }; + bool enabled_ { false }; + bool stopping_ { false }; + bool shutdown_complete_ { false }; + bool cancel_requested_ { false }; + bool force_refresh_ { false }; + bool disable_acknowledged_ { true }; + std::thread worker_; + std::shared_ptr sink_; + std::shared_ptr published_; + LocalRecord local_; + std::unordered_map endpoint_cache_; + std::unordered_map in_flight_; + size_t local_in_flight_ { 0 }; + std::unordered_map region_in_flight_; +}; + +MySQLAwsLocalityManager::MySQLAwsLocalityManager(AwsLocalityManagerConfig config) + : impl_(new Impl(std::move(config))) {} + +MySQLAwsLocalityManager::~MySQLAwsLocalityManager() = default; + +void MySQLAwsLocalityManager::configure( + std::vector hostgroups) { + impl_->configure(std::move(hostgroups)); +} + +void MySQLAwsLocalityManager::set_enabled(bool enabled) { + impl_->set_enabled(enabled); +} + +void MySQLAwsLocalityManager::request_refresh() { + impl_->request_refresh(); +} + +std::shared_ptr MySQLAwsLocalityManager::snapshot() const { + return impl_->snapshot(); +} + +std::vector MySQLAwsLocalityManager::diagnostic_rows() const { + return impl_->diagnostic_rows(); +} + +void MySQLAwsLocalityManager::shutdown() { + impl_->shutdown(); +} diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index e73c69801b..f41ac40f12 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -22,6 +22,7 @@ "aws_iam_session_state_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], "aws_iam_provider_boundary_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], "basic-t" : [ "legacy-g1","mariadb10-galera-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql84-gr-g1","mysql90-g1","mysql90-gr-g1","mysql93-g1","mysql93-gr-g1","mysql95-g1","mysql95-gr-g1" ], "c_tokenizer_unit-t" : [ "unit-tests-g1" ], diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index 1d2fd33afd..4d811b7b55 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -419,7 +419,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ @@ -941,6 +941,12 @@ aws_locality_policy_unit-t: aws_locality_policy_unit-t.cpp \ $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ $(IDIRS) $(OPT) -lpthread -o $@ +aws_locality_manager_unit-t: aws_locality_manager_unit-t.cpp \ + $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ + $(IDIRS) $(OPT) -lpthread -ldl -o $@ + test_aws_iam_metrics-t: ../test_aws_iam_metrics-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) $(CXX) $< $(TEST_HELPERS_OBJ) -I$(TEST_HELPERS_DIR) \ $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ diff --git a/test/tap/tests/unit/aws_locality_manager_unit-t.cpp b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp new file mode 100644 index 0000000000..f40ed2c51a --- /dev/null +++ b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp @@ -0,0 +1,512 @@ +#include "tap.h" + +#include "Aws_Locality_Manager.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +namespace { + +struct FakeProviderState { + struct Pending { + AwsMetadataRequestHandle handle; + AwsMetadataRequest request; + std::weak_ptr sink; + }; + + std::mutex mutex; + std::condition_variable cv; + std::vector requests; + std::vector canceled; + uint64_t next_handle { 1 }; + bool shut_down { false }; + bool destroyed { false }; +}; + +class FakeProvider final : public AwsMetadataProvider { +public: + explicit FakeProvider(std::shared_ptr state) + : state_(std::move(state)) {} + + ~FakeProvider() override { + std::lock_guard lock(state_->mutex); + state_->destroyed = true; + state_->cv.notify_all(); + } + + AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) override { + std::lock_guard lock(state_->mutex); + const AwsMetadataRequestHandle handle { state_->next_handle++ }; + state_->requests.push_back({handle, request, std::move(sink)}); + state_->cv.notify_all(); + return handle; + } + + void cancel(AwsMetadataRequestHandle handle) override { + std::lock_guard lock(state_->mutex); + state_->canceled.push_back(handle.value); + state_->cv.notify_all(); + } + + void shutdown() override { + std::lock_guard lock(state_->mutex); + state_->shut_down = true; + state_->cv.notify_all(); + } + +private: + std::shared_ptr state_; +}; + +void destroy_fake_provider(AwsMetadataProvider* provider) { + delete provider; +} + +bool wait_for_request_count( + const std::shared_ptr& state, + size_t count) { + std::unique_lock lock(state->mutex); + return state->cv.wait_for(lock, 2s, [&] { + return state->requests.size() >= count; + }); +} + +std::vector requests_copy( + const std::shared_ptr& state) { + std::lock_guard lock(state->mutex); + return state->requests; +} + +bool complete_request( + const std::shared_ptr& state, + AwsMetadataRequestKind kind, + const std::string& region, + AwsMetadataResult result, + size_t occurrence = 0) { + FakeProviderState::Pending pending; + { + std::lock_guard lock(state->mutex); + size_t found = 0; + for (const auto& item : state->requests) { + if (item.request.kind == kind && item.request.region == region) { + if (found++ == occurrence) { + pending = item; + break; + } + } + } + } + if (pending.handle.value == 0) { + return false; + } + if (auto sink = pending.sink.lock()) { + AwsMetadataCompletion completion; + completion.opaque_id = pending.request.opaque_id; + completion.generation = pending.request.generation; + completion.result = std::move(result); + sink->post(std::move(completion)); + return true; + } + return false; +} + +template +bool wait_until(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) { + return true; + } + std::this_thread::yield(); + } + return predicate(); +} + +AwsLocalityHostgroupConfig make_hostgroup( + uint32_t id, + double region_multiplier, + double az_multiplier, + uint32_t refresh, + uint32_t stale, + std::initializer_list endpoints) { + AwsLocalityHostgroupConfig config; + config.hostgroup_id = id; + config.policy.valid = true; + config.policy.same_region_multiplier = region_multiplier; + config.policy.same_az_multiplier = az_multiplier; + config.policy.refresh_interval_seconds = refresh; + config.policy.stale_ttl_seconds = stale; + for (const char* endpoint : endpoints) { + config.backends.push_back(recognize_rds_endpoint(id, endpoint, 3306)); + } + return config; +} + +const AwsLocalitySnapshotEntry* lookup( + const std::shared_ptr& snapshot, + uint32_t hostgroup_id, + const char* hostname) { + return snapshot->find(hostgroup_id, hostname, 3306); +} + +} // namespace + +int main() { + plan(41); + + auto provider_state = std::make_shared(); + ok(install_global_aws_metadata_provider( + new FakeProvider(provider_state), destroy_fake_provider, nullptr), + "metadata provider can be installed without an SDK dependency in core"); + { + auto lease = acquire_global_aws_metadata_provider(); + ok(lease && lease.get() != nullptr, + "installed provider is available through a retained lease"); + } + + std::atomic steady_seconds { 0 }; + std::atomic wall_seconds { 1700000000 }; + AwsLocalityManagerConfig manager_config; + manager_config.steady_clock = [&] { + return std::chrono::steady_clock::time_point( + std::chrono::seconds(steady_seconds.load())); + }; + manager_config.wall_clock = [&] { + return std::chrono::system_clock::time_point( + std::chrono::seconds(wall_seconds.load())); + }; + manager_config.request_timeout = 5s; + + MySQLAwsLocalityManager manager(manager_config); + const auto east_one = "db-1.abcdefghijkl.us-east-1.rds.amazonaws.com"; + const auto east_missing = "db-missing.abcdefghijkl.us-east-1.rds.amazonaws.com"; + const auto west_one = "db-2.abcdefghijkl.eu-west-1.rds.amazonaws.com"; + manager.configure({ + make_hostgroup(10, 2.0, 4.0, 300, 1800, {east_one, east_missing}), + make_hostgroup(11, 1.5, 3.0, 60, 120, {east_one, west_one}), + }); + const auto disabled_snapshot = manager.snapshot(); + const auto* disabled_entry = lookup(disabled_snapshot, 10, east_one); + ok(disabled_entry && disabled_entry->status == AwsLocalityMetadataStatus::disabled && + disabled_entry->multiplier == 1.0, + "configured manager publishes neutral diagnostic rows while disabled"); + + manager.set_enabled(true); + ok(wait_for_request_count(provider_state, 3), + "enable dispatches local identity and one coalesced request per Region"); + const auto initial_requests = requests_copy(provider_state); + const auto east_request = std::find_if(initial_requests.begin(), initial_requests.end(), + [](const auto& pending) { + return pending.request.kind == AwsMetadataRequestKind::rds_region && + pending.request.region == "us-east-1"; + }); + ok(east_request != initial_requests.end() && east_request->request.endpoints.size() == 2, + "duplicate endpoint registrations across hostgroups are coalesced"); + ok(east_request != initial_requests.end() && + east_request->request.deadline == manager_config.steady_clock() + 5s, + "provider request receives the configured monotonic deadline"); + + AwsMetadataResult local_result; + local_result.status = AwsMetadataStatus::ok; + local_result.local = {"us-east-1", "us-east-1a", "111122223333"}; + ok(complete_request(provider_state, AwsMetadataRequestKind::local_location, + "", std::move(local_result)), "local identity completion is accepted"); + + AwsMetadataResult east_result; + east_result.status = AwsMetadataStatus::ok; + east_result.endpoints.push_back({east_one, 3306, AwsEndpointType::instance, + "us-east-1", "us-east-1a", "111122223333"}); + ok(complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "us-east-1", std::move(east_result)), "east Region completion is accepted"); + + AwsMetadataResult west_result; + west_result.status = AwsMetadataStatus::ok; + west_result.endpoints.push_back({west_one, 3306, AwsEndpointType::instance, + "eu-west-1", "eu-west-1a", "111122223333"}); + ok(complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "eu-west-1", std::move(west_result)), "west Region completion is accepted"); + + ok(wait_until([&] { + auto snapshot = manager.snapshot(); + const auto* entry = lookup(snapshot, 10, east_one); + return entry != nullptr && entry->status == AwsLocalityMetadataStatus::fresh; + }), "immutable snapshot is published after normalized completions"); + + auto snapshot = manager.snapshot(); + const auto* hg10_east = lookup(snapshot, 10, east_one); + const auto* hg11_east = lookup(snapshot, 11, east_one); + const auto* hg11_west = lookup(snapshot, 11, west_one); + const auto* missing = lookup(snapshot, 10, east_missing); + ok(hg10_east && hg10_east->locality == AwsLocalityClass::same_az && + hg10_east->multiplier == 4.0, + "same-AZ result uses only the hostgroup's AZ multiplier"); + ok(hg11_east && hg11_east->multiplier == 3.0, + "one backend can use a different policy in another hostgroup"); + ok(hg11_west && hg11_west->locality == AwsLocalityClass::remote && + hg11_west->multiplier == 1.0, + "known remote backend remains neutral"); + ok(missing && missing->status == AwsLocalityMetadataStatus::error && + missing->failure_category == "endpoint_not_found" && missing->multiplier == 1.0, + "authoritative endpoint-not-found is redacted and neutral"); + ok(snapshot->effective_weight(10, east_one, 3306, 10) == 40 && + snapshot->effective_weight(11, west_one, 3306, 30) == 30, + "snapshot computes temporary weights without mutating configuration"); + + steady_seconds.store(61); + wall_seconds.store(1700000061); + manager.request_refresh(); + ok(wait_for_request_count(provider_state, 6), + "shortest configured refresh interval drives the next coalesced cycle"); + snapshot = manager.snapshot(); + hg10_east = lookup(snapshot, 10, east_one); + hg11_east = lookup(snapshot, 11, east_one); + ok(hg10_east && hg10_east->status == AwsLocalityMetadataStatus::fresh && + hg11_east && hg11_east->status == AwsLocalityMetadataStatus::stale, + "freshness is evaluated independently for each hostgroup policy"); + ok(hg11_east && hg11_east->multiplier == 3.0, + "stale but unexpired metadata remains active"); + AwsMetadataResult refresh_failure; + refresh_failure.status = AwsMetadataStatus::timeout; + const bool failed_refreshes_delivered = + complete_request(provider_state, AwsMetadataRequestKind::local_location, + "", refresh_failure, 1) && + complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "us-east-1", refresh_failure, 1) && + complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "eu-west-1", std::move(refresh_failure), 1); + const auto failed_refresh_snapshot = manager.snapshot(); + ok(failed_refreshes_delivered && + lookup(failed_refresh_snapshot, 11, east_one)->status == AwsLocalityMetadataStatus::stale, + "failed refresh preserves the last successful value within stale TTL"); + + steady_seconds.store(121); + wall_seconds.store(1700000121); + manager.request_refresh(); + ok(wait_for_request_count(provider_state, 9), + "forced scheduler wake dispatches another due refresh cycle"); + snapshot = manager.snapshot(); + hg11_east = lookup(snapshot, 11, east_one); + ok(hg11_east && hg11_east->status == AwsLocalityMetadataStatus::expired && + hg11_east->multiplier == 1.0, + "expired metadata becomes neutral"); + + const size_t canceled_before_reload = provider_state->canceled.size(); + manager.configure({ + make_hostgroup(12, 2.0, 4.0, 300, 1800, + {"db-new.abcdefghijkl.us-east-1.rds.amazonaws.com"}), + }); + ok(wait_until([&] { + std::lock_guard lock(provider_state->mutex); + return provider_state->canceled.size() > canceled_before_reload; + }), "configuration generation change cancels obsolete requests"); + const uint64_t new_generation = manager.snapshot()->generation; + ok(new_generation > snapshot->generation, + "configuration reload advances the immutable snapshot generation"); + + AwsMetadataResult late_result; + late_result.status = AwsMetadataStatus::ok; + late_result.endpoints.push_back({east_one, 3306, AwsEndpointType::instance, + "us-east-1", "us-east-1a", "111122223333"}); + ok(complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "us-east-1", std::move(late_result), 1), + "obsolete provider callback can still reach the completion sink safely"); + ok(manager.snapshot()->find(10, east_one, 3306) == nullptr, + "completion from an obsolete generation cannot repopulate removed policy state"); + + manager.set_enabled(false); + ok(wait_until([&] { + const auto current = manager.snapshot(); + const auto* entry = lookup(current, 12, + "db-new.abcdefghijkl.us-east-1.rds.amazonaws.com"); + return entry && entry->status == AwsLocalityMetadataStatus::disabled && + entry->multiplier == 1.0; + }), "disabling the master switch retains neutral diagnostic rows"); + const size_t requests_while_disabled = requests_copy(provider_state).size(); + manager.request_refresh(); + std::this_thread::yield(); + ok(requests_copy(provider_state).size() == requests_while_disabled, + "disabled manager does not dispatch provider work"); + + manager.set_enabled(true); + ok(wait_for_request_count(provider_state, requests_while_disabled + 2), + "re-enabling resumes local and regional discovery"); + manager.shutdown(); + const auto shutdown_snapshot = manager.snapshot(); + ok(lookup(shutdown_snapshot, 12, + "db-new.abcdefghijkl.us-east-1.rds.amazonaws.com") != nullptr, + "manager shutdown retains its final neutral diagnostic snapshot"); + + std::atomic global_shutdown_done { false }; + auto held_lease = acquire_global_aws_metadata_provider(); + std::thread shutdown_thread([&] { + shutdown_global_aws_metadata_provider(); + global_shutdown_done.store(true); + }); + ok(wait_until([&] { + return !acquire_global_aws_metadata_provider(); + }) && !global_shutdown_done.load(), + "global shutdown rejects new leases while waiting for an active lease"); + held_lease = {}; + shutdown_thread.join(); + ok(global_shutdown_done.load(), + "global shutdown completes after the final provider lease drains"); + ok(provider_state->shut_down && provider_state->destroyed, + "provider shutdown and destruction occur after leases drain"); + + MySQLAwsLocalityManager absent_provider_manager(manager_config); + const auto absent_endpoint = "db-absent.abcdefghijkl.us-east-1.rds.amazonaws.com"; + absent_provider_manager.configure({ + make_hostgroup(20, 2.0, 4.0, 300, 1800, {absent_endpoint}), + }); + absent_provider_manager.set_enabled(true); + ok(wait_until([&] { + const auto current = absent_provider_manager.snapshot(); + const auto* entry = lookup(current, 20, absent_endpoint); + return entry && entry->status == AwsLocalityMetadataStatus::error && + entry->failure_category == "provider_unavailable"; + }), "missing plugin provider is reported with a fixed neutral category"); + + auto replacement_state = std::make_shared(); + ok(install_global_aws_metadata_provider( + new FakeProvider(replacement_state), destroy_fake_provider, nullptr), + "provider registry accepts a replacement after complete shutdown"); + absent_provider_manager.request_refresh(); + ok(wait_for_request_count(replacement_state, 2), + "manager acquires the replacement provider on its next refresh"); + absent_provider_manager.shutdown(); + + MySQLAwsLocalityManager concurrent_manager(manager_config); + std::vector concurrent_hostgroups; + concurrent_hostgroups.reserve(100); + std::vector concurrent_regions; + std::vector concurrent_endpoints; + for (uint32_t index = 1; index <= 100; ++index) { + const std::string region = "us-test-" + std::to_string(index); + const std::string endpoint = "db-" + std::to_string(index) + + ".abcdefghijkl." + region + ".rds.amazonaws.com"; + AwsLocalityHostgroupConfig hostgroup; + hostgroup.hostgroup_id = 100 + index; + hostgroup.policy = {true, 2.0, 4.0, 300, 1800}; + hostgroup.backends.push_back(recognize_rds_endpoint( + hostgroup.hostgroup_id, endpoint, 3306)); + concurrent_regions.push_back(region); + concurrent_endpoints.push_back(endpoint); + concurrent_hostgroups.push_back(std::move(hostgroup)); + } + concurrent_manager.configure(std::move(concurrent_hostgroups)); + concurrent_manager.set_enabled(true); + ok(wait_for_request_count(replacement_state, 103), + "100-Region fixture publishes one request per distinct Region"); + AwsMetadataResult concurrent_local; + concurrent_local.status = AwsMetadataStatus::ok; + concurrent_local.local = {"us-test-1", "us-test-1a", "111122223333"}; + complete_request(replacement_state, AwsMetadataRequestKind::local_location, + "", std::move(concurrent_local), 1); + std::atomic posted_completions { 0 }; + std::vector producers; + producers.reserve(100); + for (size_t index = 0; index < 100; ++index) { + producers.emplace_back([&, index] { + AwsMetadataResult result; + result.status = AwsMetadataStatus::ok; + result.endpoints.push_back({concurrent_endpoints[index], 3306, + AwsEndpointType::instance, concurrent_regions[index], + concurrent_regions[index] + "a", "111122223333"}); + if (complete_request(replacement_state, + AwsMetadataRequestKind::rds_region, concurrent_regions[index], + std::move(result))) { + posted_completions.fetch_add(1); + } + }); + } + for (auto& producer : producers) { + producer.join(); + } + ok(posted_completions.load() == 100, + "100 concurrent completion producers all reach the shared sink"); + ok(wait_until([&] { + const auto current = concurrent_manager.snapshot(); + return current->entries.size() == 100 && + std::all_of(current->entries.begin(), current->entries.end(), + [](const auto& item) { + return item.second.status == AwsLocalityMetadataStatus::fresh; + }); + }), "concurrent completions publish one consistent immutable snapshot"); + concurrent_manager.shutdown(); + + std::mutex completion_hook_mutex; + std::condition_variable completion_hook_cv; + bool completion_hook_entered = false; + bool release_completion = false; + AwsLocalityManagerConfig blocking_config = manager_config; + blocking_config.before_completion = [&] { + std::unique_lock lock(completion_hook_mutex); + completion_hook_entered = true; + completion_hook_cv.notify_all(); + completion_hook_cv.wait(lock, [&] { return release_completion; }); + }; + MySQLAwsLocalityManager blocking_manager(blocking_config); + const auto blocking_endpoint = "db-block.abcdefghijkl.ap-block-1.rds.amazonaws.com"; + blocking_manager.configure({ + make_hostgroup(300, 2.0, 4.0, 300, 1800, {blocking_endpoint}), + }); + blocking_manager.set_enabled(true); + ok(wait_for_request_count(replacement_state, 105), + "callback-shutdown fixture dispatches local and regional requests"); + std::atomic callback_returned { false }; + std::thread callback_thread([&] { + AwsMetadataResult result; + result.status = AwsMetadataStatus::ok; + result.local = {"ap-block-1", "ap-block-1a", "111122223333"}; + complete_request(replacement_state, AwsMetadataRequestKind::local_location, + "", std::move(result), 2); + callback_returned.store(true); + }); + { + std::unique_lock lock(completion_hook_mutex); + completion_hook_cv.wait(lock, [&] { return completion_hook_entered; }); + } + std::mutex shutdown_mutex; + std::condition_variable shutdown_cv; + bool shutdown_done = false; + std::thread blocking_shutdown([&] { + blocking_manager.shutdown(); + std::lock_guard lock(shutdown_mutex); + shutdown_done = true; + shutdown_cv.notify_all(); + }); + bool shutdown_finished_early = false; + { + std::unique_lock lock(shutdown_mutex); + shutdown_finished_early = shutdown_cv.wait_for(lock, 100ms, [&] { + return shutdown_done; + }); + } + ok(!shutdown_finished_early, + "manager shutdown waits for an active completion callback"); + { + std::lock_guard lock(completion_hook_mutex); + release_completion = true; + completion_hook_cv.notify_all(); + } + callback_thread.join(); + blocking_shutdown.join(); + ok(callback_returned.load() && shutdown_done, + "callback and shutdown finish without accessing detached manager state"); + + shutdown_global_aws_metadata_provider(); + + return exit_status(); +} From 6713cab2880807b0ce9186e4901cb47380cc8718 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 20:22:24 +0000 Subject: [PATCH 06/17] feat(mysql): load AWS locality configuration --- include/Base_HostGroups_Manager.h | 6 + include/MySQL_HostGroups_Manager.h | 13 ++ include/MySQL_Thread.h | 3 + include/proxysql_structs.h | 6 + lib/Admin_FlushVariables.cpp | 11 + lib/BaseHGC.cpp | 3 + lib/MySQL_HostGroups_Manager.cpp | 77 +++++++ lib/MySQL_Thread.cpp | 12 ++ test/tap/groups/groups.json | 1 + test/tap/tests/unit/Makefile | 2 +- .../tests/unit/aws_locality_config_unit-t.cpp | 196 ++++++++++++++++++ 11 files changed, 329 insertions(+), 1 deletion(-) create mode 100644 test/tap/tests/unit/aws_locality_config_unit-t.cpp diff --git a/include/Base_HostGroups_Manager.h b/include/Base_HostGroups_Manager.h index b4734a5b87..43d778b965 100644 --- a/include/Base_HostGroups_Manager.h +++ b/include/Base_HostGroups_Manager.h @@ -19,6 +19,9 @@ class MetricsCollector; #include "proxysql.h" #include "cpp.h" #include "GTID_Server_Data.h" +#ifdef PROXYSQL40 +#include "Aws_Locality_Types.h" +#endif #include @@ -311,6 +314,9 @@ class BaseHGC { // MySQL Host Group Container char * comment; char * ignore_session_variables_text; // this is the original version (text format) of ignore_session_variables char * aws_iam_region; +#ifdef PROXYSQL40 + AwsLocalityPolicy aws_locality_policy; +#endif uint32_t max_num_online_servers; uint32_t throttle_connections_per_sec; int32_t monitor_slave_lag_when_null; diff --git a/include/MySQL_HostGroups_Manager.h b/include/MySQL_HostGroups_Manager.h index 60ff5a8ce6..79b44e0f23 100644 --- a/include/MySQL_HostGroups_Manager.h +++ b/include/MySQL_HostGroups_Manager.h @@ -2,6 +2,9 @@ #define PROXYSQL_MYSQL_HOSTGROUPS_MANAGER_H #include "proxysql.h" #include "MySQL_Backend_Auth.h" +#ifdef PROXYSQL40 +#include "Aws_Locality_Manager.h" +#endif #include "cpp.h" #include "proxysql_gtid.h" @@ -625,6 +628,9 @@ class MySQL_HostGroups_Manager : public Base_HostGroups_Manager { * present, distinguishing between 'READER' and 'WRITER' hostgroups. */ std::unordered_map> hostgroup_server_mapping; +#ifdef PROXYSQL40 + std::unique_ptr aws_locality_manager_; +#endif /** * @brief Holds the previous computed checksum for 'mysql_servers'. * @details Used to check if the servers checksums has changed during 'commit', if a change is detected, @@ -885,6 +891,13 @@ class MySQL_HostGroups_Manager : public Base_HostGroups_Manager { MySQL_HostGroups_Manager(); ~MySQL_HostGroups_Manager(); void init(); +#ifdef PROXYSQL40 + void refresh_aws_locality_configuration(); + void set_aws_locality_awareness_enabled(bool enabled); + MySQLAwsLocalityManager* aws_locality_manager() const { + return aws_locality_manager_.get(); + } +#endif #if 0 void wrlock(); void wrunlock(); diff --git a/include/MySQL_Thread.h b/include/MySQL_Thread.h index 25de6f075a..6cbf434aba 100644 --- a/include/MySQL_Thread.h +++ b/include/MySQL_Thread.h @@ -645,6 +645,9 @@ class MySQL_Threads_Handler bool passthrough_auth_empty_password; bool passthrough_auth_unknown_users; bool passthrough_auth_require_tls; +#ifdef PROXYSQL40 + bool aws_locality_awareness; +#endif int passthrough_default_hg; int passthrough_auth_cache_ttl_s; int passthrough_auth_max_inflight_probes; diff --git a/include/proxysql_structs.h b/include/proxysql_structs.h index 86d924f80b..ac7ff1cd10 100644 --- a/include/proxysql_structs.h +++ b/include/proxysql_structs.h @@ -1299,6 +1299,9 @@ __thread bool mysql_thread___passthrough_auth_enabled; __thread bool mysql_thread___passthrough_auth_empty_password; __thread bool mysql_thread___passthrough_auth_unknown_users; __thread bool mysql_thread___passthrough_auth_require_tls; +#ifdef PROXYSQL40 +__thread bool mysql_thread___aws_locality_awareness; +#endif __thread int mysql_thread___passthrough_default_hg; __thread int mysql_thread___passthrough_auth_cache_ttl_s; __thread int mysql_thread___passthrough_auth_max_inflight_probes; @@ -1656,6 +1659,9 @@ extern __thread bool mysql_thread___passthrough_auth_enabled; extern __thread bool mysql_thread___passthrough_auth_empty_password; extern __thread bool mysql_thread___passthrough_auth_unknown_users; extern __thread bool mysql_thread___passthrough_auth_require_tls; +#ifdef PROXYSQL40 +extern __thread bool mysql_thread___aws_locality_awareness; +#endif extern __thread int mysql_thread___passthrough_default_hg; extern __thread int mysql_thread___passthrough_auth_cache_ttl_s; extern __thread int mysql_thread___passthrough_auth_max_inflight_probes; diff --git a/lib/Admin_FlushVariables.cpp b/lib/Admin_FlushVariables.cpp index 6a00f448d9..d195106c7e 100644 --- a/lib/Admin_FlushVariables.cpp +++ b/lib/Admin_FlushVariables.cpp @@ -615,8 +615,19 @@ FlushVariableStats ProxySQL_Admin::flush_mysql_variables___database_to_runtime(S ASSERT_SQLITE_OK(rc, db); } } +#ifdef PROXYSQL40 + const bool aws_locality_awareness_enabled = + GloMTH->get_variable_int("aws_locality_awareness") != 0; +#endif GloMTH->wrunlock(); +#ifdef PROXYSQL40 + if (MyHGM != nullptr) { + MyHGM->set_aws_locality_awareness_enabled( + aws_locality_awareness_enabled); + } +#endif + { // NOTE: 'GloMTH->wrunlock()' should have been called before this point to avoid possible // deadlocks. See issue #3847. diff --git a/lib/BaseHGC.cpp b/lib/BaseHGC.cpp index c452c5f25e..f9f55ed1cf 100644 --- a/lib/BaseHGC.cpp +++ b/lib/BaseHGC.cpp @@ -83,6 +83,9 @@ void BaseHGC::reset_attributes() { attributes.ignore_session_variables_text = NULL; free(attributes.aws_iam_region); attributes.aws_iam_region = NULL; +#ifdef PROXYSQL40 + attributes.aws_locality_policy = {}; +#endif if (attributes.ignore_session_variables_json) { delete attributes.ignore_session_variables_json; attributes.ignore_session_variables_json = NULL; diff --git a/lib/MySQL_HostGroups_Manager.cpp b/lib/MySQL_HostGroups_Manager.cpp index a9c7216489..96907e0435 100644 --- a/lib/MySQL_HostGroups_Manager.cpp +++ b/lib/MySQL_HostGroups_Manager.cpp @@ -690,6 +690,9 @@ hg_metrics_map = std::make_tuple( ); MySQL_HostGroups_Manager::MySQL_HostGroups_Manager() { +#ifdef PROXYSQL40 + aws_locality_manager_ = std::make_unique(); +#endif status.client_connections=0; status.client_connections_prim_pass=0; status.client_connections_addl_pass=0; @@ -802,6 +805,11 @@ void MySQL_HostGroups_Manager::init() { } void MySQL_HostGroups_Manager::shutdown() { +#ifdef PROXYSQL40 + if (aws_locality_manager_) { + aws_locality_manager_->shutdown(); + } +#endif queue.add(NULL); HGCU_thread->join(); delete HGCU_thread; @@ -811,6 +819,11 @@ void MySQL_HostGroups_Manager::shutdown() { } MySQL_HostGroups_Manager::~MySQL_HostGroups_Manager() { +#ifdef PROXYSQL40 + if (aws_locality_manager_) { + aws_locality_manager_->shutdown(); + } +#endif while (MyHostGroups->len) { MyHGC *myhgc=(MyHGC *)MyHostGroups->remove_index_fast(0); delete myhgc; @@ -830,6 +843,43 @@ MySQL_HostGroups_Manager::~MySQL_HostGroups_Manager() { pthread_mutex_destroy(&lock); } +#ifdef PROXYSQL40 +void MySQL_HostGroups_Manager::refresh_aws_locality_configuration() { + std::vector hostgroups; + + wrlock(); + for (unsigned int i = 0; i < MyHostGroups->len; ++i) { + MyHGC* hostgroup = static_cast(MyHostGroups->index(i)); + if (!hostgroup->attributes.aws_locality_policy.valid) { + continue; + } + + AwsLocalityHostgroupConfig config; + config.hostgroup_id = hostgroup->hid; + config.policy = hostgroup->attributes.aws_locality_policy; + config.backends.reserve(hostgroup->mysrvs->servers->len); + for (unsigned int j = 0; j < hostgroup->mysrvs->servers->len; ++j) { + MySrvC* server = static_cast(hostgroup->mysrvs->servers->index(j)); + config.backends.emplace_back( + recognize_rds_endpoint(hostgroup->hid, server->address, server->port), + server->weight); + } + hostgroups.emplace_back(std::move(config)); + } + wrunlock(); + + if (aws_locality_manager_) { + aws_locality_manager_->configure(std::move(hostgroups)); + } +} + +void MySQL_HostGroups_Manager::set_aws_locality_awareness_enabled(bool enabled) { + if (aws_locality_manager_) { + aws_locality_manager_->set_enabled(enabled); + } +} +#endif + void MySQL_HostGroups_Manager::p_update_mysql_error_counter(p_mysql_error_type err_type, unsigned int hid, char* address, uint16_t port, unsigned int code) { p_hg_dyn_counter::metric metric = p_hg_dyn_counter::mysql_error; if (err_type == p_mysql_error_type::proxysql) { @@ -1641,6 +1691,9 @@ bool MySQL_HostGroups_Manager::commit( update_aws_rds_bgd_hosts_monitor_resultset(); wrunlock(); +#ifdef PROXYSQL40 + refresh_aws_locality_configuration(); +#endif unsigned long long curtime2=monotonic_time(); curtime1 = curtime1/1000; curtime2 = curtime2/1000; @@ -6242,11 +6295,35 @@ void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc) const uint32_t hid = myhgc->hid; free(myhgc->attributes.aws_iam_region); myhgc->attributes.aws_iam_region = NULL; +#ifdef PROXYSQL40 + myhgc->attributes.aws_locality_policy = {}; +#endif if (hostgroup_settings[0] != '\0') { try { nlohmann::json j = nlohmann::json::parse(hostgroup_settings); +#ifdef PROXYSQL40 + const auto aws = j.find("aws"); + if (aws != j.end()) { + if (!aws->is_object()) { + proxy_error("Invalid AWS locality policy field 'aws' for hostgroup %u. Value rejected.\n", hid); + } else { + const auto locality = aws->find("locality_awareness"); + if (locality != aws->end()) { + AwsLocalityPolicyError error; + myhgc->attributes.aws_locality_policy = parse_aws_locality_policy( + *locality, hid, error); + if (!myhgc->attributes.aws_locality_policy.valid) { + proxy_error( + "Invalid AWS locality policy field '%s' for hostgroup %u. Value rejected.\n", + error.field.c_str(), hid); + } + } + } + } +#endif + const auto handle_warnings_check = [](int8_t handle_warnings) -> bool { return handle_warnings == 0 || handle_warnings == 1; }; const int8_t handle_warnings = j_get_srv_default_int_val(j, hid, "handle_warnings", handle_warnings_check); myhgc->attributes.handle_warnings = handle_warnings; diff --git a/lib/MySQL_Thread.cpp b/lib/MySQL_Thread.cpp index e82bf8c97a..bae60243f8 100644 --- a/lib/MySQL_Thread.cpp +++ b/lib/MySQL_Thread.cpp @@ -511,6 +511,9 @@ static char * mysql_thread_variables_names[]= { (char *)"passthrough_auth_empty_password", (char *)"passthrough_auth_unknown_users", (char *)"passthrough_auth_require_tls", +#ifdef PROXYSQL40 + (char *)"aws_locality_awareness", +#endif (char *)"passthrough_default_hg", (char *)"passthrough_default_schema", (char *)"passthrough_auth_cache_ttl_s", @@ -1541,6 +1544,9 @@ MySQL_Threads_Handler::MySQL_Threads_Handler() { variables.passthrough_auth_empty_password = true; variables.passthrough_auth_unknown_users = false; variables.passthrough_auth_require_tls = true; +#ifdef PROXYSQL40 + variables.aws_locality_awareness = false; +#endif variables.passthrough_default_hg = 0; variables.passthrough_default_schema = strdup((char *)""); variables.passthrough_auth_cache_ttl_s = 0; @@ -2892,6 +2898,9 @@ char ** MySQL_Threads_Handler::get_variables_list() { VariablesPointers_bool["passthrough_auth_empty_password"] = make_tuple(&variables.passthrough_auth_empty_password, false); VariablesPointers_bool["passthrough_auth_unknown_users"] = make_tuple(&variables.passthrough_auth_unknown_users, false); VariablesPointers_bool["passthrough_auth_require_tls"] = make_tuple(&variables.passthrough_auth_require_tls, false); +#ifdef PROXYSQL40 + VariablesPointers_bool["aws_locality_awareness"] = make_tuple(&variables.aws_locality_awareness, false); +#endif #ifdef PROXYSQL31 VariablesPointers_bool["caching_sha2_password_auto_generate_rsa_keys"] = make_tuple(&variables.caching_sha2_password_auto_generate_rsa_keys, false); @@ -5146,6 +5155,9 @@ void MySQL_Thread::refresh_variables() { REFRESH_VARIABLE_BOOL(passthrough_auth_empty_password); REFRESH_VARIABLE_BOOL(passthrough_auth_unknown_users); REFRESH_VARIABLE_BOOL(passthrough_auth_require_tls); +#ifdef PROXYSQL40 + REFRESH_VARIABLE_BOOL(aws_locality_awareness); +#endif REFRESH_VARIABLE_INT(passthrough_default_hg); REFRESH_VARIABLE_INT(passthrough_auth_cache_ttl_s); REFRESH_VARIABLE_INT(passthrough_auth_max_inflight_probes); diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index f41ac40f12..a093c7c175 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -23,6 +23,7 @@ "aws_iam_provider_boundary_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_config_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], "basic-t" : [ "legacy-g1","mariadb10-galera-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql84-gr-g1","mysql90-g1","mysql90-gr-g1","mysql93-g1","mysql93-gr-g1","mysql95-g1","mysql95-gr-g1" ], "c_tokenizer_unit-t" : [ "unit-tests-g1" ], diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index 4d811b7b55..ef4f20d245 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -419,7 +419,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t aws_locality_config_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ diff --git a/test/tap/tests/unit/aws_locality_config_unit-t.cpp b/test/tap/tests/unit/aws_locality_config_unit-t.cpp new file mode 100644 index 0000000000..f91a90999b --- /dev/null +++ b/test/tap/tests/unit/aws_locality_config_unit-t.cpp @@ -0,0 +1,196 @@ +#include "tap.h" +#include "test_globals.h" + +#include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Thread.h" +#include "proxysql_utils.h" + +#include +#include +#include +#include + +void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc); + +namespace { + +bool contains_variable(char** variables, const char* expected) { + for (char** item = variables; item != nullptr && *item != nullptr; ++item) { + if (strcmp(*item, expected) == 0) { + return true; + } + } + return false; +} + +void free_variables(char** variables) { + if (variables == nullptr) { + return; + } + for (char** item = variables; *item != nullptr; ++item) { + free(*item); + } + free(variables); +} + +std::string capture_invalid_policy_log(MyHGC& hostgroup) { + FILE* captured = tmpfile(); + if (captured == nullptr) { + return {}; + } + fflush(stderr); + const int saved_stderr = dup(STDERR_FILENO); + if (saved_stderr < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) { + if (saved_stderr >= 0) { + close(saved_stderr); + } + fclose(captured); + return {}; + } + + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":2.0,"same_az_multiplier":"FAKE_SECRET_MULTIPLIER"}}})", + &hostgroup); + fflush(stderr); + dup2(saved_stderr, STDERR_FILENO); + close(saved_stderr); + + std::string output; + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) { + output += buffer; + } + fclose(captured); + return output; +} + +} // namespace + +int main() { + plan(21); + + MyHGC hostgroup(42); + init_myhgc_hostgroup_settings( + R"({"aws_iam_region":"us-east-1","aws":{"locality_awareness":{"same_region_multiplier":2.5,"same_az_multiplier":4.75}}})", + &hostgroup); + ok(hostgroup.attributes.aws_locality_policy.valid, + "valid nested AWS locality policy is owned by the hostgroup"); + ok(hostgroup.attributes.aws_locality_policy.same_region_multiplier == 2.5 && + hostgroup.attributes.aws_locality_policy.same_az_multiplier == 4.75, + "hostgroup retains both floating-point multipliers"); + ok(hostgroup.attributes.aws_locality_policy.refresh_interval_seconds == 300 && + hostgroup.attributes.aws_locality_policy.stale_ttl_seconds == 1800, + "hostgroup policy receives the documented timing defaults"); + ok(hostgroup.attributes.aws_iam_region != nullptr && + strcmp(hostgroup.attributes.aws_iam_region, "us-east-1") == 0, + "locality policy does not alter the independent IAM authentication Region"); + + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":1.0,"same_az_multiplier":10.0,"refresh_interval_seconds":30,"stale_ttl_seconds":604800}}})", + &hostgroup); + ok(hostgroup.attributes.aws_locality_policy.valid && + hostgroup.attributes.aws_locality_policy.refresh_interval_seconds == 30 && + hostgroup.attributes.aws_locality_policy.stale_ttl_seconds == 604800, + "inclusive policy bounds survive hostgroup parsing"); + + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":5.0,"same_az_multiplier":4.0}}})", + &hostgroup); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "invalid reload clears the previously accepted policy"); + init_myhgc_hostgroup_settings(R"({"aws":[]})", &hostgroup); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "non-object AWS settings remain disabled"); + init_myhgc_hostgroup_settings("{}", &hostgroup); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "removing locality settings removes the runtime policy"); + + const std::string diagnostics = capture_invalid_policy_log(hostgroup); + ok(diagnostics.find("hostgroup 42") != std::string::npos && + diagnostics.find("same_az_multiplier") != std::string::npos, + "invalid locality diagnostic identifies only hostgroup and field"); + ok(diagnostics.find("FAKE_SECRET_MULTIPLIER") == std::string::npos, + "invalid locality diagnostic never prints the hostgroup JSON payload"); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "rejected diagnostic case leaves no active locality policy"); + + test_globals_init(); + { + MySQL_Threads_Handler handler; + char** variables = handler.get_variables_list(); +#ifdef PROXYSQL31 + handler.set_variable("caching_sha2_password_auto_generate_rsa_keys", "false"); + handler.set_variable("caching_sha2_password_private_key_path", ""); + handler.set_variable("caching_sha2_password_public_key_path", ""); +#endif + mf_unique_ptr default_value { handler.get_variable("aws_locality_awareness") }; + ok(contains_variable(variables, "aws_locality_awareness"), + "v4 MySQL variable list exposes aws_locality_awareness"); + ok(default_value != nullptr && strcmp(default_value.get(), "false") == 0, + "aws_locality_awareness defaults to false"); + ok(handler.set_variable("aws_locality_awareness", "true") && + handler.get_variable_int("aws_locality_awareness") == 1, + "master switch accepts true"); + ok(handler.set_variable("AWS_LOCALITY_AWARENESS", "0") && + handler.get_variable_int("aws_locality_awareness") == 0, + "master switch is case-insensitive and accepts zero"); + ok(!handler.set_variable("aws_locality_awareness", "yes") && + handler.get_variable_int("aws_locality_awareness") == 0, + "master switch rejects non-boolean spelling"); + handler.set_variable("aws_locality_awareness", "1"); + ok(handler.commit().rejected_variables.empty() && + handler.get_variable_int("aws_locality_awareness") == 1, + "MySQL variable commit preserves the accepted master switch"); + free_variables(variables); + } + test_globals_cleanup(); + + auto refresh_method = &MySQL_HostGroups_Manager::refresh_aws_locality_configuration; + ok(refresh_method != nullptr, + "Hostgroup Manager exposes the post-commit locality refresh boundary"); + + GloVars.prometheus_registry = std::make_shared(); + { + MySQL_HostGroups_Manager manager; + srv_info_t info; + info.addr = "db.abcdef.us-east-1.rds.amazonaws.com"; + info.port = 3306; + info.kind = "AWS locality test"; + srv_opts_t options; + options.weigth = 7; + options.max_conns = 10; + options.use_ssl = 1; + + manager.wrlock(); + manager.create_new_server_in_hg(420, info, options); + MyHGC* configured_hostgroup = manager.MyHGC_find(420); + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":2.0,"same_az_multiplier":4.0}}})", + configured_hostgroup); + MySrvC* configured_server = static_cast( + configured_hostgroup->mysrvs->servers->index(0)); + manager.wrunlock(); + + manager.refresh_aws_locality_configuration(); + auto snapshot = manager.aws_locality_manager()->snapshot(); + const auto* entry = snapshot->find(420, info.addr, info.port); + ok(entry != nullptr && entry->configured_weight == 7, + "post-commit refresh copies hostgroup policy and backend identity into the manager"); + + manager.set_aws_locality_awareness_enabled(true); + manager.set_aws_locality_awareness_enabled(false); + ok(!manager.aws_locality_manager()->snapshot()->enabled && + configured_server->weight == 7, + "master-switch transitions never mutate the configured runtime server weight"); + + init_myhgc_hostgroup_settings("{}", configured_hostgroup); + manager.refresh_aws_locality_configuration(); + ok(manager.aws_locality_manager()->snapshot()->entries.empty(), + "removing the hostgroup policy removes its manager configuration"); + } + GloVars.prometheus_registry.reset(); + + return exit_status(); +} From e1ee397df50580556a0401702438edb0896ec037 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 20:37:59 +0000 Subject: [PATCH 07/17] feat(mysql): apply AWS locality during backend selection --- include/Aws_Locality_Manager.h | 12 + lib/Aws_Locality_Manager.cpp | 33 ++ lib/MyHGC.cpp | 78 +++- lib/MySQL_Thread.cpp | 153 +++++++ test/tap/groups/groups.json | 1 + test/tap/tests/unit/Makefile | 2 +- .../unit/aws_locality_selection_unit-t.cpp | 381 ++++++++++++++++++ 7 files changed, 656 insertions(+), 4 deletions(-) create mode 100644 test/tap/tests/unit/aws_locality_selection_unit-t.cpp diff --git a/include/Aws_Locality_Manager.h b/include/Aws_Locality_Manager.h index a3fc63ffda..45a380184d 100644 --- a/include/Aws_Locality_Manager.h +++ b/include/Aws_Locality_Manager.h @@ -5,12 +5,14 @@ #include "json_fwd.hpp" #include +#include #include #include #include #include #include #include +#include #include AwsLocalityPolicy parse_aws_locality_policy( @@ -31,6 +33,12 @@ uint64_t aws_locality_effective_weight( int64_t configured_weight, double multiplier); +uint64_t aws_locality_saturating_add(uint64_t lhs, uint64_t rhs); +size_t aws_locality_weighted_index( + const uint64_t* weights, + size_t count, + uint64_t random_value); + using AwsMetadataProviderDestroyFn = void (*)(AwsMetadataProvider*); class AwsMetadataProviderLease { @@ -83,6 +91,7 @@ struct AwsLocalitySnapshot { uint64_t generation { 0 }; bool enabled { false }; std::unordered_map entries; + std::unordered_set hostgroups; const AwsLocalitySnapshotEntry* find( uint32_t hostgroup_id, @@ -93,6 +102,9 @@ struct AwsLocalitySnapshot { std::string_view hostname, uint16_t port, int64_t configured_weight) const; + bool has_hostgroup(uint32_t hostgroup_id) const { + return hostgroups.find(hostgroup_id) != hostgroups.end(); + } }; struct AwsLocalityManagerConfig { diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp index d27cf8505e..4feff419e4 100644 --- a/lib/Aws_Locality_Manager.cpp +++ b/lib/Aws_Locality_Manager.cpp @@ -249,6 +249,38 @@ uint64_t aws_locality_effective_weight( return static_cast(product); } +uint64_t aws_locality_saturating_add(uint64_t lhs, uint64_t rhs) { + const uint64_t maximum = std::numeric_limits::max(); + return maximum - lhs < rhs ? maximum : lhs + rhs; +} + +size_t aws_locality_weighted_index( + const uint64_t* weights, + size_t count, + uint64_t random_value) { + if (weights == nullptr || count == 0) { + return count; + } + + uint64_t total = 0; + for (size_t i = 0; i < count; ++i) { + total = aws_locality_saturating_add(total, weights[i]); + } + if (total == 0) { + return count; + } + + const uint64_t target = random_value % total; + uint64_t cumulative = 0; + for (size_t i = 0; i < count; ++i) { + cumulative = aws_locality_saturating_add(cumulative, weights[i]); + if (target < cumulative) { + return i; + } + } + return count; +} + namespace { std::mutex metadata_provider_mutex; @@ -983,6 +1015,7 @@ class MySQLAwsLocalityManager::Impl { next->enabled = enabled_; const auto now = config_.steady_clock(); for (const auto& hostgroup : hostgroups_) { + next->hostgroups.insert(hostgroup.hostgroup_id); for (const auto& backend : hostgroup.backends) { auto entry = build_entry_locked(hostgroup, backend, now); next->entries.emplace(snapshot_key(entry.hostgroup_id, diff --git a/lib/MyHGC.cpp b/lib/MyHGC.cpp index 7f3e1d4dbb..6deec50afd 100644 --- a/lib/MyHGC.cpp +++ b/lib/MyHGC.cpp @@ -31,9 +31,36 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ MySrvC **mysrvcCandidates = mysrvcCandidates_static; unsigned int num_candidates = 0; bool max_connections_reached = false; + bool use_aws_locality = false; +#ifdef PROXYSQL40 + std::shared_ptr aws_locality_snapshot; + if (mysql_thread___aws_locality_awareness && MyHGM != nullptr && + MyHGM->aws_locality_manager() != nullptr) { + aws_locality_snapshot = MyHGM->aws_locality_manager()->snapshot(); + use_aws_locality = aws_locality_snapshot != nullptr && + aws_locality_snapshot->enabled && + aws_locality_snapshot->has_hostgroup(hid); + } +#endif if (l>32) { mysrvcCandidates = (MySrvC **)malloc(sizeof(MySrvC *)*l); } + auto candidate_weight_sum = [&]() -> uint64_t { +#ifdef PROXYSQL40 + if (use_aws_locality) { + uint64_t effective_sum = 0; + for (unsigned int candidate = 0; candidate < num_candidates; ++candidate) { + MySrvC* server = mysrvcCandidates[candidate]; + effective_sum = aws_locality_saturating_add( + effective_sum, + aws_locality_snapshot->effective_weight( + hid, server->address, server->port, server->weight)); + } + return effective_sum; + } +#endif + return sum; + }; if (l) { //int j=0; for (j=0; j32) { free(mysrvcCandidates); @@ -275,7 +302,7 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ unsigned int New_sum=sum; - if (New_sum==0) { + if (candidate_weight_sum()==0) { proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, "Returning MySrvC NULL because no backend ONLINE or with weight\n"); if (l>32) { free(mysrvcCandidates); @@ -327,6 +354,51 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ } } +#ifdef PROXYSQL40 + if (use_aws_locality) { + uint64_t locality_weights_static[32]; + uint64_t* locality_weights = locality_weights_static; + if (num_candidates > 32) { + locality_weights = static_cast( + malloc(sizeof(uint64_t) * num_candidates)); + } + for (j = 0; j < num_candidates; ++j) { + mysrvc = mysrvcCandidates[j]; + locality_weights[j] = aws_locality_snapshot->effective_weight( + hid, mysrvc->address, mysrvc->port, mysrvc->weight); + } + const uint64_t random_value = + (static_cast(rand_fast()) << 32) | + static_cast(rand_fast()); + const size_t selected = aws_locality_weighted_index( + locality_weights, num_candidates, random_value); + if (num_candidates > 32) { + free(locality_weights); + } + if (selected < num_candidates) { + mysrvc = mysrvcCandidates[selected]; + proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, + "Returning MySrvC %p, server %s:%d with AWS locality weighting\n", + mysrvc, mysrvc->address, mysrvc->port); + if (l>32) { + free(mysrvcCandidates); + } +#ifdef TEST_AURORA + array_mysrvc_cands += num_candidates; +#endif // TEST_AURORA + return mysrvc; + } + proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, + "Returning MySrvC NULL because AWS locality weights are zero\n"); + if (l>32) { + free(mysrvcCandidates); + } +#ifdef TEST_AURORA + array_mysrvc_cands += num_candidates; +#endif // TEST_AURORA + return NULL; + } +#endif unsigned int k; k=rand_fast()%New_sum; diff --git a/lib/MySQL_Thread.cpp b/lib/MySQL_Thread.cpp index bae60243f8..7bae61638a 100644 --- a/lib/MySQL_Thread.cpp +++ b/lib/MySQL_Thread.cpp @@ -6839,6 +6839,18 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local( std::vector parents; // this is a vector of srvers that needs to be excluded in case gtid_uuid is used MySQL_Connection *c=NULL; MySQL_Connection *client_conn = sess->client_myds->myconn; + bool use_aws_locality = false; +#ifdef PROXYSQL40 + std::shared_ptr aws_locality_snapshot; + if (mysql_thread___aws_locality_awareness && MyHGM != nullptr && + MyHGM->aws_locality_manager() != nullptr) { + aws_locality_snapshot = MyHGM->aws_locality_manager()->snapshot(); + use_aws_locality = aws_locality_snapshot != nullptr && + aws_locality_snapshot->enabled && + aws_locality_snapshot->has_hostgroup(_hid); + } +#endif + if (!use_aws_locality) { for (i=0; ilen;) { c = (MySQL_Connection *) cached_connections->index(i); const char *candidate_username = @@ -6929,6 +6941,147 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local( ++i; } return NULL; + } + +#ifdef PROXYSQL40 + struct AwsLocalityParentCandidate { + MySrvC* parent; + MySQL_Connection* connection; + uint64_t weight; + }; + AwsLocalityParentCandidate candidates_static[32]; + AwsLocalityParentCandidate* candidates = candidates_static; + const unsigned int candidate_capacity = cached_connections->len; + std::vector candidates_dynamic; + if (candidate_capacity > 32) { + candidates_dynamic.resize(candidate_capacity); + candidates = candidates_dynamic.data(); + } + unsigned int num_candidates = 0; + + for (i = 0; i < cached_connections->len;) { + c = static_cast(cached_connections->index(i)); + const char* candidate_username = + c->userinfo != nullptr ? c->userinfo->username : nullptr; + const char* requested_username = client_conn->userinfo->username; + const bool same_username = candidate_username != nullptr && + requested_username != nullptr && + strcmp(candidate_username, requested_username) == 0; + if (c->parent->myhgc->hid == _hid && same_username && + (c->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + c->requires_CHANGE_USER(client_conn, requested_type)))) { + cached_connections->remove_index_fast(i); + c->send_quit = false; + MyHGM->destroy_MyConn_from_pool(c); + continue; + } + if (c->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + c->requires_CHANGE_USER(client_conn, requested_type))) { + ++i; + continue; + } + if (!c->healthy || !c->reusable) { + ++i; + continue; + } + if (check_session_track_backoff) { + session_track_backoff_until = + c->parent->session_track_backoff_until.load(std::memory_order_relaxed); + if (session_track_backoff_until > curtime) { + ++i; + continue; + } + } + if (c->parent->myhgc->hid != _hid || + !client_conn->match_tracked_options(c)) { + ++i; + continue; + } + + MySrvC* parent = c->parent; + if (find(parents.begin(), parents.end(), parent) != parents.end()) { + ++i; + continue; + } + bool parent_already_selected = false; + for (unsigned int candidate = 0; candidate < num_candidates; ++candidate) { + if (candidates[candidate].parent == parent) { + parent_already_selected = true; + break; + } + } + if (parent_already_selected) { + ++i; + continue; + } + if (gtid_uuid != nullptr && + !MyHGM->gtid_exists(parent, gtid_uuid, gtid_trxid)) { + parents.push_back(parent); + ++i; + continue; + } + if (c->requires_CHANGE_USER(client_conn, requested_type)) { + ++i; + continue; + } + char* schema = client_conn->userinfo->schemaname; + if (strcmp(c->userinfo->schemaname, schema) != 0) { + ++i; + continue; + } + unsigned int not_match = 0; + c->number_of_matching_session_variables(client_conn, not_match); + if (not_match != 0) { + ++i; + continue; + } + if (max_lag_ms >= 0 && + static_cast(max_lag_ms) < + (parent->aws_aurora_current_lag_us / 1000)) { + status_variables.stvar[st_var_aws_aurora_replicas_skipped_during_query]++; + ++i; + continue; + } + + candidates[num_candidates++] = { + parent, + c, + aws_locality_snapshot->effective_weight( + _hid, parent->address, parent->port, parent->weight) + }; + ++i; + } + + uint64_t locality_weights_static[32]; + uint64_t* locality_weights = locality_weights_static; + std::vector locality_weights_dynamic; + if (num_candidates > 32) { + locality_weights_dynamic.resize(num_candidates); + locality_weights = locality_weights_dynamic.data(); + } + for (i = 0; i < num_candidates; ++i) { + locality_weights[i] = candidates[i].weight; + } + const uint64_t random_value = + (static_cast(rand_fast()) << 32) | + static_cast(rand_fast()); + const size_t selected = aws_locality_weighted_index( + locality_weights, num_candidates, random_value); + MySQL_Connection* selected_connection = + selected < num_candidates ? candidates[selected].connection : nullptr; + if (selected_connection == nullptr) { + return NULL; + } + for (i = 0; i < cached_connections->len; ++i) { + if (cached_connections->index(i) == selected_connection) { + return static_cast( + cached_connections->remove_index_fast(i)); + } + } +#endif + return NULL; } diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index a093c7c175..e9b55bf4e2 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -24,6 +24,7 @@ "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_config_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_selection_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], "basic-t" : [ "legacy-g1","mariadb10-galera-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql84-gr-g1","mysql90-g1","mysql90-gr-g1","mysql93-g1","mysql93-gr-g1","mysql95-g1","mysql95-gr-g1" ], "c_tokenizer_unit-t" : [ "unit-tests-g1" ], diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index ef4f20d245..1c5445f572 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -419,7 +419,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t aws_locality_config_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t aws_locality_config_unit-t aws_locality_selection_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ diff --git a/test/tap/tests/unit/aws_locality_selection_unit-t.cpp b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp new file mode 100644 index 0000000000..dfb4c03b53 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp @@ -0,0 +1,381 @@ +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "Aws_Locality_Manager.h" +#include "MySQL_Data_Stream.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Logger.hpp" +#include "MySQL_Thread.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +extern MySQL_HostGroups_Manager* MyHGM; +extern MySQL_Threads_Handler* GloMTH; +extern MySQL_Logger* GloMyLogger; + +void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc); + +namespace { + +constexpr unsigned int kHostgroup = 740; +constexpr const char* kUser = "locality_user"; +constexpr const char* kSchema = "locality_schema"; +constexpr const char* kLocal = "db-local.abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char* kRegional = "cluster-regional.abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char* kRemote = "db-remote.abcdefghijkl.eu-west-1.rds.amazonaws.com"; + +struct ProviderState { + struct Pending { + AwsMetadataRequestHandle handle; + AwsMetadataRequest request; + std::weak_ptr sink; + }; + + std::mutex mutex; + std::condition_variable cv; + std::vector pending; + uint64_t next_handle { 1 }; +}; + +class FakeProvider final : public AwsMetadataProvider { +public: + explicit FakeProvider(std::shared_ptr state) + : state_(std::move(state)) {} + + AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) override { + std::lock_guard lock(state_->mutex); + AwsMetadataRequestHandle handle { state_->next_handle++ }; + state_->pending.push_back({handle, request, std::move(sink)}); + state_->cv.notify_all(); + return handle; + } + + void cancel(AwsMetadataRequestHandle) override {} + void shutdown() override {} + +private: + std::shared_ptr state_; +}; + +void destroy_provider(AwsMetadataProvider* provider) { + delete provider; +} + +bool wait_for_requests(const std::shared_ptr& state, size_t count) { + std::unique_lock lock(state->mutex); + return state->cv.wait_for(lock, 2s, [&] { return state->pending.size() >= count; }); +} + +bool complete( + const std::shared_ptr& state, + AwsMetadataRequestKind kind, + const char* region, + AwsMetadataResult result) { + ProviderState::Pending selected; + { + std::lock_guard lock(state->mutex); + for (const auto& pending : state->pending) { + if (pending.request.kind == kind && pending.request.region == region) { + selected = pending; + break; + } + } + } + if (selected.handle.value == 0) return false; + auto sink = selected.sink.lock(); + if (!sink) return false; + AwsMetadataCompletion completion; + completion.opaque_id = selected.request.opaque_id; + completion.generation = selected.request.generation; + completion.result = std::move(result); + sink->post(std::move(completion)); + return true; +} + +template +bool wait_until(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) return true; + std::this_thread::yield(); + } + return predicate(); +} + +MySrvC* add_server(const char* hostname, int64_t weight) { + srv_info_t info; + info.addr = hostname; + info.port = 3306; + info.kind = "AWS locality selection test"; + srv_opts_t options; + options.weigth = weight; + options.max_conns = 100; + options.use_ssl = 1; + + MyHGM->wrlock(); + const int result = MyHGM->create_new_server_in_hg(kHostgroup, info, options); + MyHGC* hostgroup = MyHGM->MyHGC_find(kHostgroup); + MySrvC* server = nullptr; + if (hostgroup != nullptr) { + for (unsigned int i = 0; i < hostgroup->mysrvs->cnt(); ++i) { + MySrvC* candidate = hostgroup->mysrvs->idx(i); + if (strcmp(candidate->address, hostname) == 0) { + server = candidate; + break; + } + } + } + MyHGM->wrunlock(); + if (result != 0 || server == nullptr) { + BAIL_OUT("failed to add locality server %s", hostname); + } + return server; +} + +class SessionFixture { +public: + explicit SessionFixture(MySQL_Thread& worker) { + session = new MySQL_Session(); + session->thread = &worker; + session->connections_handler = true; + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->attach_connection(frontend); + frontend_stream->myprot.init(&frontend_stream, frontend->userinfo, session); + session->client_myds = frontend_stream; + frontend->userinfo->set( + const_cast(kUser), const_cast("password"), + const_cast(kSchema), nullptr); + frontend->set_backend_auth_type(MySQLBackendAuthType::PASSWORD); + } + + ~SessionFixture() { delete session; } + + MySQL_Session* session { nullptr }; + MySQL_Data_Stream* frontend_stream { nullptr }; + MySQL_Connection* frontend { nullptr }; +}; + +MySQL_Connection* make_connection(MySrvC* server, int fd) { + MySQL_Connection* connection = new MySQL_Connection(); + connection->mysql = mysql_init(nullptr); + if (connection->mysql == nullptr) BAIL_OUT("mysql_init failed"); + connection->ret_mysql = connection->mysql; + connection->mysql->charset = mariadb_get_charset_by_name("utf8mb4"); + connection->parent = server; + connection->userinfo->set( + const_cast(kUser), const_cast("password"), + const_cast(kSchema), nullptr); + connection->set_backend_auth_type(MySQLBackendAuthType::PASSWORD); + connection->healthy = true; + connection->reusable = true; + connection->send_quit = false; + connection->fd = fd; + connection->async_state_machine = ASYNC_IDLE; + server->ConnectionsUsed->add(connection); + return connection; +} + +} // namespace + +int main() { + plan(20); + ok(aws_locality_saturating_add( + std::numeric_limits::max() - 2, 5) == + std::numeric_limits::max(), + "locality weight sums saturate instead of overflowing"); + const uint64_t lottery_weights[] = {40, 40, 30}; + ok(aws_locality_weighted_index(lottery_weights, 3, 0) == 0 && + aws_locality_weighted_index(lottery_weights, 3, 39) == 0 && + aws_locality_weighted_index(lottery_weights, 3, 40) == 1 && + aws_locality_weighted_index(lottery_weights, 3, 79) == 1 && + aws_locality_weighted_index(lottery_weights, 3, 80) == 2, + "shared locality lottery uses exact cumulative weight boundaries"); + const uint64_t zero_weights[] = {0, 0}; + ok(aws_locality_weighted_index(zero_weights, 2, 17) == 2, + "shared locality lottery rejects an all-zero candidate set"); + ok(test_init_minimal() == 0 && test_init_auth() == 0 && test_init_query_processor() == 0 && + test_init_hostgroups() == 0, + "minimal worker and Hostgroup Manager fixtures initialize"); + GloMyLogger = new MySQL_Logger(); + + auto provider_state = std::make_shared(); + ok(install_global_aws_metadata_provider( + new FakeProvider(provider_state), destroy_provider, nullptr), + "fake metadata provider installs through the production lease registry"); + + MySrvC* local = add_server(kLocal, 10); + MySrvC* regional = add_server(kRegional, 20); + MySrvC* remote = add_server(kRemote, 30); + MyHGC* hostgroup = MyHGM->MyHGC_find(kHostgroup); + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":2.0,"same_az_multiplier":4.0}}})", + hostgroup); + MyHGM->refresh_aws_locality_configuration(); + MyHGM->set_aws_locality_awareness_enabled(true); + ok(wait_for_requests(provider_state, 3), + "selection fixture requests local identity and both backend Regions"); + + AwsMetadataResult local_result; + local_result.status = AwsMetadataStatus::ok; + local_result.local = {"us-east-1", "us-east-1a", "111122223333"}; + ok(complete(provider_state, AwsMetadataRequestKind::local_location, "", + std::move(local_result)), "local identity completion is accepted"); + + AwsMetadataResult east; + east.status = AwsMetadataStatus::ok; + east.endpoints.push_back({kLocal, 3306, AwsEndpointType::instance, + "us-east-1", "us-east-1a", "111122223333"}); + east.endpoints.push_back({kRegional, 3306, AwsEndpointType::cluster, + "us-east-1", "us-east-1a", "111122223333"}); + ok(complete(provider_state, AwsMetadataRequestKind::rds_region, "us-east-1", + std::move(east)), "same-Region endpoint completion is accepted"); + + AwsMetadataResult west; + west.status = AwsMetadataStatus::ok; + west.endpoints.push_back({kRemote, 3306, AwsEndpointType::instance, + "eu-west-1", "eu-west-1a", "111122223333"}); + ok(complete(provider_state, AwsMetadataRequestKind::rds_region, "eu-west-1", + std::move(west)), "remote endpoint completion is accepted"); + + ok(wait_until([&] { + auto snapshot = MyHGM->aws_locality_manager()->snapshot(); + const auto* entry = snapshot->find(kHostgroup, kRemote, 3306); + return entry != nullptr && entry->status == AwsLocalityMetadataStatus::fresh; + }), "fresh immutable selection snapshot is published"); + + auto snapshot = MyHGM->aws_locality_manager()->snapshot(); + ok(snapshot->effective_weight(kHostgroup, kLocal, 3306, 10) == 40 && + snapshot->effective_weight(kHostgroup, kRegional, 3306, 20) == 40 && + snapshot->effective_weight(kHostgroup, kRemote, 3306, 30) == 30, + "literal same-AZ, same-Region, and remote weights evaluate to 40/40/30"); + const auto* regional_entry = snapshot->find(kHostgroup, kRegional, 3306); + ok(regional_entry != nullptr && + regional_entry->locality == AwsLocalityClass::same_region && + regional_entry->multiplier == 2.0, + "cluster endpoints receive Region bias only, even when their reported AZ matches"); + + GloMTH->set_variable("aws_locality_awareness", "true"); + GloMTH->num_threads = 1; + { + MySQL_Thread worker; + if (!worker.init()) BAIL_OUT("worker init failed"); + worker.curtime = 10000000; + SessionFixture session(worker); + + int global_counts[3] = {0, 0, 0}; + for (int i = 0; i < 12000; ++i) { + MySrvC* selected = hostgroup->get_random_MySrvC(nullptr, 0, -1, session.session); + if (selected == local) ++global_counts[0]; + else if (selected == regional) ++global_counts[1]; + else if (selected == remote) ++global_counts[2]; + } + const double local_share = static_cast(global_counts[0]) / 12000.0; + const double regional_share = static_cast(global_counts[1]) / 12000.0; + const double remote_share = static_cast(global_counts[2]) / 12000.0; + ok(local_share > 0.31 && local_share < 0.42 && + regional_share > 0.31 && regional_share < 0.42 && + remote_share > 0.20 && remote_share < 0.34, + "global server lottery follows effective 40:40:30 weights (%.3f/%.3f/%.3f)", + local_share, regional_share, remote_share); + ok(local->weight == 10 && regional->weight == 20 && remote->weight == 30, + "global locality selection never mutates configured server weights"); + + MySQL_Connection* local_connection = make_connection(local, 100); + std::vector remote_connections; + for (int i = 0; i < 5; ++i) { + remote_connections.push_back(make_connection(remote, 200 + i)); + } + worker.push_MyConn_local(local_connection); + for (auto* connection : remote_connections) worker.push_MyConn_local(connection); + + int local_parent = 0; + int remote_parent = 0; + for (int i = 0; i < 6000; ++i) { + MySQL_Connection* selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + if (selected == nullptr) BAIL_OUT("local selection returned no candidate"); + if (selected->parent == local) ++local_parent; + else if (selected->parent == remote) ++remote_parent; + worker.push_MyConn_local(selected); + } + const double local_parent_share = static_cast(local_parent) / 6000.0; + ok(local_parent_share > 0.53 && local_parent_share < 0.62 && + local_parent + remote_parent == 6000, + "local cache chooses 40:30 weighted parents, not the 1:5 connection count (%.3f local)", + local_parent_share); + + remote->aws_aurora_current_lag_us = 5000; + MySQL_Connection* selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, 1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == local, + "replication-lag eligibility excludes a remote parent before locality weighting"); + worker.push_MyConn_local(selected); + remote->aws_aurora_current_lag_us = 0; + + MySQL_Connection* incompatible = make_connection(regional, 300); + incompatible->userinfo->set( + const_cast("other_user"), const_cast("password"), + const_cast(kSchema), nullptr); + worker.push_MyConn_local(incompatible); + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent != regional, + "authentication incompatibility excludes a parent before locality weighting"); + worker.push_MyConn_local(selected); + + local_connection->healthy = false; + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == remote, + "health eligibility excludes a preferred local connection before locality weighting"); + worker.push_MyConn_local(selected); + local_connection->healthy = true; + + local->session_track_backoff_until.store(worker.curtime + 1, std::memory_order_relaxed); + mysql_thread___session_track_variables = session_track_variables::ENFORCED; + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == remote, + "session-capability backoff excludes a local parent before locality weighting"); + worker.push_MyConn_local(selected); + mysql_thread___session_track_variables = session_track_variables::DISABLED; + local->session_track_backoff_until.store(0, std::memory_order_relaxed); + + local_connection->options.client_flag |= CLIENT_FOUND_ROWS; + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == remote, + "session option incompatibility is enforced before locality weighting"); + worker.push_MyConn_local(selected); + local_connection->options.client_flag &= ~CLIENT_FOUND_ROWS; + } + MyHGM->set_aws_locality_awareness_enabled(false); + test_cleanup_hostgroups(); + shutdown_global_aws_metadata_provider(); + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_query_processor(); + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +} From 2d9a77c25ec73fb6dd120b24b19ea8ac808f8db9 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 20:59:06 +0000 Subject: [PATCH 08/17] feat(aws): provide asynchronous locality metadata --- include/ProxySQL_Plugin.h | 12 +- lib/ProxySQL_PluginManager.cpp | 13 + plugins/aws/src/aws_locality_provider.cpp | 800 ++++++++++++++++++ plugins/aws/src/aws_locality_provider.h | 218 +++++ src/main.cpp | 9 + test/tap/tests/unit/Makefile | 7 + .../tests/unit/aws_locality_plugin_unit-t.cpp | 411 +++++++++ 7 files changed, 1468 insertions(+), 2 deletions(-) create mode 100644 plugins/aws/src/aws_locality_provider.cpp create mode 100644 plugins/aws/src/aws_locality_provider.h create mode 100644 test/tap/tests/unit/aws_locality_plugin_unit-t.cpp diff --git a/include/ProxySQL_Plugin.h b/include/ProxySQL_Plugin.h index 5226a0ba1c..527bf03b10 100644 --- a/include/ProxySQL_Plugin.h +++ b/include/ProxySQL_Plugin.h @@ -15,6 +15,7 @@ class SQLite3DB; class SQLite3_result; class AwsIamTokenSource; +class AwsMetadataProvider; namespace prometheus { class Registry; } // Descriptor ABI version the plugin was compiled for. Plugins set @@ -40,8 +41,10 @@ namespace prometheus { class Registry; } // trailing field — matching the pre-ABI-4 behaviour. // ABI 5: ProxySQL_PluginServices gains AWS IAM provider installation and // sizing callbacks. They are live only during normal plugin init. -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 5u; -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 5u; +// ABI 6: ProxySQL_PluginServices gains the general AWS metadata-provider +// installation callback used by locality discovery. +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 6u; +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 6u; enum class ProxySQL_PluginDBKind : uint8_t { admin_db = 0, @@ -242,6 +245,9 @@ using proxysql_plugin_install_aws_iam_token_source_cb = using proxysql_plugin_get_aws_iam_limits_cb = void (*)(size_t *max_total_waiters, size_t *max_waiters_per_key); + +using proxysql_plugin_install_aws_metadata_provider_cb = + bool (*)(AwsMetadataProvider *, void (*)(AwsMetadataProvider *), void *module_handle); #endif /* PROXYSQL40 */ // Services provided to plugins across the four-phase lifecycle. @@ -308,6 +314,8 @@ struct ProxySQL_PluginServices { // ABI-5 extension. Both are null outside plugin init(). proxysql_plugin_install_aws_iam_token_source_cb install_aws_iam_token_source; proxysql_plugin_get_aws_iam_limits_cb get_aws_iam_limits; + // ABI-6 extension. Null outside normal plugin init(). + proxysql_plugin_install_aws_metadata_provider_cb install_aws_metadata_provider; #endif /* PROXYSQL40 */ }; diff --git a/lib/ProxySQL_PluginManager.cpp b/lib/ProxySQL_PluginManager.cpp index 2abb4c2a8b..5aa0a20f00 100644 --- a/lib/ProxySQL_PluginManager.cpp +++ b/lib/ProxySQL_PluginManager.cpp @@ -6,6 +6,7 @@ #include "ProxySQL_PluginManager.h" #include "Aws_Iam_Provider.h" +#include "Aws_Locality_Manager.h" #include "MySQL_Thread.h" #include @@ -196,6 +197,17 @@ void get_aws_iam_limits_service(size_t *max_total_waiters, size_t *max_waiters_p if (max_total_waiters != nullptr) *max_total_waiters = maximum; if (max_waiters_per_key != nullptr) *max_waiters_per_key = maximum; } + +bool install_aws_metadata_provider_service( + AwsMetadataProvider *provider, + void (*destroy)(AwsMetadataProvider *), + void *module_handle) { + if (g_registry_target == nullptr) { + proxy_warning("AWS metadata provider installation attempted outside plugin init phase\n"); + return false; + } + return install_global_aws_metadata_provider(provider, destroy, module_handle); +} #endif /* PROXYSQL40 */ SQLite3DB* get_admindb_service() { @@ -324,6 +336,7 @@ ProxySQL_PluginManager::ProxySQL_PluginManager() { services_.register_runtime_view = ®ister_runtime_view_service; services_.install_aws_iam_token_source = &install_aws_iam_token_source_service; services_.get_aws_iam_limits = &get_aws_iam_limits_service; + services_.install_aws_metadata_provider = &install_aws_metadata_provider_service; // Phase-B (register_schemas) services: same layout as init(), but DB // handle getters and the query-hook registrar are stubbed -- see the diff --git a/plugins/aws/src/aws_locality_provider.cpp b/plugins/aws/src/aws_locality_provider.cpp new file mode 100644 index 0000000000..ce4416528e --- /dev/null +++ b/plugins/aws/src/aws_locality_provider.cpp @@ -0,0 +1,800 @@ +#include "aws_locality_provider.h" + +#include "json.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef PROXYSQL_AWS_SDK_PROVIDER +#include "curl/curl.h" + +#include +#include +#include +#include +#include +#include +#include +#endif + +using nlohmann::json; + +namespace { + +const char* fixed_failure_category(AwsMetadataStatus status) { + switch (status) { + case AwsMetadataStatus::ok: return ""; + case AwsMetadataStatus::provider_unavailable: return "provider_unavailable"; + case AwsMetadataStatus::access_denied: return "access_denied"; + case AwsMetadataStatus::throttled: return "throttled"; + case AwsMetadataStatus::imds_unavailable: return "imds_unavailable"; + case AwsMetadataStatus::timeout: return "timeout"; + case AwsMetadataStatus::cancelled: return "cancelled"; + case AwsMetadataStatus::invalid_response: return "invalid_response"; + case AwsMetadataStatus::shutdown: return "cancelled"; + } + return "invalid_response"; +} + +AwsMetadataResult fixed_failure(AwsMetadataStatus status) { + AwsMetadataResult result; + result.status = status; + result.failure_category = fixed_failure_category(status); + return result; +} + +void normalize_failure(AwsMetadataResult& result) { + result.failure_category = fixed_failure_category(result.status); + if (result.status != AwsMetadataStatus::ok) { + result.local = {}; + result.endpoints.clear(); + } +} + +std::string normalized_hostname(const std::string& input) { + if (input.empty()) return {}; + std::string hostname = input; + if (hostname.back() == '.') hostname.pop_back(); + if (hostname.empty() || hostname.back() == '.') return {}; + for (char& character : hostname) { + const unsigned char value = static_cast(character); + if (!(std::isalnum(value) || character == '-' || character == '.')) return {}; + character = static_cast(std::tolower(value)); + } + return hostname; +} + +bool valid_location_value(const std::string& value, size_t maximum) { + if (value.empty() || value.size() > maximum) return false; + for (const unsigned char character : value) { + if (!(std::isalnum(character) || character == '-')) return false; + } + return true; +} + +bool valid_account_id(const std::string& value) { + return value.size() == 12 && std::all_of(value.begin(), value.end(), + [](unsigned char character) { return std::isdigit(character); }); +} + +std::string account_from_rds_arn( + const std::string& arn, + const std::string& expected_region) { + size_t begin = 0; + std::string fields[6]; + for (size_t field = 0; field < 5; ++field) { + const size_t end = arn.find(':', begin); + if (end == std::string::npos) return {}; + fields[field] = arn.substr(begin, end - begin); + begin = end + 1; + } + fields[5] = arn.substr(begin); + if (fields[0] != "arn" || fields[2] != "rds" || + fields[3] != expected_region || !valid_account_id(fields[4])) { + return {}; + } + return fields[4]; +} + +AwsEndpointType endpoint_type(const std::string& input) { + std::string value; + value.reserve(input.size()); + for (const unsigned char character : input) { + value.push_back(static_cast(std::toupper(character))); + } + if (value == "WRITER") return AwsEndpointType::cluster; + if (value == "READER") return AwsEndpointType::reader; + if (value == "CUSTOM") return AwsEndpointType::custom; + return AwsEndpointType::unknown; +} + +bool expired( + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled, + AwsMetadataResult& result) { + if (cancelled()) { + result = fixed_failure(AwsMetadataStatus::cancelled); + return true; + } + if (std::chrono::steady_clock::now() >= deadline) { + result = fixed_failure(AwsMetadataStatus::timeout); + return true; + } + return false; +} + +} // namespace + +class AwsSdkMetadataProvider::Impl { +public: + struct Job { + AwsMetadataRequestHandle handle; + AwsMetadataRequest request; + std::weak_ptr sink; + std::atomic cancelled { false }; + }; + + Impl( + std::shared_ptr backend, + AwsMetadataProviderConfig config) + : backend_(std::move(backend)), config_(std::move(config)) { + if (config_.worker_count == 0) config_.worker_count = 1; + if (config_.worker_count > 2) config_.worker_count = 2; + if (config_.max_pending < config_.worker_count) { + config_.max_pending = config_.worker_count; + } + workers_.reserve(config_.worker_count); + for (size_t i = 0; i < config_.worker_count; ++i) { + workers_.emplace_back([this] { worker_loop(); }); + } + } + + ~Impl() { shutdown(); } + + AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) { + AwsMetadataResult immediate; + bool deliver = false; + { + std::lock_guard lock(mutex_); + if (stopping_ || backend_ == nullptr) return {}; + if (config_.steady_clock() >= request.deadline) { + immediate = fixed_failure(AwsMetadataStatus::timeout); + deliver = true; + } else if (jobs_.size() >= config_.max_pending) { + immediate = fixed_failure(AwsMetadataStatus::throttled); + deliver = true; + } else { + auto job = std::make_shared(); + job->handle.value = next_handle_++; + job->request = request; + job->sink = std::move(sink); + jobs_.emplace(job->handle.value, job); + queue_.push_back(job); + cv_.notify_one(); + return job->handle; + } + } + if (deliver) deliver_immediate(request, std::move(sink), std::move(immediate)); + return {}; + } + + void cancel(AwsMetadataRequestHandle handle) { + if (handle.value == 0) return; + std::lock_guard lock(mutex_); + const auto found = jobs_.find(handle.value); + if (found != jobs_.end()) found->second->cancelled.store(true); + cv_.notify_all(); + } + + void shutdown() { + { + std::unique_lock lock(mutex_); + if (shutdown_complete_) return; + if (shutdown_started_) { + cv_.wait(lock, [&] { return shutdown_complete_; }); + return; + } + shutdown_started_ = true; + stopping_ = true; + for (const auto& item : jobs_) item.second->cancelled.store(true); + cv_.notify_all(); + } + for (auto& worker : workers_) { + if (worker.joinable()) worker.join(); + } + { + std::unique_lock lock(mutex_); + cv_.wait(lock, [&] { return active_callbacks_ == 0; }); + jobs_.clear(); + queue_.clear(); + shutdown_complete_ = true; + cv_.notify_all(); + } + } + +private: + void deliver_immediate( + const AwsMetadataRequest& request, + std::weak_ptr weak_sink, + AwsMetadataResult result) { + auto sink = weak_sink.lock(); + if (!sink) return; + { + std::lock_guard lock(mutex_); + if (stopping_) return; + ++active_callbacks_; + } + AwsMetadataCompletion completion; + completion.opaque_id = request.opaque_id; + completion.generation = request.generation; + completion.result = std::move(result); + try { + sink->post(std::move(completion)); + } catch (...) { + // Plugin callbacks must not escape across the provider ABI boundary. + } + { + std::lock_guard lock(mutex_); + --active_callbacks_; + cv_.notify_all(); + } + } + + void worker_loop() { + for (;;) { + std::shared_ptr job; + { + std::unique_lock lock(mutex_); + cv_.wait(lock, [&] { return stopping_ || !queue_.empty(); }); + if (stopping_ && queue_.empty()) return; + job = queue_.front(); + queue_.pop_front(); + } + + AwsMetadataResult result; + if (job->cancelled.load()) { + result = fixed_failure(AwsMetadataStatus::cancelled); + } else if (config_.steady_clock() >= job->request.deadline) { + result = fixed_failure(AwsMetadataStatus::timeout); + } else { + try { + result = backend_->discover(job->request, + [job, this] { + return job->cancelled.load() || stopping_.load(); + }); + } catch (...) { + result = fixed_failure(AwsMetadataStatus::provider_unavailable); + } + if (config_.steady_clock() >= job->request.deadline) { + result = fixed_failure(AwsMetadataStatus::timeout); + } + } + normalize_failure(result); + + std::shared_ptr sink; + { + std::lock_guard lock(mutex_); + jobs_.erase(job->handle.value); + if (!stopping_ && !job->cancelled.load()) { + sink = job->sink.lock(); + if (sink) ++active_callbacks_; + } + } + if (sink) { + AwsMetadataCompletion completion; + completion.opaque_id = job->request.opaque_id; + completion.generation = job->request.generation; + completion.result = std::move(result); + try { + sink->post(std::move(completion)); + } catch (...) { + // Keep worker and shutdown bookkeeping intact on a bad consumer. + } + std::lock_guard lock(mutex_); + --active_callbacks_; + cv_.notify_all(); + } + } + } + + std::shared_ptr backend_; + AwsMetadataProviderConfig config_; + std::mutex mutex_; + std::condition_variable cv_; + std::deque> queue_; + std::unordered_map> jobs_; + std::vector workers_; + std::atomic stopping_ { false }; + bool shutdown_started_ { false }; + bool shutdown_complete_ { false }; + uint64_t next_handle_ { 1 }; + size_t active_callbacks_ { 0 }; +}; + +AwsSdkMetadataProvider::AwsSdkMetadataProvider( + std::shared_ptr backend, + AwsMetadataProviderConfig config) + : impl_(new Impl(std::move(backend), std::move(config))) {} + +AwsSdkMetadataProvider::~AwsSdkMetadataProvider() = default; + +AwsMetadataRequestHandle AwsSdkMetadataProvider::request( + const AwsMetadataRequest& request, + std::weak_ptr sink) { + return impl_->request(request, std::move(sink)); +} + +void AwsSdkMetadataProvider::cancel(AwsMetadataRequestHandle handle) { + impl_->cancel(handle); +} + +void AwsSdkMetadataProvider::shutdown() { + impl_->shutdown(); +} + +AwsLocalLocation aws_locality_environment_location( + const AwsLocalityEnvironmentGetter& getenv_value) { + AwsLocalLocation location; + if (!getenv_value) return location; + location.region = getenv_value("AWS_REGION"); + if (!valid_location_value(location.region, 64)) { + location.region = getenv_value("AWS_DEFAULT_REGION"); + } + if (!valid_location_value(location.region, 64)) { + location.region.clear(); + return location; + } + location.availability_zone = getenv_value("AWS_AVAILABILITY_ZONE"); + if (!valid_location_value(location.availability_zone, 64)) { + location.availability_zone.clear(); + } + location.account_id = getenv_value("AWS_ACCOUNT_ID"); + if (!valid_account_id(location.account_id)) location.account_id.clear(); + return location; +} + +AwsLocalityLocalDiscovery::AwsLocalityLocalDiscovery( + std::shared_ptr transport, + AwsLocalityEnvironmentGetter getenv_value) + : transport_(std::move(transport)), getenv_value_(std::move(getenv_value)) {} + +AwsMetadataResult AwsLocalityLocalDiscovery::discover( + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) const { + AwsMetadataResult result; + if (expired(deadline, cancelled, result)) return result; + + auto fallback = [&](AwsMetadataStatus failure_status) { + AwsMetadataResult fallback_result; + fallback_result.local = aws_locality_environment_location(getenv_value_); + if (!fallback_result.local.region.empty()) { + fallback_result.status = AwsMetadataStatus::ok; + return fallback_result; + } + return fixed_failure(failure_status); + }; + + if (!transport_) return fallback(AwsMetadataStatus::imds_unavailable); + AwsImdsResponse token = transport_->put_token(deadline, cancelled); + if (expired(deadline, cancelled, result)) return result; + if (!token.transport_ok || token.status_code != 200 || token.body.empty() || + token.body.size() > 4096) { + if (!token.body.empty()) OPENSSL_cleanse(&token.body[0], token.body.size()); + return fallback(AwsMetadataStatus::imds_unavailable); + } + + AwsImdsResponse document = transport_->get_identity_document( + token.body, deadline, cancelled); + OPENSSL_cleanse(&token.body[0], token.body.size()); + if (expired(deadline, cancelled, result)) return result; + if (!document.transport_ok || document.status_code != 200) { + return fallback(AwsMetadataStatus::imds_unavailable); + } + if (document.body.empty() || document.body.size() > 16384) { + return fallback(AwsMetadataStatus::invalid_response); + } + + try { + const json identity = json::parse(document.body); + if (!identity.is_object() || !identity.contains("region") || + !identity["region"].is_string()) { + return fallback(AwsMetadataStatus::invalid_response); + } + result.local.region = identity["region"].get(); + if (!valid_location_value(result.local.region, 64)) { + return fallback(AwsMetadataStatus::invalid_response); + } + if (identity.contains("availabilityZone") && + identity["availabilityZone"].is_string()) { + result.local.availability_zone = + identity["availabilityZone"].get(); + if (!valid_location_value(result.local.availability_zone, 64)) { + result.local.availability_zone.clear(); + } + } + if (identity.contains("accountId") && identity["accountId"].is_string()) { + result.local.account_id = identity["accountId"].get(); + if (!valid_account_id(result.local.account_id)) result.local.account_id.clear(); + } + result.status = AwsMetadataStatus::ok; + return result; + } catch (...) { + return fallback(AwsMetadataStatus::invalid_response); + } +} + +AwsLocalityRdsDiscovery::AwsLocalityRdsDiscovery( + std::shared_ptr api) + : api_(std::move(api)) {} + +AwsMetadataResult AwsLocalityRdsDiscovery::discover( + const AwsMetadataRequest& request, + const AwsLocalityCancelPredicate& cancelled) const { + AwsMetadataResult result; + if (request.region.empty() || !api_) { + return fixed_failure(AwsMetadataStatus::invalid_response); + } + if (expired(request.deadline, cancelled, result)) return result; + result.status = AwsMetadataStatus::ok; + + std::unordered_map endpoint_indices; + auto append = [&](const std::string& hostname_input, int port, + AwsEndpointType type, const std::string& az, const std::string& account) { + const std::string hostname = normalized_hostname(hostname_input); + if (hostname.empty() || type == AwsEndpointType::unknown || + port < 0 || port > 65535) return; + AwsMetadataEndpoint endpoint; + endpoint.hostname = hostname; + endpoint.port = static_cast(port); + endpoint.endpoint_type = type; + endpoint.region = request.region; + endpoint.availability_zone = az; + endpoint.account_id = account; + const std::string key = hostname + "\n" + std::to_string(port); + const auto found = endpoint_indices.find(key); + if (found == endpoint_indices.end()) { + endpoint_indices.emplace(key, result.endpoints.size()); + result.endpoints.push_back(std::move(endpoint)); + } else { + result.endpoints[found->second] = std::move(endpoint); + } + }; + + auto fail = [&](AwsMetadataStatus status) { + result = fixed_failure(status); + return result; + }; + auto check_page = [&](AwsMetadataStatus status) { + if (cancelled()) return AwsMetadataStatus::cancelled; + if (std::chrono::steady_clock::now() >= request.deadline) { + return AwsMetadataStatus::timeout; + } + return status; + }; + + std::set markers; + std::string marker; + for (;;) { + AwsRdsInstancesPage page = api_->describe_instances( + request.region, marker, request.deadline, cancelled); + const AwsMetadataStatus status = check_page(page.status); + if (status != AwsMetadataStatus::ok) return fail(status); + for (const auto& instance : page.instances) { + if (instance.port <= 0 || instance.port > 65535) continue; + append(instance.endpoint, instance.port, AwsEndpointType::instance, + instance.availability_zone, + account_from_rds_arn(instance.arn, request.region)); + } + if (page.next_marker.empty()) break; + if (!markers.insert(page.next_marker).second) { + return fail(AwsMetadataStatus::invalid_response); + } + marker = std::move(page.next_marker); + } + + std::unordered_map cluster_accounts; + markers.clear(); + marker.clear(); + for (;;) { + AwsRdsClustersPage page = api_->describe_clusters( + request.region, marker, request.deadline, cancelled); + const AwsMetadataStatus status = check_page(page.status); + if (status != AwsMetadataStatus::ok) return fail(status); + for (const auto& cluster : page.clusters) { + const std::string account = account_from_rds_arn(cluster.arn, request.region); + if (!cluster.identifier.empty()) cluster_accounts[cluster.identifier] = account; + if (cluster.port > 0 && cluster.port <= 65535) { + append(cluster.endpoint, cluster.port, AwsEndpointType::cluster, {}, account); + append(cluster.reader_endpoint, cluster.port, AwsEndpointType::reader, {}, account); + } + for (const auto& custom : cluster.custom_endpoints) { + append(custom, 0, AwsEndpointType::custom, {}, account); + } + } + if (page.next_marker.empty()) break; + if (!markers.insert(page.next_marker).second) { + return fail(AwsMetadataStatus::invalid_response); + } + marker = std::move(page.next_marker); + } + + markers.clear(); + marker.clear(); + for (;;) { + AwsRdsClusterEndpointsPage page = api_->describe_cluster_endpoints( + request.region, marker, request.deadline, cancelled); + const AwsMetadataStatus status = check_page(page.status); + if (status != AwsMetadataStatus::ok) return fail(status); + for (const auto& endpoint : page.endpoints) { + const auto account = cluster_accounts.find(endpoint.cluster_identifier); + append(endpoint.endpoint, 0, endpoint_type(endpoint.endpoint_type), {}, + account == cluster_accounts.end() ? std::string() : account->second); + } + if (page.next_marker.empty()) break; + if (!markers.insert(page.next_marker).second) { + return fail(AwsMetadataStatus::invalid_response); + } + marker = std::move(page.next_marker); + } + + return result; +} + +AwsLocalityCompositeDiscovery::AwsLocalityCompositeDiscovery( + std::shared_ptr imds, + AwsLocalityEnvironmentGetter getenv_value, + std::shared_ptr rds) + : local_(std::move(imds), std::move(getenv_value)), + rds_(std::move(rds)) {} + +AwsMetadataResult AwsLocalityCompositeDiscovery::discover( + const AwsMetadataRequest& request, + const AwsLocalityCancelPredicate& cancelled) { + if (request.kind == AwsMetadataRequestKind::local_location) { + return local_.discover(request.deadline, cancelled); + } + if (request.kind == AwsMetadataRequestKind::rds_region) { + return rds_.discover(request, cancelled); + } + return fixed_failure(AwsMetadataStatus::invalid_response); +} + +#ifdef PROXYSQL_AWS_SDK_PROVIDER +namespace { + +struct CurlResponseContext { + std::string* body; + size_t maximum; + bool overflow { false }; +}; + +size_t imds_write_callback(char* data, size_t size, size_t count, void* opaque) { + auto* context = static_cast(opaque); + if (size != 0 && count > context->maximum / size) { + context->overflow = true; + return 0; + } + const size_t bytes = size * count; + if (bytes > context->maximum - std::min(context->maximum, context->body->size())) { + context->overflow = true; + return 0; + } + context->body->append(data, bytes); + return bytes; +} + +struct CurlProgressContext { + std::chrono::steady_clock::time_point deadline; + const AwsLocalityCancelPredicate* cancelled; +}; + +int imds_progress_callback(void* opaque, curl_off_t, curl_off_t, curl_off_t, curl_off_t) { + auto* context = static_cast(opaque); + return (*(context->cancelled))() || + std::chrono::steady_clock::now() >= context->deadline; +} + +AwsImdsResponse imds_request( + const char* url, + const char* method, + struct curl_slist* headers, + size_t maximum, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) { + AwsImdsResponse response; + if (cancelled() || std::chrono::steady_clock::now() >= deadline) return response; + CURL* handle = curl_easy_init(); + if (handle == nullptr) return response; + CurlResponseContext write_context {&response.body, maximum}; + CurlProgressContext progress_context {deadline, &cancelled}; + const auto remaining = std::chrono::duration_cast( + deadline - std::chrono::steady_clock::now()).count(); + const long timeout = static_cast(std::max(1, remaining)); + curl_easy_setopt(handle, CURLOPT_URL, url); + curl_easy_setopt(handle, CURLOPT_CUSTOMREQUEST, method); + curl_easy_setopt(handle, CURLOPT_HTTPHEADER, headers); + curl_easy_setopt(handle, CURLOPT_NOBODY, 0L); + curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 0L); + curl_easy_setopt(handle, CURLOPT_NOPROXY, "*"); + curl_easy_setopt(handle, CURLOPT_PROXY, ""); + curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, std::min(timeout, 500L)); + curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, timeout); + curl_easy_setopt(handle, CURLOPT_NOSIGNAL, 1L); + curl_easy_setopt(handle, CURLOPT_WRITEFUNCTION, &imds_write_callback); + curl_easy_setopt(handle, CURLOPT_WRITEDATA, &write_context); + curl_easy_setopt(handle, CURLOPT_XFERINFOFUNCTION, &imds_progress_callback); + curl_easy_setopt(handle, CURLOPT_XFERINFODATA, &progress_context); + curl_easy_setopt(handle, CURLOPT_NOPROGRESS, 0L); + const CURLcode code = curl_easy_perform(handle); + curl_easy_getinfo(handle, CURLINFO_RESPONSE_CODE, &response.status_code); + curl_easy_cleanup(handle); + response.transport_ok = code == CURLE_OK && !write_context.overflow; + if (!response.transport_ok) response.body.clear(); + return response; +} + +template +AwsMetadataStatus map_sdk_error(const Error& error) { + const int type = static_cast(error.GetErrorType()); + const int code = static_cast(error.GetResponseCode()); + if (type == static_cast(Aws::Client::CoreErrors::ACCESS_DENIED) || code == 401 || code == 403) { + return AwsMetadataStatus::access_denied; + } + if (type == static_cast(Aws::Client::CoreErrors::THROTTLING) || + type == static_cast(Aws::Client::CoreErrors::SLOW_DOWN) || code == 429) { + return AwsMetadataStatus::throttled; + } + if (type == static_cast(Aws::Client::CoreErrors::REQUEST_TIMEOUT) || + type == static_cast(Aws::Client::CoreErrors::NETWORK_CONNECTION)) { + return AwsMetadataStatus::timeout; + } + return AwsMetadataStatus::provider_unavailable; +} + +} // namespace + +AwsImdsResponse AwsCurlImdsTransport::put_token( + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) { + struct curl_slist* headers = nullptr; + headers = curl_slist_append(headers, "X-aws-ec2-metadata-token-ttl-seconds: 21600"); + AwsImdsResponse response = imds_request( + "http://169.254.169.254/latest/api/token", "PUT", headers, + 4096, deadline, cancelled); + curl_slist_free_all(headers); + return response; +} + +AwsImdsResponse AwsCurlImdsTransport::get_identity_document( + const std::string& token, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) { + struct curl_slist* headers = nullptr; + std::string token_header = "X-aws-ec2-metadata-token: " + token; + headers = curl_slist_append(headers, token_header.c_str()); + AwsImdsResponse response = imds_request( + "http://169.254.169.254/latest/dynamic/instance-identity/document", + "GET", headers, 16384, deadline, cancelled); + curl_slist_free_all(headers); + if (!token_header.empty()) { + OPENSSL_cleanse(&token_header[0], token_header.size()); + } + return response; +} + +std::shared_ptr AwsSdkRdsDiscoveryApi::client_for_region( + const std::string& region) { + std::lock_guard lock(clients_mutex_); + const auto found = clients_.find(region); + if (found != clients_.end()) return found->second; + Aws::Client::ClientConfiguration config; + config.region = region.c_str(); + config.connectTimeoutMs = 500; + config.requestTimeoutMs = 4000; + config.httpRequestTimeoutMs = 4000; + config.retryStrategy = std::make_shared(2, 50); + auto client = std::make_shared(config); + clients_.emplace(region, client); + return client; +} + +AwsRdsInstancesPage AwsSdkRdsDiscoveryApi::describe_instances( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) { + AwsRdsInstancesPage page; + if (cancelled()) { page.status = AwsMetadataStatus::cancelled; return page; } + if (std::chrono::steady_clock::now() >= deadline) { + page.status = AwsMetadataStatus::timeout; + return page; + } + Aws::RDS::Model::DescribeDBInstancesRequest request; + if (!marker.empty()) request.SetMarker(marker.c_str()); + const auto outcome = client_for_region(region)->DescribeDBInstances(request); + if (!outcome.IsSuccess()) { + page.status = map_sdk_error(outcome.GetError()); + return page; + } + page.status = AwsMetadataStatus::ok; + page.next_marker = outcome.GetResult().GetMarker().c_str(); + for (const auto& instance : outcome.GetResult().GetDBInstances()) { + const auto& endpoint = instance.GetEndpoint(); + page.instances.push_back({endpoint.GetAddress().c_str(), endpoint.GetPort(), + instance.GetAvailabilityZone().c_str(), instance.GetDBInstanceArn().c_str()}); + } + return page; +} + +AwsRdsClustersPage AwsSdkRdsDiscoveryApi::describe_clusters( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) { + AwsRdsClustersPage page; + if (cancelled()) { page.status = AwsMetadataStatus::cancelled; return page; } + if (std::chrono::steady_clock::now() >= deadline) { + page.status = AwsMetadataStatus::timeout; + return page; + } + Aws::RDS::Model::DescribeDBClustersRequest request; + if (!marker.empty()) request.SetMarker(marker.c_str()); + const auto outcome = client_for_region(region)->DescribeDBClusters(request); + if (!outcome.IsSuccess()) { + page.status = map_sdk_error(outcome.GetError()); + return page; + } + page.status = AwsMetadataStatus::ok; + page.next_marker = outcome.GetResult().GetMarker().c_str(); + for (const auto& cluster : outcome.GetResult().GetDBClusters()) { + std::vector custom; + for (const auto& endpoint : cluster.GetCustomEndpoints()) { + custom.emplace_back(endpoint.c_str()); + } + page.clusters.push_back({cluster.GetDBClusterIdentifier().c_str(), + cluster.GetEndpoint().c_str(), cluster.GetReaderEndpoint().c_str(), + cluster.GetPort(), std::move(custom), cluster.GetDBClusterArn().c_str()}); + } + return page; +} + +AwsRdsClusterEndpointsPage AwsSdkRdsDiscoveryApi::describe_cluster_endpoints( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) { + AwsRdsClusterEndpointsPage page; + if (cancelled()) { page.status = AwsMetadataStatus::cancelled; return page; } + if (std::chrono::steady_clock::now() >= deadline) { + page.status = AwsMetadataStatus::timeout; + return page; + } + Aws::RDS::Model::DescribeDBClusterEndpointsRequest request; + if (!marker.empty()) request.SetMarker(marker.c_str()); + const auto outcome = client_for_region(region)->DescribeDBClusterEndpoints(request); + if (!outcome.IsSuccess()) { + page.status = map_sdk_error(outcome.GetError()); + return page; + } + page.status = AwsMetadataStatus::ok; + page.next_marker = outcome.GetResult().GetMarker().c_str(); + for (const auto& endpoint : outcome.GetResult().GetDBClusterEndpoints()) { + page.endpoints.push_back({endpoint.GetEndpoint().c_str(), + endpoint.GetEndpointType().c_str(), endpoint.GetDBClusterIdentifier().c_str()}); + } + return page; +} +#endif // PROXYSQL_AWS_SDK_PROVIDER diff --git a/plugins/aws/src/aws_locality_provider.h b/plugins/aws/src/aws_locality_provider.h new file mode 100644 index 0000000000..8d6643c848 --- /dev/null +++ b/plugins/aws/src/aws_locality_provider.h @@ -0,0 +1,218 @@ +#ifndef PROXYSQL_AWS_LOCALITY_PROVIDER_H +#define PROXYSQL_AWS_LOCALITY_PROVIDER_H + +#include "Aws_Locality_Types.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace Aws { namespace RDS { class RDSClient; } } + +using AwsLocalityCancelPredicate = std::function; +using AwsLocalityEnvironmentGetter = std::function; + +struct AwsMetadataProviderConfig { + using SteadyClock = std::function; + + size_t worker_count { 2 }; + size_t max_pending { 256 }; + SteadyClock steady_clock { + [] { return std::chrono::steady_clock::now(); } + }; +}; + +class AwsLocalityDiscoveryBackend { +public: + virtual AwsMetadataResult discover( + const AwsMetadataRequest& request, + const AwsLocalityCancelPredicate& cancelled) = 0; + virtual ~AwsLocalityDiscoveryBackend() = default; +}; + +class AwsSdkMetadataProvider final : public AwsMetadataProvider { +public: + explicit AwsSdkMetadataProvider( + std::shared_ptr backend, + AwsMetadataProviderConfig config = {}); + ~AwsSdkMetadataProvider() override; + + AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) override; + void cancel(AwsMetadataRequestHandle handle) override; + void shutdown() override; + + AwsSdkMetadataProvider(const AwsSdkMetadataProvider&) = delete; + AwsSdkMetadataProvider& operator=(const AwsSdkMetadataProvider&) = delete; + +private: + class Impl; + std::unique_ptr impl_; +}; + +struct AwsImdsResponse { + bool transport_ok { false }; + long status_code { 0 }; + std::string body; +}; + +class AwsImdsTransport { +public: + virtual AwsImdsResponse put_token( + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) = 0; + virtual AwsImdsResponse get_identity_document( + const std::string& token, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) = 0; + virtual ~AwsImdsTransport() = default; +}; + +AwsLocalLocation aws_locality_environment_location( + const AwsLocalityEnvironmentGetter& getenv_value); + +class AwsLocalityLocalDiscovery { +public: + AwsLocalityLocalDiscovery( + std::shared_ptr transport, + AwsLocalityEnvironmentGetter getenv_value); + + AwsMetadataResult discover( + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) const; + +private: + std::shared_ptr transport_; + AwsLocalityEnvironmentGetter getenv_value_; +}; + +struct AwsRdsInstanceRecord { + std::string endpoint; + int port { 0 }; + std::string availability_zone; + std::string arn; +}; + +struct AwsRdsClusterRecord { + std::string identifier; + std::string endpoint; + std::string reader_endpoint; + int port { 0 }; + std::vector custom_endpoints; + std::string arn; +}; + +struct AwsRdsClusterEndpointRecord { + std::string endpoint; + std::string endpoint_type; + std::string cluster_identifier; +}; + +struct AwsRdsInstancesPage { + AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; + std::vector instances; + std::string next_marker; +}; + +struct AwsRdsClustersPage { + AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; + std::vector clusters; + std::string next_marker; +}; + +struct AwsRdsClusterEndpointsPage { + AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; + std::vector endpoints; + std::string next_marker; +}; + +class AwsRdsDiscoveryApi { +public: + virtual AwsRdsInstancesPage describe_instances( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) = 0; + virtual AwsRdsClustersPage describe_clusters( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) = 0; + virtual AwsRdsClusterEndpointsPage describe_cluster_endpoints( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) = 0; + virtual ~AwsRdsDiscoveryApi() = default; +}; + +class AwsLocalityRdsDiscovery { +public: + explicit AwsLocalityRdsDiscovery(std::shared_ptr api); + + AwsMetadataResult discover( + const AwsMetadataRequest& request, + const AwsLocalityCancelPredicate& cancelled) const; + +private: + std::shared_ptr api_; +}; + +class AwsCurlImdsTransport final : public AwsImdsTransport { +public: + AwsImdsResponse put_token( + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) override; + AwsImdsResponse get_identity_document( + const std::string& token, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) override; +}; + +class AwsSdkRdsDiscoveryApi final : public AwsRdsDiscoveryApi { +public: + AwsRdsInstancesPage describe_instances( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) override; + AwsRdsClustersPage describe_clusters( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) override; + AwsRdsClusterEndpointsPage describe_cluster_endpoints( + const std::string& region, + const std::string& marker, + std::chrono::steady_clock::time_point deadline, + const AwsLocalityCancelPredicate& cancelled) override; + +private: + std::shared_ptr client_for_region(const std::string& region); + std::mutex clients_mutex_; + std::unordered_map> clients_; +}; + +class AwsLocalityCompositeDiscovery final : public AwsLocalityDiscoveryBackend { +public: + AwsLocalityCompositeDiscovery( + std::shared_ptr imds, + AwsLocalityEnvironmentGetter getenv_value, + std::shared_ptr rds); + + AwsMetadataResult discover( + const AwsMetadataRequest& request, + const AwsLocalityCancelPredicate& cancelled) override; + +private: + AwsLocalityLocalDiscovery local_; + AwsLocalityRdsDiscovery rds_; +}; + +#endif // PROXYSQL_AWS_LOCALITY_PROVIDER_H diff --git a/src/main.cpp b/src/main.cpp index 64d2b6e1e7..e80bad85fd 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -43,6 +43,7 @@ using json = nlohmann::json; #include "Web_Interface.hpp" #ifdef PROXYSQL40 #include "ProxySQL_PluginManager.h" +#include "Aws_Locality_Manager.h" #endif /* PROXYSQL40 */ #include "proxysql_utils.h" #include "PgSQL_Monitor.hpp" @@ -1808,6 +1809,14 @@ bool ProxySQL_Main_init_phase3___start_all() { void ProxySQL_Main_init_phase4___shutdown() { cpu_timer t; ProxySQL_Main_join_all_threads(); +#ifdef PROXYSQL40 + // The locality manager can retain a provider lease between refreshes. + // Drain it before shutting down the plugin-owned provider registry. + if (MyHGM != nullptr && MyHGM->aws_locality_manager() != nullptr) { + MyHGM->aws_locality_manager()->shutdown(); + } + shutdown_global_aws_metadata_provider(); +#endif shutdown_global_aws_iam_token_source(); //write(GloAdmin->pipefd[1], &GloAdmin->pipefd[1], 1); // write a random byte diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index 1c5445f572..c2db0ed79c 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -947,6 +947,13 @@ aws_locality_manager_unit-t: aws_locality_manager_unit-t.cpp \ $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ $(IDIRS) $(OPT) -lpthread -ldl -o $@ +aws_locality_plugin_unit-t: aws_locality_plugin_unit-t.cpp \ + $(PROXYSQL_PATH)/plugins/aws/src/aws_locality_provider.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(PROXYSQL_PATH)/plugins/aws/src/aws_locality_provider.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ + $(IDIRS) -I$(PROXYSQL_PATH)/plugins/aws/src $(OPT) -lpthread -lcrypto -o $@ + test_aws_iam_metrics-t: ../test_aws_iam_metrics-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) $(CXX) $< $(TEST_HELPERS_OBJ) -I$(TEST_HELPERS_DIR) \ $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ diff --git a/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp b/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp new file mode 100644 index 0000000000..3e9d7d2103 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp @@ -0,0 +1,411 @@ +#include "tap.h" + +#include "aws_locality_provider.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +namespace { + +class CapturingSink final : public AwsMetadataCompletionSink { +public: + void post(AwsMetadataCompletion&& completion) override { + std::lock_guard lock(mutex_); + completions_.push_back(std::move(completion)); + cv_.notify_all(); + } + + bool wait_for(size_t count) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, 2s, [&] { return completions_.size() >= count; }); + } + + std::vector snapshot() const { + std::lock_guard lock(mutex_); + return completions_; + } + +private: + mutable std::mutex mutex_; + std::condition_variable cv_; + std::vector completions_; +}; + +class BlockingBackend final : public AwsLocalityDiscoveryBackend { +public: + AwsMetadataResult discover( + const AwsMetadataRequest& request, + const AwsLocalityCancelPredicate& cancelled) override { + const int active = active_.fetch_add(1) + 1; + int observed = max_active_.load(); + while (active > observed && + !max_active_.compare_exchange_weak(observed, active)) {} + { + std::unique_lock lock(mutex_); + started_.push_back(request.opaque_id); + cv_.notify_all(); + cv_.wait(lock, [&] { return released_ || cancelled(); }); + } + active_.fetch_sub(1); + AwsMetadataResult result; + result.status = cancelled() ? AwsMetadataStatus::cancelled : AwsMetadataStatus::ok; + result.failure_category = "FAKE_SECRET_RAW_ERROR"; + result.local.region = request.region; + return result; + } + + bool wait_started(size_t count) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, 2s, [&] { return started_.size() >= count; }); + } + + void release() { + std::lock_guard lock(mutex_); + released_ = true; + cv_.notify_all(); + } + + int max_active() const { return max_active_.load(); } + +private: + std::atomic active_ { 0 }; + std::atomic max_active_ { 0 }; + std::mutex mutex_; + std::condition_variable cv_; + std::vector started_; + bool released_ { false }; +}; + +class FakeImdsTransport final : public AwsImdsTransport { +public: + AwsImdsResponse put_token( + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { + ++token_calls; + return token; + } + + AwsImdsResponse get_identity_document( + const std::string& supplied_token, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { + ++document_calls; + seen_token = supplied_token; + return document; + } + + AwsImdsResponse token { true, 200, "imds-token" }; + AwsImdsResponse document { true, 200, + R"({"region":"us-east-1","availabilityZone":"us-east-1b","accountId":"111122223333"})" }; + int token_calls { 0 }; + int document_calls { 0 }; + std::string seen_token; +}; + +class FakeRdsApi final : public AwsRdsDiscoveryApi { +public: + AwsRdsInstancesPage describe_instances( + const std::string&, const std::string& marker, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { + ++instance_calls; + if (marker.empty()) { + AwsRdsInstancesPage page; + page.status = AwsMetadataStatus::ok; + page.next_marker = "instances-2"; + page.instances.push_back({ + "DB-ONE.ABCDEFGHIJKL.US-EAST-1.RDS.AMAZONAWS.COM.", 3306, + "us-east-1a", "arn:aws:rds:us-east-1:111122223333:db:one"}); + page.instances.push_back({ + "missing-port.abcdefghijkl.us-east-1.rds.amazonaws.com", 0, + "us-east-1a", "arn:aws:rds:us-east-1:111122223333:db:missing"}); + return page; + } + AwsRdsInstancesPage page; + page.status = AwsMetadataStatus::ok; + page.instances.push_back({ + "db-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306, + "us-east-1d", "arn:aws:rds:us-east-1:111122223333:db:one"}); + page.instances.push_back({ + "db-two.abcdefghijkl.us-east-1.rds.amazonaws.com", 3307, + "us-east-1c", "arn:aws:rds:us-east-1:444455556666:db:two"}); + return page; + } + + AwsRdsClustersPage describe_clusters( + const std::string&, const std::string&, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { + ++cluster_calls; + AwsRdsClustersPage page; + page.status = AwsMetadataStatus::ok; + page.clusters.push_back({ + "cluster-one", "cluster-one.abcdefghijkl.us-east-1.rds.amazonaws.com", + "cluster-ro-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306, + {"custom-one.abcdefghijkl.us-east-1.rds.amazonaws.com"}, + "arn:aws:rds:us-east-1:111122223333:cluster:cluster-one"}); + return page; + } + + AwsRdsClusterEndpointsPage describe_cluster_endpoints( + const std::string&, const std::string&, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { + ++endpoint_calls; + AwsRdsClusterEndpointsPage page; + page.status = AwsMetadataStatus::ok; + page.endpoints.push_back({ + "custom-two.abcdefghijkl.us-east-1.rds.amazonaws.com", + "CUSTOM", "cluster-one"}); + return page; + } + + int instance_calls { 0 }; + int cluster_calls { 0 }; + int endpoint_calls { 0 }; +}; + +const AwsMetadataEndpoint* find_endpoint( + const AwsMetadataResult& result, + const std::string& hostname, + uint16_t port) { + for (const auto& endpoint : result.endpoints) { + if (endpoint.hostname == hostname && endpoint.port == port) return &endpoint; + } + return nullptr; +} + +} // namespace + +int main() { + plan(31); + + auto backend = std::make_shared(); + AwsSdkMetadataProvider provider(backend, AwsMetadataProviderConfig {2, 3}); + auto sink = std::make_shared(); + std::vector handles; + for (uint64_t id = 1; id <= 4; ++id) { + AwsMetadataRequest request; + request.kind = AwsMetadataRequestKind::rds_region; + request.opaque_id = id; + request.generation = 17; + request.region = "us-east-1"; + request.deadline = std::chrono::steady_clock::now() + 2s; + handles.push_back(provider.request(request, sink)); + } + ok(backend->wait_started(2), "provider starts its two bounded workers"); + ok(backend->max_active() == 2, "provider never exceeds two concurrent discoveries"); + ok(handles[0].value != 0 && handles[1].value != 0 && handles[2].value != 0, + "accepted work receives cancellable handles"); + ok(handles[3].value == 0 && sink->wait_for(1), + "bounded queue rejects excess work immediately"); + auto completions = sink->snapshot(); + ok(completions[0].opaque_id == 4 && completions[0].generation == 17 && + completions[0].result.status == AwsMetadataStatus::throttled, + "queue rejection preserves request identity with a fixed category"); + ok(completions[0].result.failure_category == "throttled" && + completions[0].result.failure_category.find("FAKE_SECRET") == std::string::npos, + "provider never forwards a backend's raw failure text"); + provider.cancel(handles[2]); + backend->release(); + ok(sink->wait_for(3), "accepted non-cancelled work completes"); + completions = sink->snapshot(); + bool saw_one = false; + bool saw_two = false; + bool saw_three = false; + for (const auto& completion : completions) { + saw_one = saw_one || (completion.opaque_id == 1 && completion.generation == 17); + saw_two = saw_two || (completion.opaque_id == 2 && completion.generation == 17); + saw_three = saw_three || completion.opaque_id == 3; + } + ok(saw_one && saw_two && !saw_three, + "queued cancellation suppresses its callback without affecting other work"); + ok(saw_one && saw_two, + "successful asynchronous completions preserve opaque ID and generation"); + + AwsMetadataRequest expired; + expired.opaque_id = 9; + expired.generation = 18; + expired.deadline = std::chrono::steady_clock::now() - 1ms; + ok(provider.request(expired, sink).value == 0 && sink->wait_for(4), + "already-expired work is rejected without entering the backend"); + completions = sink->snapshot(); + ok(completions.back().opaque_id == 9 && + completions.back().result.status == AwsMetadataStatus::timeout && + completions.back().result.failure_category == "timeout", + "deadline rejection is normalized and preserves identity"); + provider.shutdown(); + ok(provider.request(expired, sink).value == 0, + "shutdown permanently rejects new metadata work"); + + const std::unordered_map env { + {"AWS_REGION", "us-west-2"}, + {"AWS_DEFAULT_REGION", "eu-west-1"}, + {"AWS_AVAILABILITY_ZONE", "us-west-2b"}, + {"AWS_ACCOUNT_ID", "999900001111"}, + }; + const auto env_getter = [&](const char* name) { + const auto found = env.find(name); + return found == env.end() ? std::string() : found->second; + }; + AwsLocalLocation environment = aws_locality_environment_location(env_getter); + ok(environment.region == "us-west-2", + "AWS_REGION takes precedence over AWS_DEFAULT_REGION"); + ok(environment.availability_zone == "us-west-2b" && + environment.account_id == "999900001111", + "environment fallback retains optional AZ and account assertion"); + const auto partial_getter = [](const char* name) { + return std::string(name) == "AWS_DEFAULT_REGION" ? "eu-central-1" : ""; + }; + environment = aws_locality_environment_location(partial_getter); + ok(environment.region == "eu-central-1" && + environment.availability_zone.empty() && environment.account_id.empty(), + "partial environment fallback keeps Region while leaving AZ/account unknown"); + + auto imds = std::make_shared(); + AwsLocalityLocalDiscovery local_discovery(imds, env_getter); + const auto never_cancelled = [] { return false; }; + AwsMetadataResult local_result = local_discovery.discover( + std::chrono::steady_clock::now() + 1s, never_cancelled); + ok(local_result.status == AwsMetadataStatus::ok && + local_result.local.region == "us-east-1" && + local_result.local.availability_zone == "us-east-1b" && + local_result.local.account_id == "111122223333", + "IMDSv2 identity document supplies Region, AZ, and account"); + ok(imds->token_calls == 1 && imds->document_calls == 1 && + imds->seen_token == "imds-token", + "IMDSv2 token is required for the identity-document request"); + imds->token = {false, 0, "FAKE_SECRET_TRANSPORT_ERROR"}; + local_result = local_discovery.discover( + std::chrono::steady_clock::now() + 1s, never_cancelled); + ok(local_result.status == AwsMetadataStatus::ok && + local_result.local.region == "us-west-2", + "IMDS unavailability falls back to process environment"); + ok(local_result.failure_category.empty(), + "successful environment fallback exposes no IMDS transport detail"); + imds->token = {true, 200, "imds-token"}; + imds->document = {true, 200, "{not-json-FAKE_SECRET}"}; + AwsLocalityLocalDiscovery invalid_local(imds, + [](const char*) { return std::string(); }); + local_result = invalid_local.discover( + std::chrono::steady_clock::now() + 1s, never_cancelled); + ok(local_result.status == AwsMetadataStatus::invalid_response && + local_result.failure_category == "invalid_response", + "malformed IMDS data without fallback returns only a fixed category"); + + auto rds_api = std::make_shared(); + AwsLocalityRdsDiscovery rds_discovery(rds_api); + AwsMetadataRequest rds_request; + rds_request.kind = AwsMetadataRequestKind::rds_region; + rds_request.region = "us-east-1"; + rds_request.deadline = std::chrono::steady_clock::now() + 2s; + AwsMetadataResult rds_result = rds_discovery.discover(rds_request, never_cancelled); + ok(rds_result.status == AwsMetadataStatus::ok, + "paginated RDS discovery completes successfully"); + ok(rds_api->instance_calls == 2 && rds_api->cluster_calls == 1 && + rds_api->endpoint_calls == 1, + "RDS instances, clusters, and cluster endpoints are all paginated"); + const auto* instance_one = find_endpoint(rds_result, + "db-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); + ok(instance_one != nullptr && instance_one->endpoint_type == AwsEndpointType::instance && + instance_one->availability_zone == "us-east-1d" && + instance_one->account_id == "111122223333", + "later duplicate instance metadata replaces the earlier page after failover"); + size_t instance_one_count = 0; + for (const auto& endpoint : rds_result.endpoints) { + if (endpoint.hostname == "db-one.abcdefghijkl.us-east-1.rds.amazonaws.com" && + endpoint.port == 3306) ++instance_one_count; + } + ok(instance_one_count == 1, + "duplicate paginated endpoint metadata is emitted exactly once"); + const auto* instance_two = find_endpoint(rds_result, + "db-two.abcdefghijkl.us-east-1.rds.amazonaws.com", 3307); + ok(instance_two != nullptr && instance_two->account_id == "444455556666", + "later instance pages and cross-account identities are retained"); + ok(find_endpoint(rds_result, + "missing-port.abcdefghijkl.us-east-1.rds.amazonaws.com", 0) == nullptr, + "instance metadata without a valid port is rejected"); + const auto* cluster = find_endpoint(rds_result, + "cluster-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); + const auto* reader = find_endpoint(rds_result, + "cluster-ro-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); + ok(cluster != nullptr && cluster->endpoint_type == AwsEndpointType::cluster && + reader != nullptr && reader->endpoint_type == AwsEndpointType::reader && + cluster->availability_zone.empty() && reader->availability_zone.empty(), + "cluster writer and reader endpoints have Region/account but no stable AZ"); + const auto* custom_one = find_endpoint(rds_result, + "custom-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 0); + const auto* custom_two = find_endpoint(rds_result, + "custom-two.abcdefghijkl.us-east-1.rds.amazonaws.com", 0); + ok(custom_one != nullptr && custom_two != nullptr && + custom_one->endpoint_type == AwsEndpointType::custom && + custom_two->endpoint_type == AwsEndpointType::custom && + custom_two->account_id == "111122223333", + "custom endpoints from both cluster APIs inherit cluster account and no port"); + + class RepeatingApi final : public AwsRdsDiscoveryApi { + public: + AwsRdsInstancesPage describe_instances(const std::string&, const std::string&, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { + AwsRdsInstancesPage page; + page.status = AwsMetadataStatus::ok; + page.next_marker = "same"; + return page; + } + AwsRdsClustersPage describe_clusters(const std::string&, const std::string&, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { return {}; } + AwsRdsClusterEndpointsPage describe_cluster_endpoints( + const std::string&, const std::string&, + std::chrono::steady_clock::time_point, + const AwsLocalityCancelPredicate&) override { return {}; } + }; + AwsLocalityRdsDiscovery repeating(std::make_shared()); + rds_result = repeating.discover(rds_request, never_cancelled); + ok(rds_result.status == AwsMetadataStatus::invalid_response && + rds_result.failure_category == "invalid_response", + "repeated pagination markers fail with a fixed category"); + + class ShutdownBackend final : public AwsLocalityDiscoveryBackend { + public: + AwsMetadataResult discover(const AwsMetadataRequest&, + const AwsLocalityCancelPredicate& cancelled) override { + started.store(true); + while (!cancelled()) std::this_thread::yield(); + AwsMetadataResult result; + result.status = AwsMetadataStatus::cancelled; + return result; + } + std::atomic started { false }; + }; + auto shutdown_backend = std::make_shared(); + AwsSdkMetadataProvider shutdown_provider( + shutdown_backend, AwsMetadataProviderConfig {2, 8}); + auto shutdown_sink = std::make_shared(); + AwsMetadataRequest shutdown_request; + shutdown_request.opaque_id = 44; + shutdown_request.deadline = std::chrono::steady_clock::now() + 2s; + shutdown_provider.request(shutdown_request, shutdown_sink); + while (!shutdown_backend->started.load()) std::this_thread::yield(); + std::thread shutdown_one([&] { shutdown_provider.shutdown(); }); + std::thread shutdown_two([&] { shutdown_provider.shutdown(); }); + shutdown_one.join(); + shutdown_two.join(); + ok(shutdown_sink->snapshot().empty(), + "concurrent shutdown is idempotent, cancels work, and publishes no late callback"); + ok(shutdown_provider.request(shutdown_request, shutdown_sink).value == 0, + "post-shutdown requests remain rejected"); + + return exit_status(); +} From 39dab690bdef64ae26c85ad3bee266c15b7095d0 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 21:08:29 +0000 Subject: [PATCH 09/17] feat(stats): expose AWS locality decisions --- include/MySQL_HostGroups_Manager.h | 4 + include/ProxySQL_Plugin.h | 11 +- lib/MySQL_HostGroups_Manager.cpp | 94 +++++++ lib/ProxySQL_Admin.cpp | 14 +- lib/ProxySQL_PluginManager.cpp | 13 + test/tap/groups/groups.json | 1 + test/tap/tests/unit/Makefile | 7 +- .../tests/unit/aws_locality_stats_unit-t.cpp | 240 ++++++++++++++++++ 8 files changed, 374 insertions(+), 10 deletions(-) create mode 100644 test/tap/tests/unit/aws_locality_stats_unit-t.cpp diff --git a/include/MySQL_HostGroups_Manager.h b/include/MySQL_HostGroups_Manager.h index 79b44e0f23..de85449b96 100644 --- a/include/MySQL_HostGroups_Manager.h +++ b/include/MySQL_HostGroups_Manager.h @@ -894,6 +894,10 @@ class MySQL_HostGroups_Manager : public Base_HostGroups_Manager { #ifdef PROXYSQL40 void refresh_aws_locality_configuration(); void set_aws_locality_awareness_enabled(bool enabled); + void refresh_aws_locality_stats(SQLite3DB* statsdb) const; + static bool project_aws_locality_stats( + SQLite3DB* statsdb, + const std::vector& rows); MySQLAwsLocalityManager* aws_locality_manager() const { return aws_locality_manager_.get(); } diff --git a/include/ProxySQL_Plugin.h b/include/ProxySQL_Plugin.h index 527bf03b10..f5040102da 100644 --- a/include/ProxySQL_Plugin.h +++ b/include/ProxySQL_Plugin.h @@ -43,8 +43,10 @@ namespace prometheus { class Registry; } // sizing callbacks. They are live only during normal plugin init. // ABI 6: ProxySQL_PluginServices gains the general AWS metadata-provider // installation callback used by locality discovery. -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 6u; -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 6u; +// ABI 7: ProxySQL_PluginServices gains the MySQL-owned AWS-locality stats +// projection callback used by the AWS plugin's runtime view. +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 7u; +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 7u; enum class ProxySQL_PluginDBKind : uint8_t { admin_db = 0, @@ -248,6 +250,9 @@ using proxysql_plugin_get_aws_iam_limits_cb = using proxysql_plugin_install_aws_metadata_provider_cb = bool (*)(AwsMetadataProvider *, void (*)(AwsMetadataProvider *), void *module_handle); + +using proxysql_plugin_refresh_mysql_aws_locality_stats_cb = + void (*)(SQLite3DB *); #endif /* PROXYSQL40 */ // Services provided to plugins across the four-phase lifecycle. @@ -316,6 +321,8 @@ struct ProxySQL_PluginServices { proxysql_plugin_get_aws_iam_limits_cb get_aws_iam_limits; // ABI-6 extension. Null outside normal plugin init(). proxysql_plugin_install_aws_metadata_provider_cb install_aws_metadata_provider; + // ABI-7 extension. Live in Phase B and normal init; it performs no I/O. + proxysql_plugin_refresh_mysql_aws_locality_stats_cb refresh_mysql_aws_locality_stats; #endif /* PROXYSQL40 */ }; diff --git a/lib/MySQL_HostGroups_Manager.cpp b/lib/MySQL_HostGroups_Manager.cpp index 96907e0435..8591fe12ab 100644 --- a/lib/MySQL_HostGroups_Manager.cpp +++ b/lib/MySQL_HostGroups_Manager.cpp @@ -878,6 +878,100 @@ void MySQL_HostGroups_Manager::set_aws_locality_awareness_enabled(bool enabled) aws_locality_manager_->set_enabled(enabled); } } + +namespace { + +const char* aws_endpoint_type_name(AwsEndpointType type) { + switch (type) { + case AwsEndpointType::instance: return "instance"; + case AwsEndpointType::cluster: return "cluster"; + case AwsEndpointType::reader: return "reader"; + case AwsEndpointType::custom: return "custom"; + case AwsEndpointType::unknown: return "unknown"; + } + return "unknown"; +} + +const char* aws_locality_name(AwsLocalityClass locality) { + switch (locality) { + case AwsLocalityClass::remote: return "remote"; + case AwsLocalityClass::same_region: return "same_region"; + case AwsLocalityClass::same_az: return "same_az"; + case AwsLocalityClass::unknown: return "unknown"; + } + return "unknown"; +} + +const char* aws_metadata_status_name(AwsLocalityMetadataStatus status) { + switch (status) { + case AwsLocalityMetadataStatus::disabled: return "disabled"; + case AwsLocalityMetadataStatus::pending: return "pending"; + case AwsLocalityMetadataStatus::fresh: return "fresh"; + case AwsLocalityMetadataStatus::stale: return "stale"; + case AwsLocalityMetadataStatus::expired: return "expired"; + case AwsLocalityMetadataStatus::error: return "error"; + } + return "error"; +} + +const char* aws_account_match_name(const AwsLocalitySnapshotEntry& row) { + if (row.local.account_id.empty() || row.backend.account_id.empty()) { + return "unknown"; + } + return row.local.account_id == row.backend.account_id ? "same" : "different"; +} + +bool aws_locality_status_is_active(AwsLocalityMetadataStatus status) { + return status == AwsLocalityMetadataStatus::fresh || + status == AwsLocalityMetadataStatus::stale; +} + +} // namespace + +bool MySQL_HostGroups_Manager::project_aws_locality_stats( + SQLite3DB* statsdb, + const std::vector& rows) { + if (statsdb == nullptr || !statsdb->execute("BEGIN")) return false; + bool success = statsdb->execute("DELETE FROM stats_mysql_aws_locality"); + for (const auto& row : rows) { + if (!success) break; + const double active_multiplier = aws_locality_status_is_active(row.status) + ? row.multiplier : 1.0; + const uint64_t effective_weight = aws_locality_effective_weight( + row.configured_weight, active_multiplier); + char* query = sqlite3_mprintf( + "INSERT INTO stats_mysql_aws_locality (" + "hostgroup_id,hostname,port,endpoint_type,configured_weight," + "effective_weight,local_region,local_az,backend_region,backend_az," + "account_match,locality,active_multiplier,metadata_status," + "last_success_timestamp,last_attempt_timestamp,last_error_category) " + "VALUES (%u,'%q',%u,'%q',%lld,%llu,'%q','%q','%q','%q','%q','%q'," + "%.17g,'%q',%lld,%lld,'%q')", + row.hostgroup_id, row.hostname.c_str(), static_cast(row.port), + aws_endpoint_type_name(row.endpoint_type), + static_cast(row.configured_weight), + static_cast(effective_weight), + row.local.region.c_str(), row.local.availability_zone.c_str(), + row.backend.region.c_str(), row.backend.availability_zone.c_str(), + aws_account_match_name(row), aws_locality_name(row.locality), + active_multiplier, aws_metadata_status_name(row.status), + static_cast(row.last_success_timestamp), + static_cast(row.last_attempt_timestamp), + row.failure_category.c_str()); + success = query != nullptr && statsdb->execute(query); + sqlite3_free(query); + } + if (success) success = statsdb->execute("COMMIT"); + if (!success) statsdb->execute("ROLLBACK"); + return success; +} + +void MySQL_HostGroups_Manager::refresh_aws_locality_stats(SQLite3DB* statsdb) const { + const std::vector rows = aws_locality_manager_ + ? aws_locality_manager_->diagnostic_rows() + : std::vector(); + project_aws_locality_stats(statsdb, rows); +} #endif void MySQL_HostGroups_Manager::p_update_mysql_error_counter(p_mysql_error_type err_type, unsigned int hid, char* address, uint16_t port, unsigned int code) { diff --git a/lib/ProxySQL_Admin.cpp b/lib/ProxySQL_Admin.cpp index b7a9ce5562..e851a3253d 100644 --- a/lib/ProxySQL_Admin.cpp +++ b/lib/ProxySQL_Admin.cpp @@ -1613,10 +1613,11 @@ bool ProxySQL_Admin::GenericRefreshStatistics(const char *query_no_space, unsign } #ifdef PROXYSQL40 // Plugin-registered runtime views: if the query references any chassis- - // registered runtime view (e.g. runtime_mysqlx_users), refresh it on - // the admin path BEFORE the SELECT runs against admindb. We always - // invoke the dispatcher when the session is on the admin port; the - // chassis itself decides whether to fire any plugin's refresh +// registered runtime view (e.g. runtime_mysqlx_users or an on-demand stats +// table), refresh it BEFORE the SELECT runs. Admin sessions can project all +// three DB kinds; stats sessions receive only the stats handle, so they cannot +// trigger an admin/config projection. The chassis decides whether to fire any +// plugin's refresh // callback by per-view substring match against query_no_space, so a // query that touches no registered view is a cheap no-op (one shared // lock + N substring scans, N == registered-view count). @@ -1628,9 +1629,8 @@ bool ProxySQL_Admin::GenericRefreshStatistics(const char *query_no_space, unsign // that touches only a plugin view (e.g. SELECT * FROM runtime_mysqlx_ // users with no other runtime_* mention) still gets its projection // fired. - if (admin) { - proxysql_refresh_configured_plugin_runtime_views(query_no_space, admindb, configdb, statsdb); - } + proxysql_refresh_configured_plugin_runtime_views(query_no_space, + admin ? admindb : nullptr, admin ? configdb : nullptr, statsdb); #endif /* PROXYSQL40 */ // if (stats_mysql_processlist || stats_mysql_connection_pool || stats_mysql_query_digest || stats_mysql_query_digest_reset) { if (refresh==true) { diff --git a/lib/ProxySQL_PluginManager.cpp b/lib/ProxySQL_PluginManager.cpp index 5aa0a20f00..7b5e96c130 100644 --- a/lib/ProxySQL_PluginManager.cpp +++ b/lib/ProxySQL_PluginManager.cpp @@ -7,6 +7,7 @@ #include "ProxySQL_PluginManager.h" #include "Aws_Iam_Provider.h" #include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" #include "MySQL_Thread.h" #include @@ -208,6 +209,15 @@ bool install_aws_metadata_provider_service( } return install_global_aws_metadata_provider(provider, destroy, module_handle); } + +void refresh_mysql_aws_locality_stats_service(SQLite3DB* statsdb) { + if (statsdb == nullptr) return; + if (MyHGM != nullptr) { + MyHGM->refresh_aws_locality_stats(statsdb); + return; + } + MySQL_HostGroups_Manager::project_aws_locality_stats(statsdb, {}); +} #endif /* PROXYSQL40 */ SQLite3DB* get_admindb_service() { @@ -337,6 +347,7 @@ ProxySQL_PluginManager::ProxySQL_PluginManager() { services_.install_aws_iam_token_source = &install_aws_iam_token_source_service; services_.get_aws_iam_limits = &get_aws_iam_limits_service; services_.install_aws_metadata_provider = &install_aws_metadata_provider_service; + services_.refresh_mysql_aws_locality_stats = &refresh_mysql_aws_locality_stats_service; // Phase-B (register_schemas) services: same layout as init(), but DB // handle getters and the query-hook registrar are stubbed -- see the @@ -362,6 +373,8 @@ ProxySQL_PluginManager::ProxySQL_PluginManager() { // refresh callback won't fire until Admin handles a SELECT, by which // point admin module bootstrap has long since completed. services_phase_b_.register_runtime_view = ®ister_runtime_view_service; + services_phase_b_.refresh_mysql_aws_locality_stats = + &refresh_mysql_aws_locality_stats_service; #endif /* PROXYSQL40 */ } diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index e9b55bf4e2..f365df1671 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -25,6 +25,7 @@ "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_config_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_selection_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_stats_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], "basic-t" : [ "legacy-g1","mariadb10-galera-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql84-gr-g1","mysql90-g1","mysql90-gr-g1","mysql93-g1","mysql93-gr-g1","mysql95-g1","mysql95-gr-g1" ], "c_tokenizer_unit-t" : [ "unit-tests-g1" ], diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index c2db0ed79c..3a41f62a4a 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -419,7 +419,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t aws_locality_config_unit-t aws_locality_selection_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t aws_locality_config_unit-t aws_locality_selection_unit-t aws_locality_stats_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ @@ -965,6 +965,11 @@ test_aws_iam_backend_auth-t: ../test_aws_iam_backend_auth-t.cpp $(TEST_HELPERS_O $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ +aws_locality_stats_unit-t: aws_locality_stats_unit-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) aws_plugin_build + $(CXX) -DPROXYSQL_AWS_PLUGIN_PATH=\"$(PROXYSQL_PATH)/plugins/aws/ProxySQL_Aws_Plugin.so\" \ + $< $(TEST_HELPERS_OBJ) $(IDIRS) $(LDIRS) $(OPT) \ + $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ + ifeq ($(UNAME_S),Linux) aws_iam_session_state_unit-t: ALLOW_MULTI_DEF += -Wl,--wrap=mysql_real_connect_start -pthread aws_iam_pool_unit-t: ALLOW_MULTI_DEF += \ diff --git a/test/tap/tests/unit/aws_locality_stats_unit-t.cpp b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp new file mode 100644 index 0000000000..84cace6e53 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp @@ -0,0 +1,240 @@ +#include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" +#include "ProxySQL_PluginManager.h" +#include "sqlite3db.h" +#include "tap.h" +#include "test_globals.h" + +#include +#include +#include +#include +#include + +#ifndef PROXYSQL_AWS_PLUGIN_PATH +#error "PROXYSQL_AWS_PLUGIN_PATH must be defined" +#endif + +extern MySQL_HostGroups_Manager* MyHGM; + +namespace { + +class CountingProvider final : public AwsMetadataProvider { +public: + AwsMetadataRequestHandle request( + const AwsMetadataRequest&, + std::weak_ptr) override { + ++requests; + return {requests.load()}; + } + void cancel(AwsMetadataRequestHandle) override {} + void shutdown() override {} + std::atomic requests {0}; +}; + +CountingProvider* counting_provider = nullptr; + +void destroy_counting_provider(AwsMetadataProvider* provider) { + delete static_cast(provider); + counting_provider = nullptr; +} + +AwsLocalityHostgroupConfig disabled_hostgroup( + uint32_t hostgroup_id, const std::string& hostname, int64_t weight) { + AwsLocalityHostgroupConfig config; + config.hostgroup_id = hostgroup_id; + config.policy.valid = true; + AwsEndpointCandidate endpoint; + endpoint.recognized = true; + endpoint.hostgroup_id = hostgroup_id; + endpoint.hostname = hostname; + endpoint.port = 3306; + endpoint.region = "us-east-1"; + endpoint.partition = "aws"; + config.backends.emplace_back(std::move(endpoint), weight); + return config; +} + +AwsLocalitySnapshotEntry diagnostic_row( + uint32_t hostgroup_id, + AwsLocalityMetadataStatus status, + double multiplier, + int64_t weight) { + AwsLocalitySnapshotEntry row; + row.hostgroup_id = hostgroup_id; + row.hostname = hostgroup_id == 6 + ? "db'quoted.abcdefghijkl.us-east-1.rds.amazonaws.com" + : "db-" + std::to_string(hostgroup_id) + + ".abcdefghijkl.us-east-1.rds.amazonaws.com"; + row.port = 3306; + row.endpoint_type = hostgroup_id == 1 ? AwsEndpointType::unknown + : hostgroup_id == 2 ? AwsEndpointType::instance + : hostgroup_id == 3 ? AwsEndpointType::cluster + : hostgroup_id == 4 ? AwsEndpointType::reader + : AwsEndpointType::custom; + row.configured_weight = weight; + row.local = {"us-east-1", "us-east-1a", "111122223333"}; + row.backend.region = hostgroup_id == 4 ? "eu-west-1" : "us-east-1"; + row.backend.availability_zone = hostgroup_id == 2 ? "us-east-1a" : ""; + row.backend.account_id = hostgroup_id == 1 ? "" + : hostgroup_id == 3 ? "444455556666" : "111122223333"; + row.locality = hostgroup_id == 2 ? AwsLocalityClass::same_az + : hostgroup_id == 3 ? AwsLocalityClass::same_region + : hostgroup_id == 4 ? AwsLocalityClass::remote + : AwsLocalityClass::unknown; + row.multiplier = multiplier; + row.status = status; + row.last_success_timestamp = 1700000000 + hostgroup_id; + row.last_attempt_timestamp = 1700000100 + hostgroup_id; + row.failure_category = status == AwsLocalityMetadataStatus::stale + ? "throttled" : status == AwsLocalityMetadataStatus::error + ? "access_denied" : ""; + return row; +} + +} // namespace + +int main() { + plan(22); + + SQLite3DB statsdb; + statsdb.open((char*)":memory:", + SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX); + ok(statsdb.return_one_int( + "SELECT count(*) FROM sqlite_master WHERE name='stats_mysql_aws_locality'") == 0, + "AWS locality stats table is absent without the AWS plugin"); + + std::unique_ptr manager; + std::string error; + ok(proxysql_load_configured_plugins(manager, {PROXYSQL_AWS_PLUGIN_PATH}, error), + "real AWS plugin completes schema-registration phase"); + if (!error.empty()) diag("plugin error: %s", error.c_str()); + + const auto& tables = manager->tables(ProxySQL_PluginDBKind::stats_db); + ok(tables.size() == 1 && + std::string(tables[0].table_name) == "stats_mysql_aws_locality", + "AWS plugin registers exactly its locality table in stats DB"); + ok(manager->tables(ProxySQL_PluginDBKind::admin_db).empty() && + manager->tables(ProxySQL_PluginDBKind::config_db).empty(), + "locality diagnostics add no admin/config persistence surface"); + + if (!tables.empty()) statsdb.execute(tables[0].table_def); + ok(statsdb.return_one_int( + "SELECT count(*) FROM pragma_table_info('stats_mysql_aws_locality')") == 17, + "plugin-owned locality table has the exact 17-column schema"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM pragma_table_info('stats_mysql_aws_locality') " + "WHERE name IN ('hostgroup_id','hostname','port','endpoint_type'," + "'configured_weight','effective_weight','local_region','local_az'," + "'backend_region','backend_az','account_match','locality'," + "'active_multiplier','metadata_status','last_success_timestamp'," + "'last_attempt_timestamp','last_error_category')") == 17, + "locality stats schema exposes the documented column names"); + + std::vector rows; + rows.push_back(diagnostic_row(1, AwsLocalityMetadataStatus::pending, 4.0, 10)); + rows.push_back(diagnostic_row(2, AwsLocalityMetadataStatus::fresh, 2.5, 11)); + rows.push_back(diagnostic_row(3, AwsLocalityMetadataStatus::stale, 4.0, 12)); + rows.push_back(diagnostic_row(4, AwsLocalityMetadataStatus::expired, 5.0, 13)); + rows.push_back(diagnostic_row(5, AwsLocalityMetadataStatus::error, 6.0, 14)); + rows.push_back(diagnostic_row(6, AwsLocalityMetadataStatus::disabled, 7.0, 15)); + ok(MySQL_HostGroups_Manager::project_aws_locality_stats(&statsdb, rows), + "one retained diagnostics snapshot projects transactionally"); + ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 6, + "projection emits one row per configured backend"); + ok(statsdb.return_one_int( + "SELECT count(DISTINCT metadata_status) FROM stats_mysql_aws_locality") == 6, + "pending, fresh, stale, expired, error, and disabled are explicit"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE " + "(hostgroup_id=2 AND effective_weight=27 AND active_multiplier=2.5) OR " + "(hostgroup_id=3 AND effective_weight=48 AND active_multiplier=4.0)") == 2, + "fresh/stale rows expose integer-cast weighted multipliers"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id IN (1,4,5,6) " + "AND effective_weight=configured_weight AND active_multiplier=1.0") == 4, + "pending/expired/error/disabled rows force neutral effective weights"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE " + "(hostgroup_id=2 AND endpoint_type='instance' AND locality='same_az') OR " + "(hostgroup_id=3 AND endpoint_type='cluster' AND locality='same_region') OR " + "(hostgroup_id=4 AND endpoint_type='reader' AND locality='remote')") == 3, + "endpoint and locality classifications use stable strings"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE " + "(hostgroup_id=1 AND account_match='unknown') OR " + "(hostgroup_id=2 AND account_match='same') OR " + "(hostgroup_id=3 AND account_match='different')") == 3, + "account comparison exposes unknown/same/different without identifiers"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=3 " + "AND last_success_timestamp=1700000003 AND last_attempt_timestamp=1700000103 " + "AND last_error_category='throttled'") == 1, + "timestamps and fixed failure category survive projection"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostname=" + "'db''quoted.abcdefghijkl.us-east-1.rds.amazonaws.com'") == 1, + "projection safely quotes endpoint text"); + + GloVars.prometheus_registry = std::make_shared(); + unlink("file:mem_mydb?mode=memory"); + { + MySQL_HostGroups_Manager hostgroups; + MyHGM = &hostgroups; + counting_provider = new CountingProvider(); + ok(install_global_aws_metadata_provider( + counting_provider, &destroy_counting_provider, nullptr), + "network-request counter installs through the production registry"); + + hostgroups.aws_locality_manager()->configure({disabled_hostgroup( + 101, "first.abcdefghijkl.us-east-1.rds.amazonaws.com", 7)}); + proxysql_refresh_configured_plugin_runtime_views( + "SELECT * FROM stats_mysql_aws_locality", nullptr, nullptr, &statsdb); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=101 " + "AND metadata_status='disabled' AND configured_weight=7 " + "AND effective_weight=7") == 1, + "real plugin callback projects the MySQL manager's current snapshot"); + ok(counting_provider->requests.load() == 0, + "query-time refresh issues no metadata-provider request"); + + hostgroups.aws_locality_manager()->configure({disabled_hostgroup( + 202, "second.abcdefghijkl.us-east-1.rds.amazonaws.com", 9)}); + proxysql_refresh_configured_plugin_runtime_views( + "SELECT * FROM stats_mysql_aws_locality", nullptr, nullptr, &statsdb); + ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 1 && + statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=202") == 1, + "generation swap replaces the prior projection without mixed rows"); + + hostgroups.aws_locality_manager()->configure({}); + proxysql_refresh_configured_plugin_runtime_views( + "SELECT * FROM stats_mysql_aws_locality", nullptr, nullptr, &statsdb); + ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 0, + "no valid locality policy produces zero rows"); + ok(counting_provider->requests.load() == 0, + "repeated generation queries remain network-free"); + MyHGM = nullptr; + } + unlink("file:mem_mydb?mode=memory"); + shutdown_global_aws_metadata_provider(); + GloVars.prometheus_registry.reset(); + + statsdb.execute("PRAGMA query_only = ON"); + char* write_error = nullptr; + SQLite3_result* write_result = statsdb.execute_statement( + "INSERT INTO stats_mysql_aws_locality " + "(hostgroup_id,hostname,port,endpoint_type,configured_weight,effective_weight," + "local_region,local_az,backend_region,backend_az,account_match,locality," + "active_multiplier,metadata_status,last_success_timestamp,last_attempt_timestamp," + "last_error_category) VALUES (1,'x',3306,'unknown',1,1,'','','',''," + "'unknown','unknown',1.0,'disabled',0,0,'')", &write_error); + ok(write_error != nullptr, + "stats listener query-only mode rejects writes to the projection"); + free(write_error); + delete write_result; + statsdb.execute("PRAGMA query_only = OFF"); + + proxysql_stop_configured_plugins(manager, error); + return exit_status(); +} From 8b1a7cfe76a074ebd9d80f088acfb84ef2a1f0f2 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Thu, 13 Aug 2026 21:29:50 +0000 Subject: [PATCH 10/17] docs: document AWS locality awareness --- README.md | 2 + doc/aws-locality-awareness.md | 253 ++++++++++++++++++ ...026-08-13-aws-locality-awareness-design.md | 2 +- 3 files changed, 256 insertions(+), 1 deletion(-) create mode 100644 doc/aws-locality-awareness.md diff --git a/README.md b/README.md index 714ce1d1b0..8c35e38f82 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,8 @@ tar xzf proxysql--linux-amd64.tar.gz The archive contains `bin/proxysql`, a sample `etc/proxysql.cnf`, the `systemd/` units, and helper tools. The v4.0 build additionally ships the runtime plugins under `lib/proxysql/` (`ProxySQL_MySQLX_Plugin.so`, `ProxySQL_GenAI_Plugin.so`). +See [AWS locality-aware backend selection](doc/aws-locality-awareness.md) for +the optional external-provider contract and MySQL configuration controls. Alternatively you can also use the available repositories: diff --git a/doc/aws-locality-awareness.md b/doc/aws-locality-awareness.md new file mode 100644 index 0000000000..de8f69578b --- /dev/null +++ b/doc/aws-locality-awareness.md @@ -0,0 +1,253 @@ +# AWS locality-aware MySQL backend selection + +ProxySQL 4.0 can prefer Amazon RDS and Aurora MySQL backends in the same AWS +Region or Availability Zone (AZ) as the ProxySQL process. The feature changes +only the temporary weights used by a server-selection attempt. It never +changes `mysql_servers.weight`, `runtime_mysql_servers.weight`, saved +configuration, or ProxySQL Cluster checksums. + +The MySQL module owns the switch and hostgroup policy. The optional general +AWS plugin provides local-instance and RDS metadata asynchronously. If the +plugin is not loaded, metadata is unavailable, or metadata expires, ProxySQL +continues using the configured weights. + +## Load the AWS plugin + +The AWS plugin is built and packaged by `PROXYSQL40=1 make -j`, with the +vendored AWS SDK linked statically into `ProxySQL_Aws_Plugin.so`. It is not +loaded automatically. Add its installed path to the `plugins` list in the +ProxySQL configuration and restart ProxySQL: + +```ini +plugins = ( + "/usr/lib/proxysql/ProxySQL_Aws_Plugin.so" +) +``` + +The plugin also provides [AWS IAM database authentication](aws_iam_database_authentication.md). +Both capabilities share one AWS SDK runtime and the standard AWS credential +provider chain. + +## Grant discovery permission + +Give the ProxySQL workload read-only access to the RDS discovery APIs: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "rds:DescribeDBInstances", + "rds:DescribeDBClusters", + "rds:DescribeDBClusterEndpoints" + ], + "Resource": "*" + } + ] +} +``` + +No long-lived access keys are required or recommended. On EC2, attach an +instance profile to the instance running ProxySQL. On EKS, use IRSA or EKS Pod +Identity for the ProxySQL pod. ECS task roles and other standard AWS SDK +credential sources also work. ProxySQL adds no access-key, secret-key, or +role-ARN variables. + +IMDSv2 local-location discovery itself needs no IAM API permission. RDS API +discovery uses the workload identity above. + +## Configure a hostgroup + +First add a locality policy to the existing +`mysql_hostgroup_attributes.hostgroup_settings` JSON. Both multipliers are +required: + +```sql +INSERT INTO mysql_hostgroup_attributes(hostgroup_id, hostgroup_settings) +VALUES ( + 10, + '{ + "aws": { + "locality_awareness": { + "same_region_multiplier": 2.0, + "same_az_multiplier": 4.0, + "refresh_interval_seconds": 300, + "stale_ttl_seconds": 1800 + } + } + }' +); + +LOAD MYSQL SERVERS TO RUNTIME; +SAVE MYSQL SERVERS TO DISK; +``` + +The accepted values are: + +```text +1.0 <= same_region_multiplier <= same_az_multiplier <= 10.0 + +refresh_interval_seconds default: 300 +stale_ttl_seconds default: 1800 + +30 <= refresh_interval_seconds <= 86400 +refresh_interval_seconds <= stale_ttl_seconds <= 604800 +``` + +The multipliers are floating-point JSON numbers. For each selection attempt, +ProxySQL calculates: + +```text +remote or unknown configured_weight +same Region, different AZ int(configured_weight * same_region_multiplier) +same AZ int(configured_weight * same_az_multiplier) +``` + +Conversion to an integer truncates toward zero. The tiers are not cumulative: +a same-AZ backend gets only the AZ multiplier, not the Region multiplier times +the AZ multiplier. Weight zero stays zero. + +An invalid `aws.locality_awareness` object disables locality bias for that +hostgroup when servers are loaded. The diagnostic identifies only the +hostgroup and rejected field; it does not print the JSON value. + +## Enable the global switch + +The one process-wide control is a MySQL module variable. It defaults off: + +```sql +SET mysql-aws_locality_awareness = true; +LOAD MYSQL VARIABLES TO RUNTIME; +SAVE MYSQL VARIABLES TO DISK; +``` + +There are deliberately no ProxySQL Region, AZ, or AWS-account variables. +Every ProxySQL process discovers its own location, so ProxySQL Cluster cannot +copy one node's location to another node in a different Region or AZ. + +Disabling the variable immediately restores ordinary configured-weight +selection, cancels or supersedes outstanding locality requests, and stops new +refresh scheduling. It does not remove or rebalance existing connections. + +## How local location is determined + +ProxySQL tries the EC2 IMDSv2 instance identity document first. If IMDSv2 is +unavailable, the AWS plugin uses these process environment values: + +1. Region: `AWS_REGION`, then `AWS_DEFAULT_REGION`; +2. AZ: `AWS_AVAILABILITY_ZONE`; +3. account assertion: `AWS_ACCOUNT_ID`. + +Region is the only required fallback field. AZ without Region is ignored. A +same-AZ preference also requires the same confirmed account ID on both sides, +because identical AZ names in different accounts may represent different +physical zones. Without an account assertion, same-Region preference remains +available but same-AZ preference does not. + +In Kubernetes, inject the node Region and AZ into each ProxySQL pod from the +node's topology labels and inject the account assertion through deployment +configuration. These are per-pod environment values, not clustered ProxySQL +variables. + +## Supported backend endpoints + +The first release supports official RDS and Aurora endpoints returned by the +AWS APIs: + +- RDS DB instance and Aurora DB instance endpoints can receive Region and AZ + preference; +- Aurora or Multi-AZ cluster writer endpoints receive Region preference only; +- Aurora reader and custom endpoints receive Region preference only. + +Cluster, reader, and custom endpoints can route to several AZs, so they never +receive the same-AZ multiplier. ProxySQL normalizes the configured hostname +and requires an exact match in a paginated RDS API response. Where AWS returns +a port, it must also match. + +Custom CNAMEs, arbitrary EC2 MySQL hosts, RDS Proxy endpoints, DNS aliases, +and AWS-looking names absent from the authoritative response remain neutral. +ProxySQL does not resolve CNAMEs, scan all Regions, assume roles into other +accounts, or change hostgroup membership. + +## Selection and failure behavior + +Locality is applied only after all normal eligibility checks, including +server state, capacity, latency, GTID, replication lag, session compatibility, +and existing backoff rules. It cannot make an unhealthy or incompatible +backend eligible. + +The global Hostgroup Manager path and the thread-local idle-connection path +use the same temporary effective server weights. The local cache groups +eligible connections by parent server first, so a server does not gain more +probability merely because it has more idle connections. + +Metadata states behave as follows: + +- `pending`: discovery has not completed; configured weight is used; +- `fresh`: the active Region or AZ multiplier is used; +- `stale`: the last success remains active through the bounded stale TTL; +- `expired`: configured weight is used; +- `error`: no usable value exists, so configured weight is used; +- `disabled`: the master switch is off and configured weight is used. + +A failed refresh does not immediately discard a prior success. Once its +`stale_ttl_seconds` expires, locality becomes neutral. Missing credentials, +access denial, throttling, timeout, missing plugin, unsupported endpoint, and +IMDS failure are all fail-neutral for database traffic. + +AWS, IMDS, credential-provider, DNS, and network work runs on bounded plugin +workers. No server-selection path calls the plugin or performs network I/O. + +## Inspect locality decisions + +When—and only when—the AWS plugin is loaded during startup, it registers the +read-only `stats_mysql_aws_locality` table. Querying it materializes one +current immutable manager snapshot into the stats database: + +```sql +SELECT * +FROM stats_mysql_aws_locality +ORDER BY hostgroup_id, hostname, port; +``` + +The columns are: + +```text +hostgroup_id hostname +port endpoint_type +configured_weight effective_weight +local_region local_az +backend_region backend_az +account_match locality +active_multiplier metadata_status +last_success_timestamp last_attempt_timestamp +last_error_category +``` + +`endpoint_type` is `instance`, `cluster`, `reader`, `custom`, or `unknown`. +`account_match` is `same`, `different`, or `unknown`; account IDs themselves +are never exposed. `locality` is `same_az`, `same_region`, `remote`, or +`unknown`. Timestamps are Unix epoch seconds, or zero if no corresponding +event has occurred. + +The table query never starts or waits for metadata discovery. It reports the +latest published state and replaces the previous rows in one transaction, so +generations cannot mix. The table is not persisted, loaded to runtime, +clustered, or included in a checksum. If no hostgroup has a valid policy, it +is empty. If the master switch is off, cached location text remains visible, +but every row reports `disabled`, multiplier `1.0`, and the configured weight +as its effective weight. + +Without the AWS plugin, the table is not registered and a query returns the +normal `no such table` error. ProxySQL does not currently hot-unload configured +plugins; changing the plugin list requires a restart. + +## Roll back + +Set `mysql-aws_locality_awareness` to `false` and load MySQL variables to +runtime. This restores the existing selection path immediately without +changing server rows. To remove a policy as well, remove the +`aws.locality_awareness` object from that hostgroup's `hostgroup_settings` and +load MySQL servers to runtime. diff --git a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md index 2d41459a1c..9991ab28e4 100644 --- a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md +++ b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md @@ -537,7 +537,7 @@ last_error_category Definitions: - `endpoint_type`: `instance`, `cluster`, `reader`, `custom`, or `unknown`; -- `account_match`: `yes`, `no`, `unknown`, or `not_applicable`; +- `account_match`: `same`, `different`, or `unknown`; - `locality`: `same_az`, `same_region`, `remote`, or `unknown`; - `active_multiplier`: the multiplier currently affecting selection, otherwise `1.0`; From f8a2f56f4ba6a38eb8066090e2be9f0a0db6c8b9 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Fri, 14 Aug 2026 05:28:57 +0000 Subject: [PATCH 11/17] test: sort AWS locality group entries --- test/tap/groups/groups.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index f365df1671..38071c7273 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -19,11 +19,11 @@ "aws_iam_kill_helper_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_iam_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_iam_pool_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], - "aws_iam_session_state_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], "aws_iam_provider_boundary_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], - "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], - "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_session_state_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], "aws_locality_config_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_selection_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_locality_stats_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], From ed8f0923f8036ceaaa8c616211084fec31be08b7 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Fri, 14 Aug 2026 06:33:45 +0000 Subject: [PATCH 12/17] test: publish locality manager in config fixture --- test/tap/tests/unit/aws_locality_config_unit-t.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/tap/tests/unit/aws_locality_config_unit-t.cpp b/test/tap/tests/unit/aws_locality_config_unit-t.cpp index f91a90999b..03f5068336 100644 --- a/test/tap/tests/unit/aws_locality_config_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_config_unit-t.cpp @@ -154,6 +154,8 @@ int main() { GloVars.prometheus_registry = std::make_shared(); { MySQL_HostGroups_Manager manager; + MySQL_HostGroups_Manager *previous_hgm = MyHGM; + MyHGM = &manager; srv_info_t info; info.addr = "db.abcdef.us-east-1.rds.amazonaws.com"; info.port = 3306; @@ -189,6 +191,7 @@ int main() { manager.refresh_aws_locality_configuration(); ok(manager.aws_locality_manager()->snapshot()->entries.empty(), "removing the hostgroup policy removes its manager configuration"); + MyHGM = previous_hgm; } GloVars.prometheus_registry.reset(); From a7d0954505a980d45d7dd1ea2744a1f7dddc68d7 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Fri, 14 Aug 2026 09:29:51 +0000 Subject: [PATCH 13/17] fix(mysql): harden AWS locality lifecycle and selection --- doc/aws-locality-awareness.md | 25 +-- .../2026-08-13-aws-locality-awareness.md | 2 +- include/Aws_Locality_Manager.h | 9 +- include/Aws_Locality_Types.h | 21 +- include/MySQL_Thread.h | 9 + include/ProxySQL_Plugin.h | 12 +- lib/Aws_Locality_Manager.cpp | 169 +++++++++----- lib/MyHGC.cpp | 48 ++-- lib/MySQL_HostGroups_Manager.cpp | 3 + lib/MySQL_Thread.cpp | 211 ++++++++++-------- lib/ProxySQL_PluginManager.cpp | 9 + plugins/aws/src/aws_locality_provider.cpp | 16 +- test/tap/tests/unit/Makefile | 7 +- .../tests/unit/aws_locality_config_unit-t.cpp | 6 +- .../unit/aws_locality_manager_unit-t.cpp | 148 +++++++++++- .../tests/unit/aws_locality_plugin_unit-t.cpp | 10 +- .../tests/unit/aws_locality_policy_unit-t.cpp | 14 +- .../unit/aws_locality_selection_unit-t.cpp | 85 ++++++- .../tests/unit/aws_locality_stats_unit-t.cpp | 39 +++- 19 files changed, 592 insertions(+), 251 deletions(-) diff --git a/doc/aws-locality-awareness.md b/doc/aws-locality-awareness.md index de8f69578b..7da4f7296e 100644 --- a/doc/aws-locality-awareness.md +++ b/doc/aws-locality-awareness.md @@ -66,19 +66,18 @@ required: ```sql INSERT INTO mysql_hostgroup_attributes(hostgroup_id, hostgroup_settings) -VALUES ( - 10, - '{ - "aws": { - "locality_awareness": { - "same_region_multiplier": 2.0, - "same_az_multiplier": 4.0, - "refresh_interval_seconds": 300, - "stale_ttl_seconds": 1800 - } - } - }' -); +VALUES (10, json_object( + 'aws', json_object( + 'locality_awareness', json_object( + 'same_region_multiplier', 2.0, + 'same_az_multiplier', 4.0, + 'refresh_interval_seconds', 300, + 'stale_ttl_seconds', 1800)))) +ON CONFLICT(hostgroup_id) DO UPDATE SET + hostgroup_settings = json_set( + COALESCE(mysql_hostgroup_attributes.hostgroup_settings, '{}'), + '$.aws.locality_awareness', + json_extract(excluded.hostgroup_settings, '$.aws.locality_awareness')); LOAD MYSQL SERVERS TO RUNTIME; SAVE MYSQL SERVERS TO DISK; diff --git a/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md b/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md index 86e42cfcd4..f3744c16b4 100644 --- a/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md +++ b/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md @@ -326,7 +326,7 @@ - [ ] **Step 3: Extend the ABI and shared SDK lifetime** - Increment the plugin ABI maximum/current version to 6 and append metadata-provider installation to `ProxySQL_PluginServices`. Wire it only during plugin init. Refactor the plugin so IAM signer/token source and locality provider each retain a `std::shared_ptr`; `Aws::InitAPI` occurs once and `Aws::ShutdownAPI` occurs only after both core-owned capabilities drain. + Increment the plugin ABI maximum/current version to 6 for metadata-provider installation, to 7 for the MySQL-owned locality-stats projection callback, and to 8 for partial-init provider rollback; append those services to `ProxySQL_PluginServices`. Wire provider installation and rollback only during plugin init. Refactor the plugin so IAM signer/token source and locality provider each retain a `std::shared_ptr`; `Aws::InitAPI` occurs once and `Aws::ShutdownAPI` occurs only after both core-owned capabilities drain. - [ ] **Step 4: Implement local discovery** diff --git a/include/Aws_Locality_Manager.h b/include/Aws_Locality_Manager.h index 45a380184d..cf7b0f18e2 100644 --- a/include/Aws_Locality_Manager.h +++ b/include/Aws_Locality_Manager.h @@ -1,5 +1,5 @@ -#ifndef AWS_LOCALITY_MANAGER_H -#define AWS_LOCALITY_MANAGER_H +#ifndef __CLASS_AWS_LOCALITY_MANAGER_H +#define __CLASS_AWS_LOCALITY_MANAGER_H #include "Aws_Locality_Types.h" #include "json_fwd.hpp" @@ -90,7 +90,7 @@ struct AwsLocalitySnapshotEntry { struct AwsLocalitySnapshot { uint64_t generation { 0 }; bool enabled { false }; - std::unordered_map entries; + std::unordered_multimap entries; std::unordered_set hostgroups; const AwsLocalitySnapshotEntry* find( @@ -115,6 +115,7 @@ struct AwsLocalityManagerConfig { SteadyClock steady_clock; WallClock wall_clock; std::chrono::milliseconds request_timeout { std::chrono::seconds(5) }; + std::chrono::milliseconds disable_wait_timeout { std::chrono::milliseconds(250) }; std::function before_completion; }; @@ -137,4 +138,4 @@ class MySQLAwsLocalityManager { std::unique_ptr impl_; }; -#endif // AWS_LOCALITY_MANAGER_H +#endif // __CLASS_AWS_LOCALITY_MANAGER_H diff --git a/include/Aws_Locality_Types.h b/include/Aws_Locality_Types.h index 720d7db96c..47ff5a91f0 100644 --- a/include/Aws_Locality_Types.h +++ b/include/Aws_Locality_Types.h @@ -1,13 +1,28 @@ -#ifndef AWS_LOCALITY_TYPES_H -#define AWS_LOCALITY_TYPES_H +#ifndef __CLASS_AWS_LOCALITY_TYPES_H +#define __CLASS_AWS_LOCALITY_TYPES_H #include +#include #include #include #include +#include #include #include +inline std::string aws_locality_normalized_hostname(std::string_view input) { + if (input.empty()) return {}; + std::string hostname(input); + if (hostname.back() == '.') hostname.pop_back(); + if (hostname.empty() || hostname.back() == '.') return {}; + for (char& character : hostname) { + const unsigned char value = static_cast(character); + if (!(std::isalnum(value) || character == '-' || character == '.')) return {}; + character = static_cast(std::tolower(value)); + } + return hostname; +} + enum class AwsEndpointType : uint8_t { unknown, instance, @@ -151,4 +166,4 @@ struct AwsLocalityHostgroupConfig { std::vector backends; }; -#endif // AWS_LOCALITY_TYPES_H +#endif // __CLASS_AWS_LOCALITY_TYPES_H diff --git a/include/MySQL_Thread.h b/include/MySQL_Thread.h index 6cbf434aba..aed8b94677 100644 --- a/include/MySQL_Thread.h +++ b/include/MySQL_Thread.h @@ -205,6 +205,15 @@ class __attribute__((aligned(64))) MySQL_Thread : public Base_Thread PtrArray *cached_connections; unsigned int push_local_counter; // round-robin counter for bounded local caching: cache 1-in-N where N = mysql_threads +#ifdef PROXYSQL40 + struct AwsLocalityCachedCandidate { + MySrvC* parent; + unsigned int cached_index; + }; + // Capacity grows when a connection enters the local cache, never while a + // query is choosing a backend from that cache. + std::vector aws_locality_candidates; +#endif #ifdef IDLE_THREADS struct epoll_event events[MY_EPOLL_THREAD_MAXEVENTS]; diff --git a/include/ProxySQL_Plugin.h b/include/ProxySQL_Plugin.h index f5040102da..1b6cd36408 100644 --- a/include/ProxySQL_Plugin.h +++ b/include/ProxySQL_Plugin.h @@ -45,8 +45,10 @@ namespace prometheus { class Registry; } // installation callback used by locality discovery. // ABI 7: ProxySQL_PluginServices gains the MySQL-owned AWS-locality stats // projection callback used by the AWS plugin's runtime view. -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 7u; -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 7u; +// ABI 8: ProxySQL_PluginServices gains an IAM-provider uninstall callback +// so a plugin can roll back a partially successful init. +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 8u; +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 8u; enum class ProxySQL_PluginDBKind : uint8_t { admin_db = 0, @@ -245,6 +247,9 @@ using proxysql_plugin_register_runtime_view_cb = using proxysql_plugin_install_aws_iam_token_source_cb = bool (*)(AwsIamTokenSource *, void (*)(AwsIamTokenSource *), void *module_handle); +using proxysql_plugin_uninstall_aws_iam_token_source_cb = + bool (*)(AwsIamTokenSource *expected_source); + using proxysql_plugin_get_aws_iam_limits_cb = void (*)(size_t *max_total_waiters, size_t *max_waiters_per_key); @@ -323,6 +328,9 @@ struct ProxySQL_PluginServices { proxysql_plugin_install_aws_metadata_provider_cb install_aws_metadata_provider; // ABI-7 extension. Live in Phase B and normal init; it performs no I/O. proxysql_plugin_refresh_mysql_aws_locality_stats_cb refresh_mysql_aws_locality_stats; + // ABI-8 extension. Live only during normal plugin init and intended for + // rollback of the same plugin's partially installed IAM provider. + proxysql_plugin_uninstall_aws_iam_token_source_cb uninstall_aws_iam_token_source; #endif /* PROXYSQL40 */ }; diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp index 4feff419e4..bbc469d3d9 100644 --- a/lib/Aws_Locality_Manager.cpp +++ b/lib/Aws_Locality_Manager.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -51,7 +52,7 @@ bool read_seconds( const auto it = object.find(field); if (it == object.end()) { value = default_value; - return true; + return value >= minimum && value <= maximum; } if (!it->is_number_unsigned() && !it->is_number_integer()) { return false; @@ -66,29 +67,6 @@ bool read_seconds( return true; } -std::string normalized_hostname(std::string_view input) { - if (input.empty()) { - return {}; - } - - std::string hostname(input); - if (hostname.back() == '.') { - hostname.pop_back(); - } - if (hostname.empty() || hostname.back() == '.') { - return {}; - } - - for (char& character : hostname) { - const unsigned char value = static_cast(character); - if (!(std::isalnum(value) || character == '-' || character == '.')) { - return {}; - } - character = static_cast(std::tolower(value)); - } - return hostname; -} - bool ends_with(const std::string& value, const std::string& suffix) { return value.size() >= suffix.size() && value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0; @@ -112,21 +90,15 @@ bool valid_region(const std::string& region) { std::isdigit(static_cast(region.back())); } -bool contains_proxy_label(const std::string& prefix) { - size_t begin = 0; - while (begin < prefix.size()) { - const size_t end = prefix.find('.', begin); - const size_t length = end == std::string::npos - ? prefix.size() - begin : end - begin; - if (length >= 6 && prefix.compare(begin, 6, "proxy-") == 0) { - return true; - } - if (end == std::string::npos) { - break; - } - begin = end + 1; - } - return false; +bool is_rds_proxy_endpoint_prefix(const std::string& prefix) { + // RDS Proxy endpoints have the canonical shape + // .proxy-..rds.amazonaws.com. A normal + // DB identifier is allowed to begin with "proxy-", so only the reserved + // generated-ID label after a proxy name identifies this endpoint type. + const size_t separator = prefix.rfind('.'); + return separator != std::string::npos && separator != 0 && + prefix.size() - separator - 1 > 6 && + prefix.compare(separator + 1, 6, "proxy-") == 0; } } // namespace @@ -171,7 +143,7 @@ AwsEndpointCandidate recognize_rds_endpoint( AwsEndpointCandidate result; result.hostgroup_id = hostgroup_id; result.port = port; - result.hostname = normalized_hostname(hostname_input); + result.hostname = aws_locality_normalized_hostname(hostname_input); if (result.hostname.empty()) { return result; } @@ -197,7 +169,7 @@ AwsEndpointCandidate recognize_rds_endpoint( const std::string endpoint_prefix = before_suffix.substr(0, region_separator); result.region = before_suffix.substr(region_separator + 1); - if (contains_proxy_label(endpoint_prefix) || !valid_region(result.region)) { + if (is_rds_proxy_endpoint_prefix(endpoint_prefix) || !valid_region(result.region)) { result.region.clear(); return result; } @@ -291,17 +263,66 @@ void* metadata_provider_module = nullptr; size_t metadata_provider_leases = 0; bool metadata_provider_accepting = false; -std::string snapshot_key( +bool dns_identity_length(std::string_view input, size_t& length) { + if (input.empty()) return false; + length = input.size(); + if (input[length - 1] == '.') --length; + if (length == 0 || input[length - 1] == '.') return false; + for (size_t index = 0; index < length; ++index) { + const unsigned char value = static_cast(input[index]); + if (!(std::isalnum(value) || input[index] == '-' || input[index] == '.')) { + return false; + } + } + return true; +} + +uint64_t identity_hash( uint32_t hostgroup_id, - std::string_view hostname_input, + std::string_view hostname, uint16_t port) { - const std::string hostname = normalized_hostname(hostname_input); - return std::to_string(hostgroup_id) + "\n" + hostname + "\n" + - std::to_string(port); + uint64_t hash = 14695981039346656037ULL; + auto append = [&](unsigned char value) { + hash ^= value; + hash *= 1099511628211ULL; + }; + for (unsigned int shift = 0; shift < 32; shift += 8) { + append(static_cast(hostgroup_id >> shift)); + } + append(static_cast(port)); + append(static_cast(port >> 8)); + size_t length = 0; + const bool valid_dns = dns_identity_length(hostname, length); + append(valid_dns ? 1 : 0); + if (!valid_dns) length = hostname.size(); + for (size_t index = 0; index < length; ++index) { + const unsigned char value = static_cast(hostname[index]); + append(valid_dns ? static_cast(std::tolower(value)) : value); + } + return hash; +} + +bool same_hostname_identity(std::string_view lhs, std::string_view rhs) { + size_t lhs_length = 0; + size_t rhs_length = 0; + const bool lhs_dns = dns_identity_length(lhs, lhs_length); + const bool rhs_dns = dns_identity_length(rhs, rhs_length); + if (lhs_dns != rhs_dns) return false; + if (!lhs_dns) return lhs == rhs; + if (lhs_length != rhs_length) return false; + for (size_t index = 0; index < lhs_length; ++index) { + if (std::tolower(static_cast(lhs[index])) != + std::tolower(static_cast(rhs[index]))) { + return false; + } + } + return true; } std::string endpoint_key(std::string_view hostname_input, uint16_t port) { - return normalized_hostname(hostname_input) + "\n" + std::to_string(port); + const std::string hostname = aws_locality_normalized_hostname(hostname_input); + return (hostname.empty() ? "raw:" + std::string(hostname_input) + : "dns:" + hostname) + "\n" + std::to_string(port); } const char* failure_category(AwsMetadataStatus status) { @@ -420,8 +441,15 @@ const AwsLocalitySnapshotEntry* AwsLocalitySnapshot::find( uint32_t hostgroup_id, std::string_view hostname, uint16_t port) const { - const auto it = entries.find(snapshot_key(hostgroup_id, hostname, port)); - return it == entries.end() ? nullptr : &it->second; + const auto range = entries.equal_range(identity_hash(hostgroup_id, hostname, port)); + for (auto it = range.first; it != range.second; ++it) { + const auto& entry = it->second; + if (entry.hostgroup_id == hostgroup_id && entry.port == port && + same_hostname_identity(entry.hostname, hostname)) { + return &entry; + } + } + return nullptr; } uint64_t AwsLocalitySnapshot::effective_weight( @@ -448,8 +476,29 @@ class MySQLAwsLocalityManager::Impl { std::memory_order_release); } - ~Impl() { - shutdown(); + ~Impl() noexcept { + try { + shutdown(); + } catch (...) { + // Destructors must not let allocation/system exceptions from the + // final snapshot publication escape. Ensure the scheduler thread + // cannot make std::thread's destructor terminate the process. + try { + std::thread worker; + { + std::lock_guard lock(mutex_); + stopping_ = true; + enabled_ = false; + cancel_requested_ = true; + cv_.notify_all(); + worker = std::move(worker_); + } + if (worker.joinable()) worker.join(); + sink_->detach(); + } catch (...) { + std::terminate(); + } + } } void configure(std::vector hostgroups) { @@ -471,6 +520,9 @@ class MySQLAwsLocalityManager::Impl { ++generation_; cancel_requested_ = true; force_refresh_ = enabled_ && !hostgroups_.empty(); + if (force_refresh_) { + ensure_worker_locked(); + } publish_locked(); cv_.notify_all(); } @@ -491,7 +543,7 @@ class MySQLAwsLocalityManager::Impl { publish_locked(); cv_.notify_all(); if (!enabled_ && worker_.joinable()) { - cv_.wait(lock, [&] { + cv_.wait_for(lock, config_.disable_wait_timeout, [&] { return disable_acknowledged_ || stopping_; }); } @@ -742,7 +794,10 @@ class MySQLAwsLocalityManager::Impl { break; } if (!enabled_ || hostgroups_.empty()) { - cv_.wait(lock); + cv_.wait(lock, [&] { + return stopping_ || cancel_requested_ || force_refresh_ || + (enabled_ && !hostgroups_.empty()); + }); continue; } const auto now = config_.steady_clock(); @@ -900,11 +955,11 @@ class MySQLAwsLocalityManager::Impl { std::unordered_map returned; for (const auto& endpoint : result.endpoints) { - const std::string hostname = normalized_hostname(endpoint.hostname); + const std::string hostname = aws_locality_normalized_hostname(endpoint.hostname); if (!hostname.empty() && endpoint.region == request.region) { returned[endpoint_key(hostname, endpoint.port)] = &endpoint; if (endpoint.port == 0) { - returned[hostname + "\n0"] = &endpoint; + returned[endpoint_key(hostname, 0)] = &endpoint; } } } @@ -913,7 +968,7 @@ class MySQLAwsLocalityManager::Impl { auto& record = endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)]; record.attempt_wall = now_wall; const auto exact = returned.find(endpoint_key(endpoint.hostname, endpoint.port)); - const auto no_port = returned.find(normalized_hostname(endpoint.hostname) + "\n0"); + const auto no_port = returned.find(endpoint_key(endpoint.hostname, 0)); const AwsMetadataEndpoint* match = exact != returned.end() ? exact->second : (no_port != returned.end() ? no_port->second : nullptr); if (match == nullptr || match->endpoint_type == AwsEndpointType::unknown) { @@ -1018,7 +1073,7 @@ class MySQLAwsLocalityManager::Impl { next->hostgroups.insert(hostgroup.hostgroup_id); for (const auto& backend : hostgroup.backends) { auto entry = build_entry_locked(hostgroup, backend, now); - next->entries.emplace(snapshot_key(entry.hostgroup_id, + next->entries.emplace(identity_hash(entry.hostgroup_id, entry.hostname, entry.port), std::move(entry)); } } diff --git a/lib/MyHGC.cpp b/lib/MyHGC.cpp index 6deec50afd..df54889d7e 100644 --- a/lib/MyHGC.cpp +++ b/lib/MyHGC.cpp @@ -48,15 +48,11 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ auto candidate_weight_sum = [&]() -> uint64_t { #ifdef PROXYSQL40 if (use_aws_locality) { - uint64_t effective_sum = 0; - for (unsigned int candidate = 0; candidate < num_candidates; ++candidate) { - MySrvC* server = mysrvcCandidates[candidate]; - effective_sum = aws_locality_saturating_add( - effective_sum, - aws_locality_snapshot->effective_weight( - hid, server->address, server->port, server->weight)); - } - return effective_sum; + // Locality selection has a uniform fallback for an all-zero + // configured-weight set. These availability checks therefore only + // need to know whether an eligible candidate exists; defer snapshot + // lookups until the actual lottery below. + return num_candidates; } #endif return sum; @@ -356,24 +352,32 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ #ifdef PROXYSQL40 if (use_aws_locality) { - uint64_t locality_weights_static[32]; - uint64_t* locality_weights = locality_weights_static; - if (num_candidates > 32) { - locality_weights = static_cast( - malloc(sizeof(uint64_t) * num_candidates)); - } + uint64_t total_weight = 0; for (j = 0; j < num_candidates; ++j) { mysrvc = mysrvcCandidates[j]; - locality_weights[j] = aws_locality_snapshot->effective_weight( - hid, mysrvc->address, mysrvc->port, mysrvc->weight); + total_weight = aws_locality_saturating_add(total_weight, + aws_locality_snapshot->effective_weight( + hid, mysrvc->address, mysrvc->port, mysrvc->weight)); } const uint64_t random_value = (static_cast(rand_fast()) << 32) | static_cast(rand_fast()); - const size_t selected = aws_locality_weighted_index( - locality_weights, num_candidates, random_value); - if (num_candidates > 32) { - free(locality_weights); + size_t selected = num_candidates; + if (total_weight == 0 && num_candidates != 0) { + selected = random_value % num_candidates; + } else if (total_weight != 0) { + const uint64_t target = random_value % total_weight; + uint64_t cumulative = 0; + for (j = 0; j < num_candidates; ++j) { + mysrvc = mysrvcCandidates[j]; + cumulative = aws_locality_saturating_add(cumulative, + aws_locality_snapshot->effective_weight( + hid, mysrvc->address, mysrvc->port, mysrvc->weight)); + if (target < cumulative) { + selected = j; + break; + } + } } if (selected < num_candidates) { mysrvc = mysrvcCandidates[selected]; @@ -389,7 +393,7 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ return mysrvc; } proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, - "Returning MySrvC NULL because AWS locality weights are zero\n"); + "Returning MySrvC NULL because no AWS locality candidate is eligible\n"); if (l>32) { free(mysrvcCandidates); } diff --git a/lib/MySQL_HostGroups_Manager.cpp b/lib/MySQL_HostGroups_Manager.cpp index 8591fe12ab..29fc23d865 100644 --- a/lib/MySQL_HostGroups_Manager.cpp +++ b/lib/MySQL_HostGroups_Manager.cpp @@ -926,11 +926,14 @@ bool aws_locality_status_is_active(AwsLocalityMetadataStatus status) { status == AwsLocalityMetadataStatus::stale; } +std::mutex aws_locality_stats_projection_mutex; + } // namespace bool MySQL_HostGroups_Manager::project_aws_locality_stats( SQLite3DB* statsdb, const std::vector& rows) { + std::lock_guard projection_lock(aws_locality_stats_projection_mutex); if (statsdb == nullptr || !statsdb->execute("BEGIN")) return false; bool success = statsdb->execute("DELETE FROM stats_mysql_aws_locality"); for (const auto& row : rows) { diff --git a/lib/MySQL_Thread.cpp b/lib/MySQL_Thread.cpp index 7bae61638a..fcc8aca8d8 100644 --- a/lib/MySQL_Thread.cpp +++ b/lib/MySQL_Thread.cpp @@ -6944,21 +6944,8 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local( } #ifdef PROXYSQL40 - struct AwsLocalityParentCandidate { - MySrvC* parent; - MySQL_Connection* connection; - uint64_t weight; - }; - AwsLocalityParentCandidate candidates_static[32]; - AwsLocalityParentCandidate* candidates = candidates_static; - const unsigned int candidate_capacity = cached_connections->len; - std::vector candidates_dynamic; - if (candidate_capacity > 32) { - candidates_dynamic.resize(candidate_capacity); - candidates = candidates_dynamic.data(); - } - unsigned int num_candidates = 0; - + // Remove mode-incompatible connections before the allocation-free scoring + // passes below, so their indices remain stable throughout the lottery. for (i = 0; i < cached_connections->len;) { c = static_cast(cached_connections->index(i)); const char* candidate_username = @@ -6976,109 +6963,128 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local( MyHGM->destroy_MyConn_from_pool(c); continue; } - if (c->backend_auth_type() != requested_type || - (requested_type == MySQLBackendAuthType::AWS_IAM && - c->requires_CHANGE_USER(client_conn, requested_type))) { - ++i; - continue; - } - if (!c->healthy || !c->reusable) { - ++i; - continue; - } - if (check_session_track_backoff) { - session_track_backoff_until = - c->parent->session_track_backoff_until.load(std::memory_order_relaxed); - if (session_track_backoff_until > curtime) { - ++i; - continue; - } - } - if (c->parent->myhgc->hid != _hid || - !client_conn->match_tracked_options(c)) { - ++i; - continue; - } + ++i; + } - MySrvC* parent = c->parent; - if (find(parents.begin(), parents.end(), parent) != parents.end()) { - ++i; - continue; - } - bool parent_already_selected = false; - for (unsigned int candidate = 0; candidate < num_candidates; ++candidate) { - if (candidates[candidate].parent == parent) { - parent_already_selected = true; - break; - } - } - if (parent_already_selected) { - ++i; - continue; + auto connection_is_eligible = [&](MySQL_Connection* candidate, + bool record_lag_skip) -> bool { + if (candidate->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + candidate->requires_CHANGE_USER(client_conn, requested_type)) || + !candidate->healthy || !candidate->reusable) { + return false; } - if (gtid_uuid != nullptr && - !MyHGM->gtid_exists(parent, gtid_uuid, gtid_trxid)) { - parents.push_back(parent); - ++i; - continue; + if (check_session_track_backoff && + candidate->parent->session_track_backoff_until.load( + std::memory_order_relaxed) > curtime) { + return false; } - if (c->requires_CHANGE_USER(client_conn, requested_type)) { - ++i; - continue; + if (candidate->parent->myhgc->hid != _hid || + !client_conn->match_tracked_options(candidate) || + candidate->requires_CHANGE_USER(client_conn, requested_type)) { + return false; } char* schema = client_conn->userinfo->schemaname; - if (strcmp(c->userinfo->schemaname, schema) != 0) { - ++i; - continue; + if (strcmp(candidate->userinfo->schemaname, schema) != 0) { + return false; } unsigned int not_match = 0; - c->number_of_matching_session_variables(client_conn, not_match); + candidate->number_of_matching_session_variables(client_conn, not_match); if (not_match != 0) { - ++i; - continue; + return false; } - if (max_lag_ms >= 0 && + if (gtid_uuid == nullptr && max_lag_ms >= 0 && static_cast(max_lag_ms) < - (parent->aws_aurora_current_lag_us / 1000)) { - status_variables.stvar[st_var_aws_aurora_replicas_skipped_during_query]++; - ++i; - continue; + (candidate->parent->aws_aurora_current_lag_us / 1000)) { + if (record_lag_skip) { + status_variables.stvar[ + st_var_aws_aurora_replicas_skipped_during_query]++; + } + return false; } + return true; + }; - candidates[num_candidates++] = { - parent, - c, - aws_locality_snapshot->effective_weight( - _hid, parent->address, parent->port, parent->weight) - }; - ++i; + aws_locality_candidates.clear(); + // push_MyConn_local() grows this reusable storage before inserting into the + // cache. A direct cache mutation would violate that invariant; fail neutral + // instead of allocating in the selection path. + if (aws_locality_candidates.capacity() < cached_connections->len) { + return NULL; } - - uint64_t locality_weights_static[32]; - uint64_t* locality_weights = locality_weights_static; - std::vector locality_weights_dynamic; - if (num_candidates > 32) { - locality_weights_dynamic.resize(num_candidates); - locality_weights = locality_weights_dynamic.data(); + for (i = 0; i < cached_connections->len; ++i) { + auto* candidate = static_cast(cached_connections->index(i)); + if (connection_is_eligible(candidate, true)) { + aws_locality_candidates.push_back({candidate->parent, i}); + } } - for (i = 0; i < num_candidates; ++i) { - locality_weights[i] = candidates[i].weight; + std::sort(aws_locality_candidates.begin(), aws_locality_candidates.end(), + [](const AwsLocalityCachedCandidate& lhs, + const AwsLocalityCachedCandidate& rhs) { + if (lhs.parent != rhs.parent) { + return std::less()(lhs.parent, rhs.parent); + } + return lhs.cached_index < rhs.cached_index; + }); + + uint64_t total_weight = 0; + unsigned int num_candidates = 0; + for (i = 0; i < aws_locality_candidates.size();) { + MySrvC* parent = aws_locality_candidates[i].parent; + unsigned int next = i + 1; + while (next < aws_locality_candidates.size() && + aws_locality_candidates[next].parent == parent) { + ++next; + } + if (gtid_uuid != nullptr && !MyHGM->gtid_exists(parent, gtid_uuid, gtid_trxid)) { + i = next; + continue; + } + total_weight = aws_locality_saturating_add(total_weight, + aws_locality_snapshot->effective_weight( + _hid, parent->address, parent->port, parent->weight)); + ++num_candidates; + i = next; } const uint64_t random_value = (static_cast(rand_fast()) << 32) | static_cast(rand_fast()); - const size_t selected = aws_locality_weighted_index( - locality_weights, num_candidates, random_value); - MySQL_Connection* selected_connection = - selected < num_candidates ? candidates[selected].connection : nullptr; - if (selected_connection == nullptr) { + if (num_candidates == 0) { return NULL; } - for (i = 0; i < cached_connections->len; ++i) { - if (cached_connections->index(i) == selected_connection) { + const bool uniform_fallback = total_weight == 0; + const uint64_t target = uniform_fallback + ? random_value % num_candidates : random_value % total_weight; + uint64_t cumulative = 0; + unsigned int ordinal = 0; + for (i = 0; i < aws_locality_candidates.size();) { + const auto& candidate = aws_locality_candidates[i]; + MySrvC* parent = candidate.parent; + unsigned int next = i + 1; + while (next < aws_locality_candidates.size() && + aws_locality_candidates[next].parent == parent) { + ++next; + } + if (gtid_uuid != nullptr && !MyHGM->gtid_exists(parent, gtid_uuid, gtid_trxid)) { + i = next; + continue; + } + if (uniform_fallback) { + if (ordinal++ == target) { + return static_cast( + cached_connections->remove_index_fast(candidate.cached_index)); + } + i = next; + continue; + } + cumulative = aws_locality_saturating_add(cumulative, + aws_locality_snapshot->effective_weight( + _hid, parent->address, parent->port, parent->weight)); + if (target < cumulative) { return static_cast( - cached_connections->remove_index_fast(i)); + cached_connections->remove_index_fast(candidate.cached_index)); } + i = next; } #endif return NULL; @@ -7116,8 +7122,17 @@ void MySQL_Thread::push_MyConn_local(MySQL_Connection *c) { if (mysrvc->get_status() == MYSQL_SERVER_STATUS_ONLINE) { if (c->async_state_machine==ASYNC_IDLE) { unsigned int n = (GloMTH && GloMTH->num_threads > 0) ? GloMTH->num_threads : 1; - if ((push_local_counter++ % n) == 0) { - cached_connections->add(c); + if ((push_local_counter++ % n) == 0) { +#ifdef PROXYSQL40 + const size_t required_capacity = cached_connections->len + 1; + if (aws_locality_candidates.capacity() < required_capacity) { + const size_t grown_capacity = std::max( + 32, std::max(required_capacity, + aws_locality_candidates.capacity() * 2)); + aws_locality_candidates.reserve(grown_capacity); + } +#endif + cached_connections->add(c); return; } } diff --git a/lib/ProxySQL_PluginManager.cpp b/lib/ProxySQL_PluginManager.cpp index 7b5e96c130..b581ffb217 100644 --- a/lib/ProxySQL_PluginManager.cpp +++ b/lib/ProxySQL_PluginManager.cpp @@ -191,6 +191,14 @@ bool install_aws_iam_token_source_service( return install_global_aws_iam_token_source(source, destroy, module_handle); } +bool uninstall_aws_iam_token_source_service(AwsIamTokenSource *expected_source) { + if (g_registry_target == nullptr) { + proxy_warning("AWS IAM token source removal attempted outside plugin init phase\n"); + return false; + } + return uninstall_global_aws_iam_token_source(expected_source); +} + void get_aws_iam_limits_service(size_t *max_total_waiters, size_t *max_waiters_per_key) { const size_t maximum = GloMTH != nullptr && GloMTH->variables.max_connections > 0 ? static_cast(GloMTH->variables.max_connections) @@ -348,6 +356,7 @@ ProxySQL_PluginManager::ProxySQL_PluginManager() { services_.get_aws_iam_limits = &get_aws_iam_limits_service; services_.install_aws_metadata_provider = &install_aws_metadata_provider_service; services_.refresh_mysql_aws_locality_stats = &refresh_mysql_aws_locality_stats_service; + services_.uninstall_aws_iam_token_source = &uninstall_aws_iam_token_source_service; // Phase-B (register_schemas) services: same layout as init(), but DB // handle getters and the query-hook registrar are stubbed -- see the diff --git a/plugins/aws/src/aws_locality_provider.cpp b/plugins/aws/src/aws_locality_provider.cpp index ce4416528e..74b2de42b2 100644 --- a/plugins/aws/src/aws_locality_provider.cpp +++ b/plugins/aws/src/aws_locality_provider.cpp @@ -62,19 +62,6 @@ void normalize_failure(AwsMetadataResult& result) { } } -std::string normalized_hostname(const std::string& input) { - if (input.empty()) return {}; - std::string hostname = input; - if (hostname.back() == '.') hostname.pop_back(); - if (hostname.empty() || hostname.back() == '.') return {}; - for (char& character : hostname) { - const unsigned char value = static_cast(character); - if (!(std::isalnum(value) || character == '-' || character == '.')) return {}; - character = static_cast(std::tolower(value)); - } - return hostname; -} - bool valid_location_value(const std::string& value, size_t maximum) { if (value.empty() || value.size() > maximum) return false; for (const unsigned char character : value) { @@ -453,7 +440,7 @@ AwsMetadataResult AwsLocalityRdsDiscovery::discover( std::unordered_map endpoint_indices; auto append = [&](const std::string& hostname_input, int port, AwsEndpointType type, const std::string& az, const std::string& account) { - const std::string hostname = normalized_hostname(hostname_input); + const std::string hostname = aws_locality_normalized_hostname(hostname_input); if (hostname.empty() || type == AwsEndpointType::unknown || port < 0 || port > 65535) return; AwsMetadataEndpoint endpoint; @@ -629,7 +616,6 @@ AwsImdsResponse imds_request( curl_easy_setopt(handle, CURLOPT_NOBODY, 0L); curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 0L); curl_easy_setopt(handle, CURLOPT_NOPROXY, "*"); - curl_easy_setopt(handle, CURLOPT_PROXY, ""); curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, std::min(timeout, 500L)); curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, timeout); curl_easy_setopt(handle, CURLOPT_NOSIGNAL, 1L); diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index 3a41f62a4a..95534ba10a 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -419,7 +419,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t aws_locality_config_unit-t aws_locality_selection_unit-t aws_locality_stats_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ @@ -495,6 +495,9 @@ UNIT_TESTS += \ plugin_prometheus_unit-t \ plugin_lifecycle_unit-t \ plugin_runtime_views_unit-t \ + aws_locality_config_unit-t \ + aws_locality_selection_unit-t \ + aws_locality_stats_unit-t \ test_mysqlx_plugin_load-t \ test_mysqlx_admin_tables-t \ mysqlx_config_store_unit-t \ @@ -939,7 +942,7 @@ aws_locality_policy_unit-t: aws_locality_policy_unit-t.cpp \ $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o $(CXX) $< $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp \ $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ - $(IDIRS) $(OPT) -lpthread -o $@ + $(IDIRS) $(OPT) -lpthread -ldl -o $@ aws_locality_manager_unit-t: aws_locality_manager_unit-t.cpp \ $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o diff --git a/test/tap/tests/unit/aws_locality_config_unit-t.cpp b/test/tap/tests/unit/aws_locality_config_unit-t.cpp index 03f5068336..c31ac9f254 100644 --- a/test/tap/tests/unit/aws_locality_config_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_config_unit-t.cpp @@ -69,7 +69,7 @@ std::string capture_invalid_policy_log(MyHGC& hostgroup) { } // namespace int main() { - plan(21); + plan(20); MyHGC hostgroup(42); init_myhgc_hostgroup_settings( @@ -147,10 +147,6 @@ int main() { } test_globals_cleanup(); - auto refresh_method = &MySQL_HostGroups_Manager::refresh_aws_locality_configuration; - ok(refresh_method != nullptr, - "Hostgroup Manager exposes the post-commit locality refresh boundary"); - GloVars.prometheus_registry = std::make_shared(); { MySQL_HostGroups_Manager manager; diff --git a/test/tap/tests/unit/aws_locality_manager_unit-t.cpp b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp index f40ed2c51a..6ffc0d6176 100644 --- a/test/tap/tests/unit/aws_locality_manager_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp @@ -31,6 +31,9 @@ struct FakeProviderState { uint64_t next_handle { 1 }; bool shut_down { false }; bool destroyed { false }; + bool block_requests { false }; + bool request_entered { false }; + bool release_requests { false }; }; class FakeProvider final : public AwsMetadataProvider { @@ -47,10 +50,14 @@ class FakeProvider final : public AwsMetadataProvider { AwsMetadataRequestHandle request( const AwsMetadataRequest& request, std::weak_ptr sink) override { - std::lock_guard lock(state_->mutex); + std::unique_lock lock(state_->mutex); const AwsMetadataRequestHandle handle { state_->next_handle++ }; state_->requests.push_back({handle, request, std::move(sink)}); + state_->request_entered = true; state_->cv.notify_all(); + state_->cv.wait(lock, [&] { + return !state_->block_requests || state_->release_requests; + }); return handle; } @@ -122,6 +129,37 @@ bool complete_request( return false; } +bool complete_request_after( + const std::shared_ptr& state, + AwsMetadataRequestKind kind, + const std::string& region, + AwsMetadataResult result, + size_t first_request) { + FakeProviderState::Pending pending; + { + std::lock_guard lock(state->mutex); + for (size_t index = first_request; index < state->requests.size(); ++index) { + const auto& item = state->requests[index]; + if (item.request.kind == kind && item.request.region == region) { + pending = item; + break; + } + } + } + if (pending.handle.value == 0) { + return false; + } + if (auto sink = pending.sink.lock()) { + AwsMetadataCompletion completion; + completion.opaque_id = pending.request.opaque_id; + completion.generation = pending.request.generation; + completion.result = std::move(result); + sink->post(std::move(completion)); + return true; + } + return false; +} + template bool wait_until(Predicate predicate) { const auto deadline = std::chrono::steady_clock::now() + 2s; @@ -129,7 +167,7 @@ bool wait_until(Predicate predicate) { if (predicate()) { return true; } - std::this_thread::yield(); + std::this_thread::sleep_for(1ms); } return predicate(); } @@ -164,7 +202,7 @@ const AwsLocalitySnapshotEntry* lookup( } // namespace int main() { - plan(41); + plan(44); auto provider_state = std::make_shared(); ok(install_global_aws_metadata_provider( @@ -287,8 +325,10 @@ int main() { complete_request(provider_state, AwsMetadataRequestKind::rds_region, "eu-west-1", std::move(refresh_failure), 1); const auto failed_refresh_snapshot = manager.snapshot(); + const auto* failed_refresh_entry = lookup(failed_refresh_snapshot, 11, east_one); ok(failed_refreshes_delivered && - lookup(failed_refresh_snapshot, 11, east_one)->status == AwsLocalityMetadataStatus::stale, + failed_refresh_entry != nullptr && + failed_refresh_entry->status == AwsLocalityMetadataStatus::stale, "failed refresh preserves the last successful value within stale TTL"); steady_seconds.store(121); @@ -302,7 +342,11 @@ int main() { hg11_east->multiplier == 1.0, "expired metadata becomes neutral"); - const size_t canceled_before_reload = provider_state->canceled.size(); + size_t canceled_before_reload = 0; + { + std::lock_guard lock(provider_state->mutex); + canceled_before_reload = provider_state->canceled.size(); + } manager.configure({ make_hostgroup(12, 2.0, 4.0, 300, 1800, {"db-new.abcdefghijkl.us-east-1.rds.amazonaws.com"}), @@ -446,6 +490,17 @@ int main() { }), "concurrent completions publish one consistent immutable snapshot"); concurrent_manager.shutdown(); + const size_t startup_request_count = requests_copy(replacement_state).size(); + MySQLAwsLocalityManager startup_manager(manager_config); + startup_manager.set_enabled(true); + startup_manager.configure({ + make_hostgroup(250, 2.0, 4.0, 300, 1800, + {"db-startup.abcdefghijkl.us-east-2.rds.amazonaws.com"}), + }); + ok(wait_for_request_count(replacement_state, startup_request_count + 2), + "configuration starts discovery when the master switch was enabled first"); + startup_manager.shutdown(); + std::mutex completion_hook_mutex; std::condition_variable completion_hook_cv; bool completion_hook_entered = false; @@ -455,28 +510,32 @@ int main() { std::unique_lock lock(completion_hook_mutex); completion_hook_entered = true; completion_hook_cv.notify_all(); - completion_hook_cv.wait(lock, [&] { return release_completion; }); + completion_hook_cv.wait_for(lock, 2s, [&] { return release_completion; }); }; MySQLAwsLocalityManager blocking_manager(blocking_config); const auto blocking_endpoint = "db-block.abcdefghijkl.ap-block-1.rds.amazonaws.com"; + const size_t blocking_request_count = requests_copy(replacement_state).size(); blocking_manager.configure({ make_hostgroup(300, 2.0, 4.0, 300, 1800, {blocking_endpoint}), }); blocking_manager.set_enabled(true); - ok(wait_for_request_count(replacement_state, 105), + ok(wait_for_request_count(replacement_state, blocking_request_count + 2), "callback-shutdown fixture dispatches local and regional requests"); std::atomic callback_returned { false }; std::thread callback_thread([&] { AwsMetadataResult result; result.status = AwsMetadataStatus::ok; result.local = {"ap-block-1", "ap-block-1a", "111122223333"}; - complete_request(replacement_state, AwsMetadataRequestKind::local_location, - "", std::move(result), 2); + complete_request_after(replacement_state, + AwsMetadataRequestKind::local_location, "", std::move(result), + blocking_request_count); callback_returned.store(true); }); { std::unique_lock lock(completion_hook_mutex); - completion_hook_cv.wait(lock, [&] { return completion_hook_entered; }); + if (!completion_hook_cv.wait_for(lock, 2s, [&] { return completion_hook_entered; })) { + BAIL_OUT("completion hook did not enter before shutdown fixture deadline"); + } } std::mutex shutdown_mutex; std::condition_variable shutdown_cv; @@ -506,6 +565,75 @@ int main() { ok(callback_returned.load() && shutdown_done, "callback and shutdown finish without accessing detached manager state"); + AwsLocalityManagerConfig bounded_disable_config = manager_config; + bounded_disable_config.disable_wait_timeout = 50ms; + MySQLAwsLocalityManager bounded_disable_manager(bounded_disable_config); + { + std::lock_guard lock(replacement_state->mutex); + replacement_state->block_requests = true; + replacement_state->request_entered = false; + replacement_state->release_requests = false; + } + bounded_disable_manager.configure({ + make_hostgroup(400, 2.0, 4.0, 300, 1800, + {"db-disable.abcdefghijkl.us-east-2.rds.amazonaws.com"}), + }); + bounded_disable_manager.set_enabled(true); + { + std::unique_lock lock(replacement_state->mutex); + if (!replacement_state->cv.wait_for(lock, 2s, [&] { + return replacement_state->request_entered; + })) { + BAIL_OUT("provider request did not enter bounded-disable fixture"); + } + } + std::mutex disable_mutex; + std::condition_variable disable_cv; + bool disable_returned = false; + std::thread disable_thread([&] { + bounded_disable_manager.set_enabled(false); + std::lock_guard lock(disable_mutex); + disable_returned = true; + disable_cv.notify_all(); + }); + bool disable_returned_within_bound = false; + { + std::unique_lock lock(disable_mutex); + disable_returned_within_bound = disable_cv.wait_for(lock, 250ms, [&] { + return disable_returned; + }); + } + ok(disable_returned_within_bound, + "disabling locality does not block admin on a stalled provider call"); + { + std::lock_guard lock(replacement_state->mutex); + replacement_state->release_requests = true; + replacement_state->cv.notify_all(); + } + disable_thread.join(); + bounded_disable_manager.shutdown(); + + MySQLAwsLocalityManager invalid_identity_manager(manager_config); + AwsLocalityHostgroupConfig invalid_identity_hostgroup; + invalid_identity_hostgroup.hostgroup_id = 500; + invalid_identity_hostgroup.policy = {true, 2.0, 4.0, 300, 1800}; + AwsEndpointCandidate ipv6_identity; + ipv6_identity.hostgroup_id = 500; + ipv6_identity.hostname = "fe80::1"; + ipv6_identity.port = 3306; + AwsEndpointCandidate path_identity = ipv6_identity; + path_identity.hostname = "bad/name"; + invalid_identity_hostgroup.backends.emplace_back(ipv6_identity, 7); + invalid_identity_hostgroup.backends.emplace_back(path_identity, 9); + invalid_identity_manager.configure({std::move(invalid_identity_hostgroup)}); + const auto invalid_identity_snapshot = invalid_identity_manager.snapshot(); + const auto* ipv6_entry = invalid_identity_snapshot->find(500, "fe80::1", 3306); + const auto* path_entry = invalid_identity_snapshot->find(500, "bad/name", 3306); + ok(invalid_identity_snapshot->entries.size() == 2 && ipv6_entry != nullptr && + path_entry != nullptr && ipv6_entry->configured_weight == 7 && + path_entry->configured_weight == 9, + "rejected DNS spellings retain distinct snapshot identities"); + shutdown_global_aws_metadata_provider(); return exit_status(); diff --git a/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp b/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp index 3e9d7d2103..18df1e6cc1 100644 --- a/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp @@ -188,7 +188,7 @@ const AwsMetadataEndpoint* find_endpoint( } // namespace int main() { - plan(31); + plan(30); auto backend = std::make_shared(); AwsSdkMetadataProvider provider(backend, AwsMetadataProviderConfig {2, 3}); @@ -210,10 +210,10 @@ int main() { ok(handles[3].value == 0 && sink->wait_for(1), "bounded queue rejects excess work immediately"); auto completions = sink->snapshot(); - ok(completions[0].opaque_id == 4 && completions[0].generation == 17 && + ok(completions.size() >= 1 && completions[0].opaque_id == 4 && completions[0].generation == 17 && completions[0].result.status == AwsMetadataStatus::throttled, "queue rejection preserves request identity with a fixed category"); - ok(completions[0].result.failure_category == "throttled" && + ok(completions.size() >= 1 && completions[0].result.failure_category == "throttled" && completions[0].result.failure_category.find("FAKE_SECRET") == std::string::npos, "provider never forwards a backend's raw failure text"); provider.cancel(handles[2]); @@ -230,8 +230,6 @@ int main() { } ok(saw_one && saw_two && !saw_three, "queued cancellation suppresses its callback without affecting other work"); - ok(saw_one && saw_two, - "successful asynchronous completions preserve opaque ID and generation"); AwsMetadataRequest expired; expired.opaque_id = 9; @@ -240,7 +238,7 @@ int main() { ok(provider.request(expired, sink).value == 0 && sink->wait_for(4), "already-expired work is rejected without entering the backend"); completions = sink->snapshot(); - ok(completions.back().opaque_id == 9 && + ok(!completions.empty() && completions.back().opaque_id == 9 && completions.back().result.status == AwsMetadataStatus::timeout && completions.back().result.failure_category == "timeout", "deadline rejection is normalized and preserves identity"); diff --git a/test/tap/tests/unit/aws_locality_policy_unit-t.cpp b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp index 41d64359e0..152406e07c 100644 --- a/test/tap/tests/unit/aws_locality_policy_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp @@ -74,7 +74,13 @@ void test_policy_validation() { ok(!policy.valid && error.field == "stale_ttl_seconds", "stale TTL shorter than refresh interval is rejected"); - policy = parse_aws_locality_policy(json::array(), 20, error); + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":4.0,"refresh_interval_seconds":3600})"), + 20, error); + ok(!policy.valid && error.field == "stale_ttl_seconds", + "omitted stale TTL cannot undercut an explicit refresh interval"); + + policy = parse_aws_locality_policy(json::array(), 21, error); ok(!policy.valid && error.field == "locality_awareness", "non-object locality policy is rejected"); } @@ -107,6 +113,10 @@ void test_endpoint_recognition() { ok(!recognize_rds_endpoint(14, "db.proxy-abcdefghijkl.us-east-1.rds.amazonaws.com", 3306).recognized, "RDS Proxy endpoint is excluded from the first version"); + candidate = recognize_rds_endpoint( + 14, "proxy-app1.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); + ok(candidate.recognized && candidate.region == "us-east-1", + "an ordinary RDS identifier beginning with proxy- is not misclassified as RDS Proxy"); ok(!recognize_rds_endpoint(15, "db.internal.example.com", 3306).recognized, "custom CNAME is not inferred as RDS"); ok(!recognize_rds_endpoint(16, "db.us-east-1.rds.amazonaws.com.evil.example", 3306).recognized, @@ -171,7 +181,7 @@ void test_effective_weight() { } // namespace int main() { - plan(32); + plan(34); test_policy_validation(); test_endpoint_recognition(); test_classification(); diff --git a/test/tap/tests/unit/aws_locality_selection_unit-t.cpp b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp index dfb4c03b53..70440d0f79 100644 --- a/test/tap/tests/unit/aws_locality_selection_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp @@ -3,6 +3,7 @@ #include "test_init.h" #include "Aws_Locality_Manager.h" +#include "GTID_Server_Data.h" #include "MySQL_Data_Stream.h" #include "MySQL_HostGroups_Manager.h" #include "MySQL_Logger.hpp" @@ -10,6 +11,7 @@ #include #include +#include #include #include #include @@ -19,6 +21,26 @@ using namespace std::chrono_literals; +namespace { +thread_local bool track_hot_path_allocations = false; +thread_local size_t hot_path_allocations = 0; +} + +void* operator new(std::size_t size) { + if (track_hot_path_allocations) ++hot_path_allocations; + if (void* memory = std::malloc(size)) return memory; + throw std::bad_alloc(); +} + +void* operator new[](std::size_t size) { + return ::operator new(size); +} + +void operator delete(void* memory) noexcept { std::free(memory); } +void operator delete[](void* memory) noexcept { std::free(memory); } +void operator delete(void* memory, std::size_t) noexcept { std::free(memory); } +void operator delete[](void* memory, std::size_t) noexcept { std::free(memory); } + extern MySQL_HostGroups_Manager* MyHGM; extern MySQL_Threads_Handler* GloMTH; extern MySQL_Logger* GloMyLogger; @@ -192,18 +214,22 @@ MySQL_Connection* make_connection(MySrvC* server, int fd) { } // namespace int main() { - plan(20); + plan(28); ok(aws_locality_saturating_add( std::numeric_limits::max() - 2, 5) == std::numeric_limits::max(), "locality weight sums saturate instead of overflowing"); const uint64_t lottery_weights[] = {40, 40, 30}; - ok(aws_locality_weighted_index(lottery_weights, 3, 0) == 0 && - aws_locality_weighted_index(lottery_weights, 3, 39) == 0 && - aws_locality_weighted_index(lottery_weights, 3, 40) == 1 && - aws_locality_weighted_index(lottery_weights, 3, 79) == 1 && - aws_locality_weighted_index(lottery_weights, 3, 80) == 2, - "shared locality lottery uses exact cumulative weight boundaries"); + ok(aws_locality_weighted_index(lottery_weights, 3, 0) == 0, + "locality lottery selects the first candidate at its lower boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 39) == 0, + "locality lottery selects the first candidate at its upper boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 40) == 1, + "locality lottery selects the second candidate at its lower boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 79) == 1, + "locality lottery selects the second candidate at its upper boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 80) == 2, + "locality lottery selects the final candidate at its lower boundary"); const uint64_t zero_weights[] = {0, 0}; ok(aws_locality_weighted_index(zero_weights, 2, 17) == 2, "shared locality lottery rejects an all-zero candidate set"); @@ -293,15 +319,36 @@ int main() { local_share, regional_share, remote_share); ok(local->weight == 10 && regional->weight == 20 && remote->weight == 30, "global locality selection never mutates configured server weights"); + local->weight = 0; + regional->weight = 0; + remote->weight = 0; + ok(hostgroup->get_random_MySrvC(nullptr, 0, -1, session.session) != nullptr, + "global selection retains an eligible fallback when all configured weights are zero"); + local->weight = 10; + regional->weight = 20; + remote->weight = 30; MySQL_Connection* local_connection = make_connection(local, 100); std::vector remote_connections; for (int i = 0; i < 5; ++i) { remote_connections.push_back(make_connection(remote, 200 + i)); } + for (int i = 5; i < 33; ++i) { + remote_connections.push_back(make_connection(remote, 200 + i)); + } worker.push_MyConn_local(local_connection); for (auto* connection : remote_connections) worker.push_MyConn_local(connection); + hot_path_allocations = 0; + track_hot_path_allocations = true; + MySQL_Connection* allocation_probe = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + track_hot_path_allocations = false; + ok(allocation_probe != nullptr && hot_path_allocations == 0, + "locality selection performs no heap allocation with more than 32 cached connections"); + if (allocation_probe != nullptr) worker.push_MyConn_local(allocation_probe); + int local_parent = 0; int remote_parent = 0; for (int i = 0; i < 6000; ++i) { @@ -318,6 +365,30 @@ int main() { local_parent + remote_parent == 6000, "local cache chooses 40:30 weighted parents, not the 1:5 connection count (%.3f local)", local_parent_share); + local->weight = 0; + remote->weight = 0; + MySQL_Connection* zero_weight_selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(zero_weight_selected != nullptr, + "local cache reuses an eligible connection when all parent weights are zero"); + if (zero_weight_selected != nullptr) worker.push_MyConn_local(zero_weight_selected); + local->weight = 10; + remote->weight = 30; + + GTID_Server_Data local_gtid(nullptr, const_cast(kLocal), 0, 3306); + local_gtid.add_gtid_from_ok("aaaaaaaa-0000-1111-2222-aaaaaaaaaaaa:42"); + MyHGM->gtid_map.emplace(std::string(kLocal) + ":3306", &local_gtid); + local->aws_aurora_current_lag_us = 5000; + char gtid_uuid[] = "aaaaaaaa000011112222aaaaaaaaaaaa"; + MySQL_Connection* gtid_selected = worker.get_MyConn_local( + kHostgroup, session.session, gtid_uuid, 42, 1, + MySQLBackendAuthType::PASSWORD); + ok(gtid_selected != nullptr && gtid_selected->parent == local, + "GTID-qualified local reuse preserves the legacy max-lag exemption"); + if (gtid_selected != nullptr) worker.push_MyConn_local(gtid_selected); + local->aws_aurora_current_lag_us = 0; + MyHGM->gtid_map.erase(std::string(kLocal) + ":3306"); remote->aws_aurora_current_lag_us = 5000; MySQL_Connection* selected = worker.get_MyConn_local( diff --git a/test/tap/tests/unit/aws_locality_stats_unit-t.cpp b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp index 84cace6e53..b842d3c3d4 100644 --- a/test/tap/tests/unit/aws_locality_stats_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp @@ -8,7 +8,7 @@ #include #include #include -#include +#include #include #ifndef PROXYSQL_AWS_PLUGIN_PATH @@ -95,7 +95,10 @@ AwsLocalitySnapshotEntry diagnostic_row( } // namespace int main() { - plan(22); + plan(23); + if (test_globals_init() != 0) { + BAIL_OUT("test global initialization failed"); + } SQLite3DB statsdb; statsdb.open((char*)":memory:", @@ -176,8 +179,36 @@ int main() { "'db''quoted.abcdefghijkl.us-east-1.rds.amazonaws.com'") == 1, "projection safely quotes endpoint text"); + std::vector concurrent_rows_a; + std::vector concurrent_rows_b; + for (uint32_t index = 0; index < 200; ++index) { + concurrent_rows_a.push_back(diagnostic_row(1000 + index, + AwsLocalityMetadataStatus::fresh, 2.0, 10)); + concurrent_rows_b.push_back(diagnostic_row(2000 + index, + AwsLocalityMetadataStatus::stale, 3.0, 10)); + } + std::atomic start_concurrent_projection { false }; + std::atomic successful_projections { 0 }; + auto project_repeatedly = [&](const std::vector& projection) { + while (!start_concurrent_projection.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (unsigned int iteration = 0; iteration < 10; ++iteration) { + if (MySQL_HostGroups_Manager::project_aws_locality_stats(&statsdb, projection)) { + successful_projections.fetch_add(1, std::memory_order_relaxed); + } + } + }; + std::thread projection_a(project_repeatedly, std::cref(concurrent_rows_a)); + std::thread projection_b(project_repeatedly, std::cref(concurrent_rows_b)); + start_concurrent_projection.store(true, std::memory_order_release); + projection_a.join(); + projection_b.join(); + ok(successful_projections.load() == 20 && + statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 200, + "concurrent runtime-view refreshes serialize complete replacement transactions"); + GloVars.prometheus_registry = std::make_shared(); - unlink("file:mem_mydb?mode=memory"); { MySQL_HostGroups_Manager hostgroups; MyHGM = &hostgroups; @@ -216,7 +247,6 @@ int main() { "repeated generation queries remain network-free"); MyHGM = nullptr; } - unlink("file:mem_mydb?mode=memory"); shutdown_global_aws_metadata_provider(); GloVars.prometheus_registry.reset(); @@ -236,5 +266,6 @@ int main() { statsdb.execute("PRAGMA query_only = OFF"); proxysql_stop_configured_plugins(manager, error); + test_globals_cleanup(); return exit_status(); } From e8c542336ef147ee0da1b12f2704236d4e877ed9 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Fri, 14 Aug 2026 09:46:16 +0000 Subject: [PATCH 14/17] fix(aws): constrain IMDS curl transport --- plugins/aws/src/aws_locality_provider.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/plugins/aws/src/aws_locality_provider.cpp b/plugins/aws/src/aws_locality_provider.cpp index 74b2de42b2..7785a7e3e2 100644 --- a/plugins/aws/src/aws_locality_provider.cpp +++ b/plugins/aws/src/aws_locality_provider.cpp @@ -616,6 +616,10 @@ AwsImdsResponse imds_request( curl_easy_setopt(handle, CURLOPT_NOBODY, 0L); curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 0L); curl_easy_setopt(handle, CURLOPT_NOPROXY, "*"); + // IMDS is an HTTP-only link-local service. Restrict the handle to that + // protocol and retain a secure TLS floor if the transport URL ever changes. + curl_easy_setopt(handle, CURLOPT_PROTOCOLS_STR, "http"); + curl_easy_setopt(handle, CURLOPT_SSLVERSION, CURL_SSLVERSION_TLSv1_2); curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, std::min(timeout, 500L)); curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, timeout); curl_easy_setopt(handle, CURLOPT_NOSIGNAL, 1L); From 7b15362deb2f9995178feffc09169ecfe2a9072b Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Fri, 14 Aug 2026 12:40:23 +0000 Subject: [PATCH 15/17] fix(aws): keep IMDS transport HTTP-only --- plugins/aws/src/aws_locality_provider.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/plugins/aws/src/aws_locality_provider.cpp b/plugins/aws/src/aws_locality_provider.cpp index 7785a7e3e2..158ddb84a5 100644 --- a/plugins/aws/src/aws_locality_provider.cpp +++ b/plugins/aws/src/aws_locality_provider.cpp @@ -616,10 +616,9 @@ AwsImdsResponse imds_request( curl_easy_setopt(handle, CURLOPT_NOBODY, 0L); curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 0L); curl_easy_setopt(handle, CURLOPT_NOPROXY, "*"); - // IMDS is an HTTP-only link-local service. Restrict the handle to that - // protocol and retain a secure TLS floor if the transport URL ever changes. + // IMDS is an HTTP-only link-local service. Restrict this handle to that + // protocol; TLS options are intentionally inapplicable to this transport. curl_easy_setopt(handle, CURLOPT_PROTOCOLS_STR, "http"); - curl_easy_setopt(handle, CURLOPT_SSLVERSION, CURL_SSLVERSION_TLSv1_2); curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, std::min(timeout, 500L)); curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, timeout); curl_easy_setopt(handle, CURLOPT_NOSIGNAL, 1L); From d0c247600cf43dff8f0776986d26cf66f2c2b535 Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Fri, 14 Aug 2026 18:07:49 +0000 Subject: [PATCH 16/17] fix(aws): resolve locality Sonar findings --- lib/Aws_Locality_Manager.cpp | 8 +++++--- plugins/aws/src/aws_locality_provider.cpp | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp index bbc469d3d9..9253fed6f6 100644 --- a/lib/Aws_Locality_Manager.cpp +++ b/lib/Aws_Locality_Manager.cpp @@ -906,9 +906,11 @@ class MySQLAwsLocalityManager::Impl { } } else { auto region = region_in_flight_.find(request.region); - if (region != region_in_flight_.end() && region->second != 0 && - --region->second == 0) { - region_in_flight_.erase(region); + if (region != region_in_flight_.end() && region->second != 0) { + --region->second; + if (region->second == 0) { + region_in_flight_.erase(region); + } } } if (completion.generation != request.generation || diff --git a/plugins/aws/src/aws_locality_provider.cpp b/plugins/aws/src/aws_locality_provider.cpp index 158ddb84a5..24649d7729 100644 --- a/plugins/aws/src/aws_locality_provider.cpp +++ b/plugins/aws/src/aws_locality_provider.cpp @@ -603,7 +603,7 @@ AwsImdsResponse imds_request( const AwsLocalityCancelPredicate& cancelled) { AwsImdsResponse response; if (cancelled() || std::chrono::steady_clock::now() >= deadline) return response; - CURL* handle = curl_easy_init(); + CURL* handle = curl_easy_init(); // NOSONAR(cpp:S4423): IMDSv2 is HTTP-only; this handle is restricted to HTTP below. if (handle == nullptr) return response; CurlResponseContext write_context {&response.body, maximum}; CurlProgressContext progress_context {deadline, &cancelled}; From 5791c85d55c76b086d92f6e89564029c41e0b2bc Mon Sep 17 00:00:00 2001 From: Rene Cannao Date: Sat, 15 Aug 2026 13:51:19 +0000 Subject: [PATCH 17/17] refactor(aws): move locality provider out of tree --- doc/PLUGIN_API.md | 34 +- doc/aws-locality-awareness.md | 270 ++---- .../2026-08-13-aws-locality-awareness.md | 481 ----------- ...026-08-13-aws-locality-awareness-design.md | 679 --------------- lib/Aws_Locality_Manager.cpp | 12 + lib/ProxySQL_Admin.cpp | 10 +- plugins/aws/src/aws_locality_provider.cpp | 789 ------------------ plugins/aws/src/aws_locality_provider.h | 218 ----- test/tap/tests/unit/Makefile | 21 +- .../unit/aws_locality_manager_unit-t.cpp | 47 +- .../tests/unit/aws_locality_plugin_unit-t.cpp | 409 --------- .../unit/aws_locality_selection_unit-t.cpp | 39 +- .../tests/unit/aws_locality_stats_unit-t.cpp | 50 +- 13 files changed, 223 insertions(+), 2836 deletions(-) delete mode 100644 docs/superpowers/plans/2026-08-13-aws-locality-awareness.md delete mode 100644 docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md delete mode 100644 plugins/aws/src/aws_locality_provider.cpp delete mode 100644 plugins/aws/src/aws_locality_provider.h delete mode 100644 test/tap/tests/unit/aws_locality_plugin_unit-t.cpp diff --git a/doc/PLUGIN_API.md b/doc/PLUGIN_API.md index dc6ccbf252..80ec1cc79f 100644 --- a/doc/PLUGIN_API.md +++ b/doc/PLUGIN_API.md @@ -103,7 +103,7 @@ All types are defined in `include/ProxySQL_Plugin.h`: ```cpp struct ProxySQL_PluginDescriptor { const char *name; // Human-readable plugin name - uint32_t abi_version; // PROXYSQL_PLUGIN_ABI_VERSION (1 through 5) + uint32_t abi_version; // PROXYSQL_PLUGIN_ABI_VERSION (1 through 8) proxysql_plugin_init_cb init; // bool (*)(ProxySQL_PluginServices *) proxysql_plugin_start_cb start; // bool (*)() proxysql_plugin_stop_cb stop; // bool (*)() @@ -115,7 +115,7 @@ struct ProxySQL_PluginDescriptor { | Field | Type | Description | |--------------------|---------------|-----------------------------------------------------------| | `name` | `const char*` | Plugin identifier, used in logging. | -| `abi_version` | `uint32_t` | Set from `PROXYSQL_PLUGIN_ABI_VERSION`. Value `1` = pre-chassis descriptor (six fields). Value `2` = adds `register_schemas` (four-phase lifecycle). Value `3` = `ProxySQL_PluginServices` adds `register_runtime_view`; ABI 4 adds `db_kind` to that view; ABI 5 adds the AWS IAM provider callbacks. A v3/v3.1 core rejects `abi_version > 1`; the current PROXYSQL40 core accepts `[1, 5]`. | +| `abi_version` | `uint32_t` | Set from `PROXYSQL_PLUGIN_ABI_VERSION`. Value `1` is the pre-chassis six-field descriptor; ABI 2 adds `register_schemas`; ABI 3 adds `register_runtime_view`; ABI 4 adds the view's `db_kind`; ABI 5 adds IAM provider install/limits; ABI 6 adds metadata-provider install; ABI 7 adds the MySQL locality projection callback; ABI 8 adds IAM provider rollback. A v3/v3.1 core rejects `abi_version > 1`; the current PROXYSQL40 core accepts `[1, 8]`. | | `init` | callback | Phase D — called with live services; register tables and commands here (or finish context setup if `register_schemas` already did it). | | `start` | callback | Phase E — start threads, open sockets, load config. | | `stop` | callback | Called on shutdown. Pairs with `init`, not `start`: if `init` returned true and `start` later failed, `stop` is still called so the plugin can release resources it allocated in `init`. | @@ -128,12 +128,12 @@ Return `true` on success, `false` on failure. A `false` return from #### ABI version -`include/ProxySQL_Plugin.h` exposes `PROXYSQL_PLUGIN_ABI_VERSION` (5 under +`include/ProxySQL_Plugin.h` exposes `PROXYSQL_PLUGIN_ABI_VERSION` (8 under PROXYSQL40, undefined in pre-chassis builds — the descriptor is then a legacy six-field struct with `abi_version = 1`). Plugins MUST assign `abi_version` from this macro rather than hard-coding a literal; the core's loader uses it to detect layout skew and reject plugins built -for an unsupported ABI. ABIs 3 through 5 keep the descriptor layout +for an unsupported ABI. ABIs 3 through 8 keep the descriptor layout identical to ABI 2; each adds only tail fields to service or view structs. Plugins compiled against an older ABI therefore still load on the current core, where the trailing fields remain invisible to them. See @@ -172,6 +172,12 @@ struct ProxySQL_PluginServices { // ABI 5 tail extensions, live only while init() is running: proxysql_plugin_install_aws_iam_token_source_cb install_aws_iam_token_source; proxysql_plugin_get_aws_iam_limits_cb get_aws_iam_limits; + // ABI 6 tail extension, live only while init() is running: + proxysql_plugin_install_aws_metadata_provider_cb install_aws_metadata_provider; + // ABI 7 tail extension, live during register_schemas() and init(): + proxysql_plugin_refresh_mysql_aws_locality_stats_cb refresh_mysql_aws_locality_stats; + // ABI 8 tail extension, live only while init() is running: + proxysql_plugin_uninstall_aws_iam_token_source_cb uninstall_aws_iam_token_source; }; ``` @@ -183,6 +189,24 @@ module. A plugin must not call this callback outside `init()` or retain the service pointer after initialization. Without an installed provider, IAM authentication remains available as a policy but fails closed. +`install_aws_metadata_provider` installs an optional external asynchronous +metadata provider for MySQL locality selection. Core owns the retained lease +registry: new leases are rejected during shutdown, active manager/callback +leases drain, the provider's `shutdown()` and destroy callback run, and the +extra module handle closes last. Without an installed provider, locality uses +configured weights and reports the fixed `provider_unavailable` category. + +`refresh_mysql_aws_locality_stats` projects the current MySQL-owned immutable +locality snapshot into a schema supplied by the calling plugin. It performs no +provider or network I/O and is live during `register_schemas()` so an external +provider can register its stats table and runtime-view callback together. +Public core deliberately registers no locality table on its own. + +`uninstall_aws_iam_token_source` lets the same plugin roll back its IAM source +when a later initialization step fails. It rejects a null or different source, +stops new leases, drains active leases, destroys the source, and releases the +retained module handle before returning. + ### Service Callbacks #### `register_table` @@ -553,7 +577,7 @@ void register_stats_table(ProxySQL_PluginServices& services, - **No dependency resolution**: Plugins are loaded in the order listed in `proxysql.cnf`. If one plugin depends on another, the dependency must be listed first. -- **ABI version range**: The current core accepts `abi_version` values in `[1, 5]`. Newly built plugins should set `abi_version = PROXYSQL_PLUGIN_ABI_VERSION`. +- **ABI version range**: The current core accepts `abi_version` values in `[1, 8]`. Newly built plugins should set `abi_version = PROXYSQL_PLUGIN_ABI_VERSION`. - **Compiler coupling**: Plugins must match the ProxySQL core's C++ compiler and standard library due to `std::string` in `ProxySQL_PluginCommandResult`. diff --git a/doc/aws-locality-awareness.md b/doc/aws-locality-awareness.md index 7da4f7296e..fa16f1fc50 100644 --- a/doc/aws-locality-awareness.md +++ b/doc/aws-locality-awareness.md @@ -1,83 +1,37 @@ # AWS locality-aware MySQL backend selection -ProxySQL 4.0 can prefer Amazon RDS and Aurora MySQL backends in the same AWS -Region or Availability Zone (AZ) as the ProxySQL process. The feature changes -only the temporary weights used by a server-selection attempt. It never -changes `mysql_servers.weight`, `runtime_mysql_servers.weight`, saved -configuration, or ProxySQL Cluster checksums. - -The MySQL module owns the switch and hostgroup policy. The optional general -AWS plugin provides local-instance and RDS metadata asynchronously. If the -plugin is not loaded, metadata is unavailable, or metadata expires, ProxySQL -continues using the configured weights. - -## Load the AWS plugin - -The AWS plugin is built and packaged by `PROXYSQL40=1 make -j`, with the -vendored AWS SDK linked statically into `ProxySQL_Aws_Plugin.so`. It is not -loaded automatically. Add its installed path to the `plugins` list in the -ProxySQL configuration and restart ProxySQL: - -```ini -plugins = ( - "/usr/lib/proxysql/ProxySQL_Aws_Plugin.so" -) -``` - -The plugin also provides [AWS IAM database authentication](aws_iam_database_authentication.md). -Both capabilities share one AWS SDK runtime and the standard AWS credential -provider chain. - -## Grant discovery permission - -Give the ProxySQL workload read-only access to the RDS discovery APIs: - -```json -{ - "Version": "2012-10-17", - "Statement": [ - { - "Effect": "Allow", - "Action": [ - "rds:DescribeDBInstances", - "rds:DescribeDBClusters", - "rds:DescribeDBClusterEndpoints" - ], - "Resource": "*" - } - ] -} -``` - -No long-lived access keys are required or recommended. On EC2, attach an -instance profile to the instance running ProxySQL. On EKS, use IRSA or EKS Pod -Identity for the ProxySQL pod. ECS task roles and other standard AWS SDK -credential sources also work. ProxySQL adds no access-key, secret-key, or -role-ARN variables. +ProxySQL 4.0 can apply temporary locality multipliers while selecting eligible +Amazon RDS and Aurora MySQL backends. The MySQL module owns the configuration, +policy validation, immutable selection snapshot, and effective-weight +calculation. Locality never changes `mysql_servers.weight`, +`runtime_mysql_servers.weight`, saved configuration, or ProxySQL Cluster +checksums. -IMDSv2 local-location discovery itself needs no IAM API permission. RDS API -discovery uses the workload identity above. +Locality metadata is supplied asynchronously by an optional compatible +external provider. Without a provider, or when metadata is unavailable or too +old, selection stays neutral and uses the configured server weights. -## Configure a hostgroup +## Hostgroup policy -First add a locality policy to the existing +Add `aws.locality_awareness` to the existing `mysql_hostgroup_attributes.hostgroup_settings` JSON. Both multipliers are required: ```sql INSERT INTO mysql_hostgroup_attributes(hostgroup_id, hostgroup_settings) -VALUES (10, json_object( - 'aws', json_object( - 'locality_awareness', json_object( - 'same_region_multiplier', 2.0, - 'same_az_multiplier', 4.0, - 'refresh_interval_seconds', 300, - 'stale_ttl_seconds', 1800)))) -ON CONFLICT(hostgroup_id) DO UPDATE SET - hostgroup_settings = json_set( - COALESCE(mysql_hostgroup_attributes.hostgroup_settings, '{}'), - '$.aws.locality_awareness', - json_extract(excluded.hostgroup_settings, '$.aws.locality_awareness')); +VALUES ( + 10, + '{ + "aws": { + "locality_awareness": { + "same_region_multiplier": 2.0, + "same_az_multiplier": 4.0, + "refresh_interval_seconds": 300, + "stale_ttl_seconds": 1800 + } + } + }' +); LOAD MYSQL SERVERS TO RUNTIME; SAVE MYSQL SERVERS TO DISK; @@ -95,26 +49,13 @@ stale_ttl_seconds default: 1800 refresh_interval_seconds <= stale_ttl_seconds <= 604800 ``` -The multipliers are floating-point JSON numbers. For each selection attempt, -ProxySQL calculates: - -```text -remote or unknown configured_weight -same Region, different AZ int(configured_weight * same_region_multiplier) -same AZ int(configured_weight * same_az_multiplier) -``` - -Conversion to an integer truncates toward zero. The tiers are not cumulative: -a same-AZ backend gets only the AZ multiplier, not the Region multiplier times -the AZ multiplier. Weight zero stays zero. +Multipliers are JSON numbers. An invalid `aws.locality_awareness` object +disables locality bias for that hostgroup when servers are loaded. Diagnostics +identify the rejected field and hostgroup without logging the supplied value. -An invalid `aws.locality_awareness` object disables locality bias for that -hostgroup when servers are loaded. The diagnostic identifies only the -hostgroup and rejected field; it does not print the JSON value. +## Master switch -## Enable the global switch - -The one process-wide control is a MySQL module variable. It defaults off: +The process-wide MySQL variable defaults to `false`: ```sql SET mysql-aws_locality_awareness = true; @@ -122,131 +63,52 @@ LOAD MYSQL VARIABLES TO RUNTIME; SAVE MYSQL VARIABLES TO DISK; ``` -There are deliberately no ProxySQL Region, AZ, or AWS-account variables. -Every ProxySQL process discovers its own location, so ProxySQL Cluster cannot -copy one node's location to another node in a different Region or AZ. - -Disabling the variable immediately restores ordinary configured-weight -selection, cancels or supersedes outstanding locality requests, and stops new -refresh scheduling. It does not remove or rebalance existing connections. - -## How local location is determined - -ProxySQL tries the EC2 IMDSv2 instance identity document first. If IMDSv2 is -unavailable, the AWS plugin uses these process environment values: - -1. Region: `AWS_REGION`, then `AWS_DEFAULT_REGION`; -2. AZ: `AWS_AVAILABILITY_ZONE`; -3. account assertion: `AWS_ACCOUNT_ID`. - -Region is the only required fallback field. AZ without Region is ignored. A -same-AZ preference also requires the same confirmed account ID on both sides, -because identical AZ names in different accounts may represent different -physical zones. Without an account assertion, same-Region preference remains -available but same-AZ preference does not. - -In Kubernetes, inject the node Region and AZ into each ProxySQL pod from the -node's topology labels and inject the account assertion through deployment -configuration. These are per-pod environment values, not clustered ProxySQL -variables. +Disabling the variable immediately restores configured-weight selection, +cancels or supersedes outstanding locality requests, and stops new refresh +scheduling. It does not change existing backend connections or server rows. -## Supported backend endpoints +## Selection contract -The first release supports official RDS and Aurora endpoints returned by the -AWS APIs: +For each selection attempt, after the normal health, capacity, lag, GTID, +backoff, and session-compatibility checks, ProxySQL calculates: -- RDS DB instance and Aurora DB instance endpoints can receive Region and AZ - preference; -- Aurora or Multi-AZ cluster writer endpoints receive Region preference only; -- Aurora reader and custom endpoints receive Region preference only. - -Cluster, reader, and custom endpoints can route to several AZs, so they never -receive the same-AZ multiplier. ProxySQL normalizes the configured hostname -and requires an exact match in a paginated RDS API response. Where AWS returns -a port, it must also match. - -Custom CNAMEs, arbitrary EC2 MySQL hosts, RDS Proxy endpoints, DNS aliases, -and AWS-looking names absent from the authoritative response remain neutral. -ProxySQL does not resolve CNAMEs, scan all Regions, assume roles into other -accounts, or change hostgroup membership. - -## Selection and failure behavior - -Locality is applied only after all normal eligibility checks, including -server state, capacity, latency, GTID, replication lag, session compatibility, -and existing backoff rules. It cannot make an unhealthy or incompatible -backend eligible. - -The global Hostgroup Manager path and the thread-local idle-connection path -use the same temporary effective server weights. The local cache groups -eligible connections by parent server first, so a server does not gain more -probability merely because it has more idle connections. - -Metadata states behave as follows: - -- `pending`: discovery has not completed; configured weight is used; -- `fresh`: the active Region or AZ multiplier is used; -- `stale`: the last success remains active through the bounded stale TTL; -- `expired`: configured weight is used; -- `error`: no usable value exists, so configured weight is used; -- `disabled`: the master switch is off and configured weight is used. - -A failed refresh does not immediately discard a prior success. Once its -`stale_ttl_seconds` expires, locality becomes neutral. Missing credentials, -access denial, throttling, timeout, missing plugin, unsupported endpoint, and -IMDS failure are all fail-neutral for database traffic. - -AWS, IMDS, credential-provider, DNS, and network work runs on bounded plugin -workers. No server-selection path calls the plugin or performs network I/O. - -## Inspect locality decisions - -When—and only when—the AWS plugin is loaded during startup, it registers the -read-only `stats_mysql_aws_locality` table. Querying it materializes one -current immutable manager snapshot into the stats database: - -```sql -SELECT * -FROM stats_mysql_aws_locality -ORDER BY hostgroup_id, hostname, port; +```text +remote or unknown configured_weight +same Region, different AZ int(configured_weight * same_region_multiplier) +same AZ int(configured_weight * same_az_multiplier) ``` -The columns are: +Conversion to an integer truncates toward zero. The tiers are not cumulative, +and configured weight zero stays zero. Global hostgroup selection and +thread-local idle-connection reuse use the same immutable snapshot and never +perform provider or network work on the selection path. -```text -hostgroup_id hostname -port endpoint_type -configured_weight effective_weight -local_region local_az -backend_region backend_az -account_match locality -active_multiplier metadata_status -last_success_timestamp last_attempt_timestamp -last_error_category -``` +ProxySQL recognizes eligible RDS and Aurora endpoint shapes to form neutral +metadata requests. A compatible provider is responsible for authoritative +endpoint discovery and normalized Region, Availability Zone, and account +metadata. Custom CNAMEs, arbitrary MySQL hosts, proxy endpoints, malformed +names, and endpoints the provider cannot confirm remain neutral. + +## Provider absence and failures -`endpoint_type` is `instance`, `cluster`, `reader`, `custom`, or `unknown`. -`account_match` is `same`, `different`, or `unknown`; account IDs themselves -are never exposed. `locality` is `same_az`, `same_region`, `remote`, or -`unknown`. Timestamps are Unix epoch seconds, or zero if no corresponding -event has occurred. +The public provider interface uses retained leases so shutdown rejects new +work, drains active callbacks, joins the locality manager worker, and only then +permits provider destruction and module unload. Provider absence is exposed as +the fixed `provider_unavailable` category. -The table query never starts or waits for metadata discovery. It reports the -latest published state and replaces the previous rows in one transaction, so -generations cannot mix. The table is not persisted, loaded to runtime, -clustered, or included in a checksum. If no hostgroup has a valid policy, it -is empty. If the master switch is off, cached location text remains visible, -but every row reports `disabled`, multiplier `1.0`, and the configured weight -as its effective weight. +Metadata states are `pending`, `fresh`, `stale`, `expired`, `error`, and +`disabled`. Only `fresh` and unexpired `stale` values activate a multiplier. +All other states use multiplier `1.0`, so effective weights equal configured +weights. A failed refresh can retain the last successful value through the +bounded stale TTL; it becomes neutral after expiry. -Without the AWS plugin, the table is not registered and a query returns the -normal `no such table` error. ProxySQL does not currently hot-unload configured -plugins; changing the plugin list requires a restart. +An external provider may register a read-only runtime diagnostics table using +the plugin table/view services and the MySQL-owned projection callback. Public +core does not register `stats_mysql_aws_locality`; without a provider that +registers it, the table does not exist. -## Roll back +## Rollback Set `mysql-aws_locality_awareness` to `false` and load MySQL variables to -runtime. This restores the existing selection path immediately without -changing server rows. To remove a policy as well, remove the -`aws.locality_awareness` object from that hostgroup's `hostgroup_settings` and -load MySQL servers to runtime. +runtime. To remove a policy, delete the `aws.locality_awareness` object from +that hostgroup's `hostgroup_settings`, then load MySQL servers to runtime. diff --git a/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md b/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md deleted file mode 100644 index f3744c16b4..0000000000 --- a/docs/superpowers/plans/2026-08-13-aws-locality-awareness.md +++ /dev/null @@ -1,481 +0,0 @@ -# AWS Locality-Aware Backend Selection Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Prefer eligible RDS/Aurora backends in the ProxySQL process's AWS Region or Availability Zone using temporary selection weights, without changing configured/runtime server weights. - -**Architecture:** MySQL core parses policy, owns refresh/cache state, publishes immutable selection snapshots, and projects diagnostics. The general `aws` plugin owns the shared AWS SDK runtime, IMDSv2 and RDS calls, and an asynchronous metadata provider installed through the plugin ABI. Both global and thread-local connection selection retain one core snapshot and use one pure effective-weight helper; no hot path calls the plugin or performs network work. - -**Tech Stack:** C++17, ProxySQL plugin ABI, AWS SDK for C++ 1.11.869 (`core` and `rds`), vendored libcurl, nlohmann JSON, SQLite stats runtime views, TAP unit/integration tests, ASan and TSan. - -## Global Constraints - -- Build tier is selected only with `PROXYSQL40=1`; there is no AWS-specific build flag. -- Invoke builds with `make -j`; never hard-code `-j` inside a Makefile recipe. -- The AWS SDK remains statically linked only into `plugins/aws/ProxySQL_Aws_Plugin.so`; the ProxySQL daemon and `libproxysql.a` remain free of AWS SDK symbols and DSOs. -- The AWS SDK release stays pinned to the existing vendored 1.11.869 LFS tarball and uses ProxySQL's vendored dependencies. -- Locality never changes `mysql_servers.weight`, `runtime_mysql_servers.weight`, saved configuration, or ProxySQL Cluster checksums. -- `mysql-aws_locality_awareness` is the only global control; Region, AZ, and account identity are node-local discoveries, never synchronized variables. -- Policy lives at `mysql_hostgroup_attributes.hostgroup_settings.aws.locality_awareness`. -- Multipliers are finite JSON numbers satisfying `1.0 <= same_region_multiplier <= same_az_multiplier <= 10.0`. -- Default refresh is 300 seconds; default stale TTL is 1800 seconds; accepted bounds are `30 <= refresh <= 86400` and `refresh <= stale_ttl <= 604800`. -- Every failure is fail-neutral and all logged/provider failure data is fixed-category and redacted. -- `stats_mysql_aws_locality` exists only when the AWS plugin loads successfully and querying it never performs network discovery. - ---- - -## File Structure - -- `include/Aws_Locality_Types.h`: SDK-free provider request/result, policy, endpoint identity, classification, diagnostics, and interfaces. -- `include/Aws_Locality_Manager.h`: provider lease/registry plus the core manager public API. -- `lib/Aws_Locality_Manager.cpp`: parsing helpers, endpoint recognition, classification, arithmetic, registry lifetime, scheduler/cache, immutable publication, and diagnostic snapshots. -- `include/ProxySQL_Plugin.h`, `lib/ProxySQL_PluginManager.cpp`: ABI-6 provider installation, core snapshot projection callback, and service wiring. -- `include/MySQL_HostGroups_Manager.h`, `lib/MySQL_HostGroups_Manager.cpp`: manager ownership, hostgroup policy storage, reload registration, and diagnostics projection. -- `include/MySQL_Thread.h`, `lib/MySQL_Thread.cpp`, `lib/Admin_FlushVariables.cpp`: `mysql-aws_locality_awareness` lifecycle and manager enable/disable notification. -- `lib/MyHGC.cpp`: global Hostgroup Manager effective weighting. -- `lib/MySQL_Thread.cpp`: locality-aware parent-server selection for the thread-local connection cache. -- `plugins/aws/src/aws_plugin.cpp`: shared SDK runtime and capability installation. -- `plugins/aws/src/aws_locality_provider.h`, `plugins/aws/src/aws_locality_provider.cpp`: bounded provider, IMDSv2/environment discovery, paginated RDS discovery, normalization, and cancellation. -- `plugins/aws/Makefile`, `lib/Makefile`: new compilation units and dependencies. -- `test/tap/tests/unit/aws_locality_policy_unit-t.cpp`: policy, DNS recognition, classification, and arithmetic. -- `test/tap/tests/unit/aws_locality_manager_unit-t.cpp`: cache, refresh, generation, stale, cancellation, and concurrency. -- `test/tap/tests/unit/aws_locality_selection_unit-t.cpp`: global and local-cache selection behavior. -- `test/tap/tests/unit/aws_locality_plugin_unit-t.cpp`: fake discovery backend exercising provider queue/normalization and environment fallback. -- `test/tap/tests/unit/aws_locality_stats_unit-t.cpp`: plugin-conditional table lifecycle and query-time projection. -- `test/tap/tests/unit/Makefile`, `test/tap/groups/groups.json`: targets and CI groups. -- `doc/aws-locality-awareness.md`, `README.md`: operator configuration, permissions, behavior, and diagnostics. - ---- - -### Task 1: SDK-Free Policy, Classification, and Weight Arithmetic - -**Files:** -- Create: `include/Aws_Locality_Types.h` -- Create: `include/Aws_Locality_Manager.h` -- Create: `lib/Aws_Locality_Manager.cpp` -- Modify: `lib/Makefile` -- Create: `test/tap/tests/unit/aws_locality_policy_unit-t.cpp` -- Modify: `test/tap/tests/unit/Makefile` -- Modify: `test/tap/groups/groups.json` - -**Interfaces:** -- Produces `AwsLocalityPolicy parse_aws_locality_policy(const nlohmann::json&, uint32_t hostgroup_id, AwsLocalityPolicyError&)`. -- Produces `AwsEndpointCandidate recognize_rds_endpoint(uint32_t, std::string_view, uint16_t)`. -- Produces `AwsLocalityClass classify_aws_locality(const AwsLocalLocation&, const AwsBackendLocation&)`. -- Produces `uint64_t aws_locality_effective_weight(int64_t configured_weight, double multiplier)`. - -- [ ] **Step 1: Write the failing policy and arithmetic test** - - Use literal cases for defaults, explicit timing, missing required multipliers, wrong JSON types, NaN-like invalid values, `1.0`/`10.0` bounds, ordering, truncation, zero, and `uint64_t` saturation. Include official instance/cluster/reader/custom endpoint names in standard, GovCloud, and China partitions, plus custom CNAME/RDS Proxy/arbitrary-host negatives. - - ```cpp - AwsLocalityPolicyError error; - const auto policy = parse_aws_locality_policy(json::parse( - R"({"same_region_multiplier":2.5,"same_az_multiplier":4.75})"), 10, error); - ok(policy.valid && policy.refresh_interval_seconds == 300 && - policy.stale_ttl_seconds == 1800, "valid policy uses timing defaults"); - ok(aws_locality_effective_weight(3, 2.5) == 7, - "effective weight truncates toward zero"); - ``` - -- [ ] **Step 2: Run the focused target and verify RED** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_policy_unit-t` - - Expected: compilation fails only because the new locality types/functions are absent. - -- [ ] **Step 3: Implement the minimal pure model** - - Define SDK-free enums and structs with owned strings: - - ```cpp - enum class AwsEndpointType : uint8_t { unknown, instance, cluster, reader, custom }; - enum class AwsLocalityClass : uint8_t { unknown, remote, same_region, same_az }; - enum class AwsLocalityMetadataStatus : uint8_t { disabled, pending, fresh, stale, expired, error }; - - struct AwsLocalityPolicy { - bool valid{false}; - double same_region_multiplier{1.0}; - double same_az_multiplier{1.0}; - uint32_t refresh_interval_seconds{300}; - uint32_t stale_ttl_seconds{1800}; - }; - ``` - - Normalize hostnames by ASCII-lowercasing and removing one trailing dot. Recognition extracts only candidate Region/partition and never claims authoritative endpoint type. Use `long double` for multiplication, truncate toward zero, preserve zero, and saturate at `uint64_t::max()`. - -- [ ] **Step 4: Run focused GREEN and existing parser/selection regressions** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_policy_unit-t server_selection_unit-t aws_iam_policy_unit-t` - - Expected: all TAP plans pass with no warnings or SDK references in the focused locality binary. - -- [ ] **Step 5: Commit** - - ```bash - git add include/Aws_Locality_Types.h include/Aws_Locality_Manager.h \ - lib/Aws_Locality_Manager.cpp lib/Makefile \ - test/tap/tests/unit/aws_locality_policy_unit-t.cpp \ - test/tap/tests/unit/Makefile test/tap/groups/groups.json - git commit -m "feat(mysql): parse AWS locality policies" - ``` - ---- - -### Task 2: Provider Registry, Refresh Manager, and Immutable Snapshots - -**Files:** -- Modify: `include/Aws_Locality_Types.h` -- Modify: `include/Aws_Locality_Manager.h` -- Modify: `lib/Aws_Locality_Manager.cpp` -- Create: `test/tap/tests/unit/aws_locality_manager_unit-t.cpp` -- Modify: `test/tap/tests/unit/Makefile` -- Modify: `test/tap/groups/groups.json` - -**Interfaces:** -- Consumes Task 1 policy/candidate/classification types. -- Produces `AwsMetadataProvider`, `AwsMetadataCompletionSink`, `AwsMetadataProviderLease`, `install_global_aws_metadata_provider()`, `acquire_global_aws_metadata_provider()`, and `shutdown_global_aws_metadata_provider()`. -- Produces `MySQLAwsLocalityManager::{configure,set_enabled,snapshot,diagnostic_rows,shutdown}`. - -- [ ] **Step 1: Write a fake-provider manager test** - - Exercise a real manager against a deterministic provider that retains request IDs/generations and posts complete normalized results. Cover local discovery followed by coalesced regional requests, duplicate endpoints across hostgroups, shortest refresh interval, pending/fresh/stale/expired/error states, endpoint-not-found, provider absence, enable/disable/re-enable, invalid reload removal, late-generation rejection, cancel, provider replacement, and shutdown drain. - - ```cpp - class FakeAwsMetadataProvider final : public AwsMetadataProvider { - public: - AwsMetadataRequestHandle request(const AwsMetadataRequest& request, - std::weak_ptr sink) override; - void cancel(AwsMetadataRequestHandle handle) override; - }; - ``` - - Use an injected steady/wall clock; no sleeps for freshness assertions. Verify monotonic age calculations and wall-clock diagnostic timestamps independently. Add a deterministic callback-vs-shutdown test and a 100-repeat multi-producer publication test. - -- [ ] **Step 2: Run and verify RED** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_manager_unit-t` - - Expected: compile failure on the absent provider/manager APIs. - -- [ ] **Step 3: Implement the provider registry and lease** - - Mirror the proven IAM lease/drain contract, but keep a separate generic metadata registry. Installation transfers provider ownership plus an optional retained module handle. Shutdown disables new leases, waits for active leases, calls provider shutdown/destructor, and only then `dlclose()`s the retained module reference. - -- [ ] **Step 4: Implement manager scheduling and publication** - - `configure()` receives a copied vector of hostgroup policy/backend identities and advances a generation. A lazy scheduler thread exists only while enabled with at least one valid policy. It requests local identity and coalesces endpoint scans by Region, never holding the manager mutex across provider calls. Completions update mutable cache state under the manager mutex, then build and atomically publish `std::shared_ptr` objects. Selection snapshots contain no mutable locks or plugin pointers. - -- [ ] **Step 5: Run focused GREEN and TSan** - - Run: - - ```bash - PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_manager_unit-t - PROXYSQL40=1 NOJEMALLOC=1 WITHTSAN=1 make -C test/tap/tests/unit -j aws_locality_manager_unit-t - TSAN_OPTIONS=halt_on_error=1 test/tap/tests/unit/aws_locality_manager_unit-t - ``` - - Expected: full TAP plan passes and TSan reports no race. - -- [ ] **Step 6: Commit** - - ```bash - git add include/Aws_Locality_Types.h include/Aws_Locality_Manager.h \ - lib/Aws_Locality_Manager.cpp test/tap/tests/unit/aws_locality_manager_unit-t.cpp \ - test/tap/tests/unit/Makefile test/tap/groups/groups.json - git commit -m "feat(mysql): manage asynchronous AWS locality metadata" - ``` - ---- - -### Task 3: MySQL Variable and Hostgroup Reload Integration - -**Files:** -- Modify: `include/Base_HostGroups_Manager.h` -- Modify: `include/MySQL_HostGroups_Manager.h` -- Modify: `include/MySQL_Thread.h` -- Modify: `lib/BaseHGC.cpp` -- Modify: `lib/MySQL_HostGroups_Manager.cpp` -- Modify: `lib/MySQL_Thread.cpp` -- Modify: `lib/Admin_FlushVariables.cpp` -- Create: `test/tap/tests/unit/aws_locality_config_unit-t.cpp` -- Modify: `test/tap/tests/unit/Makefile` -- Modify: `test/tap/groups/groups.json` - -**Interfaces:** -- Consumes `MySQLAwsLocalityManager::configure()` and `set_enabled()`. -- Produces `MyHGC::attributes.aws_locality_policy` and `MySQL_HostGroups_Manager::refresh_aws_locality_configuration()`. -- Produces the dynamic boolean variable `mysql-aws_locality_awareness`, default `false`, only in `PROXYSQL40`. - -- [ ] **Step 1: Write failing configuration lifecycle tests** - - Initialize real hostgroup attributes from JSON and assert valid policy installation, exact defaults/bounds, malformed field rejection without full JSON logging, and removal of the prior policy after invalid reload. Exercise `MySQL_Threads_Handler::{set_variable,get_variable,commit}` for false/true parsing and notify a real manager after `LOAD` semantics. Compare `runtime_mysql_servers` and checksum input before/after metadata completions. - -- [ ] **Step 2: Run and verify RED** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_config_unit-t` - - Expected: compile failure on missing variable/policy fields and refresh method. - -- [ ] **Step 3: Add the variable and policy storage** - - Register `aws_locality_awareness` in the existing bool variable table, default it false, copy it to worker variables, and expose it as `mysql-aws_locality_awareness`. Compile all feature behavior under `PROXYSQL40`. Add an owned policy value to `MyHGC` and reset it on every attributes reload before parsing. - -- [ ] **Step 4: Wire load boundaries** - - At the end of `MySQL_HostGroups_Manager::commit()`, while server/hostgroup state is stable, copy valid policies and backend identities and call `configure()` after releasing the HGM lock. After MySQL variable commit and lock release, call `set_enabled()` once with the master value. Never include discovered data in generated tables/checksums. - -- [ ] **Step 5: Run focused GREEN and existing config regressions** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_config_unit-t aws_iam_connection_config_unit-t hostgroups_unit-t cluster_sync_unit-t` - -- [ ] **Step 6: Commit** - - ```bash - git add include/Base_HostGroups_Manager.h include/MySQL_HostGroups_Manager.h \ - include/MySQL_Thread.h lib/BaseHGC.cpp lib/MySQL_HostGroups_Manager.cpp \ - lib/MySQL_Thread.cpp lib/Admin_FlushVariables.cpp \ - test/tap/tests/unit/aws_locality_config_unit-t.cpp \ - test/tap/tests/unit/Makefile test/tap/groups/groups.json - git commit -m "feat(mysql): load AWS locality configuration" - ``` - ---- - -### Task 4: Global and Thread-Local Weighted Selection - -**Files:** -- Modify: `include/Aws_Locality_Manager.h` -- Modify: `include/MySQL_HostGroups_Manager.h` -- Modify: `lib/MyHGC.cpp` -- Modify: `lib/MySQL_Thread.cpp` -- Create: `test/tap/tests/unit/aws_locality_selection_unit-t.cpp` -- Modify: `test/tap/tests/unit/Makefile` -- Modify: `test/tap/groups/groups.json` - -**Interfaces:** -- Consumes immutable `AwsLocalitySnapshot` and `effective_weight(hostgroup,host,port,configured_weight)`. -- Produces identical server-level weighted selection semantics in `MyHGC::get_random_MySrvC()` and `MySQL_Thread::get_MyConn_local()`. - -- [ ] **Step 1: Write failing global/local selection tests** - - Use real `MyHGC`, `MySrvC`, and cached `MySQL_Connection` fixtures. Publish literal same-AZ/same-Region/remote/unknown metadata for weights 10/20/30 and multipliers 4.0/2.0; assert deterministic effective weights 40/40/30. Verify same-AZ requires matching account, cluster/reader/custom receive only Region bias, stale remains active, expired is neutral, and configured weights never change. - - For local cache, place multiple connections on one remote parent and one connection on one local parent. Run seeded selections and prove probability follows parent weights, not connection count. Include health, GTID, lag, auth compatibility, session state, and backoff exclusions before locality. - -- [ ] **Step 2: Run and verify RED** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_selection_unit-t` - - Expected: locality distribution assertions fail while legacy controls pass. - -- [ ] **Step 3: Integrate global selection** - - Preserve the entire existing eligibility scan. Retain one snapshot before candidate evaluation only when the thread-local master flag and hostgroup policy are active. Store a parallel `uint64_t` effective-weight array, use a saturating 64-bit sum, and run the existing lottery over those values. The inactive path retains current branches and configured-weight arithmetic. - -- [ ] **Step 4: Integrate local-cache selection** - - Keep the existing first-match implementation unchanged when locality is inactive. When active, scan compatible eligible connections, group candidate indices by `MySrvC*`, calculate one effective weight per parent, choose a parent with the same helper/lottery, then remove one best compatible connection belonging to that parent. Do not allocate/call plugins on the inactive path; reuse bounded stack storage before falling back to a vector for unusually large candidate sets. - -- [ ] **Step 5: Run focused GREEN and pool regressions** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_selection_unit-t server_selection_unit-t aws_iam_pool_unit-t connection_pool_unit-t` - -- [ ] **Step 6: Commit** - - ```bash - git add include/Aws_Locality_Manager.h include/MySQL_HostGroups_Manager.h \ - lib/MyHGC.cpp lib/MySQL_Thread.cpp \ - test/tap/tests/unit/aws_locality_selection_unit-t.cpp \ - test/tap/tests/unit/Makefile test/tap/groups/groups.json - git commit -m "feat(mysql): apply AWS locality during backend selection" - ``` - ---- - -### Task 5: AWS Plugin Metadata Provider - -**Files:** -- Create: `plugins/aws/src/aws_locality_provider.h` -- Create: `plugins/aws/src/aws_locality_provider.cpp` -- Modify: `plugins/aws/src/aws_plugin.cpp` -- Modify: `plugins/aws/Makefile` -- Modify: `include/ProxySQL_Plugin.h` -- Modify: `lib/ProxySQL_PluginManager.cpp` -- Create: `test/tap/tests/unit/aws_locality_plugin_unit-t.cpp` -- Modify: `test/tap/tests/unit/Makefile` -- Modify: `test/tap/groups/groups.json` - -**Interfaces:** -- Consumes Task 2 `AwsMetadataProvider` and ownership callbacks. -- Produces ABI-6 services `install_aws_metadata_provider` and the `aws_locality` advertised plugin capability. -- Produces `AwsSdkMetadataProvider`, with an injectable `AwsLocalityDiscoveryBackend` for deterministic tests. - -- [ ] **Step 1: Write the failing provider test** - - Drive the real bounded queue/provider with a fake discovery backend. Assert full request/result fields, two-worker bound, queue rejection, cancellation, deadline rejection before/after work, no callback after shutdown, redacted categories, and generation/opaque ID preservation. Test environment precedence (`AWS_REGION`, `AWS_DEFAULT_REGION`, `AWS_AVAILABILITY_ZONE`, `AWS_ACCOUNT_ID`) and partial fallback using scoped environment restoration. - - Build literal normalized response fixtures for RDS instances, Aurora instances, cluster writer/reader/custom endpoints, missing ports, duplicate pages, Multi-AZ/failover changes, and endpoint-not-found. Never assert on fake call existence alone; assert emitted normalized results. - -- [ ] **Step 2: Run and verify RED** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_plugin_unit-t` - - Expected: compilation fails on missing provider/backend types. - -- [ ] **Step 3: Extend the ABI and shared SDK lifetime** - - Increment the plugin ABI maximum/current version to 6 for metadata-provider installation, to 7 for the MySQL-owned locality-stats projection callback, and to 8 for partial-init provider rollback; append those services to `ProxySQL_PluginServices`. Wire provider installation and rollback only during plugin init. Refactor the plugin so IAM signer/token source and locality provider each retain a `std::shared_ptr`; `Aws::InitAPI` occurs once and `Aws::ShutdownAPI` occurs only after both core-owned capabilities drain. - -- [ ] **Step 4: Implement local discovery** - - Use IMDSv2 token `PUT /latest/api/token` with a bounded TTL header, then `GET /latest/dynamic/instance-identity/document` with the token. Use ProxySQL's vendored libcurl already linked into the plugin, enforce link-local target/timeouts/response bounds, and parse only `region`, `availabilityZone`, and `accountId`. On IMDS failure, apply the exact environment fallback order. Never log the document, token, account ID, raw curl error, or environment values. - -- [ ] **Step 5: Implement paginated RDS discovery** - - Maintain regional `Aws::RDS::RDSClient` instances under a client-map mutex. Issue paginated `DescribeDBInstances`, `DescribeDBClusters`, and `DescribeDBClusterEndpoints` requests until marker exhaustion or deadline/cancellation. Normalize only authoritative endpoint addresses and supplied ports. Map SDK errors to fixed categories (`access_denied`, `throttled`, `timeout`, `invalid_response`, `provider_unavailable`) and discard raw messages. Rate-limit logs by stable Region/endpoint/category keys without including raw responses or account identity. - -- [ ] **Step 6: Run focused GREEN, plugin loader, and secret scans** - - Run: - - ```bash - PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_plugin_unit-t aws_plugin_load_unit-t - PROXYSQL40=1 make -C plugins/aws -j - nm -C plugins/aws/ProxySQL_Aws_Plugin.so > /tmp/proxysql-aws-plugin-nm.txt - ldd plugins/aws/ProxySQL_Aws_Plugin.so > /tmp/proxysql-aws-plugin-ldd.txt - ``` - - Assert the plugin has locality symbols, has no AWS/CRT shared-library dependencies, and test/log output contains none of the fixture credentials/account IDs/tokens. - -- [ ] **Step 7: Commit** - - ```bash - git add plugins/aws/src/aws_locality_provider.h plugins/aws/src/aws_locality_provider.cpp \ - plugins/aws/src/aws_plugin.cpp plugins/aws/Makefile include/ProxySQL_Plugin.h \ - lib/ProxySQL_PluginManager.cpp test/tap/tests/unit/aws_locality_plugin_unit-t.cpp \ - test/tap/tests/unit/Makefile test/tap/groups/groups.json - git commit -m "feat(aws): discover RDS locality metadata" - ``` - ---- - -### Task 6: Plugin-Conditional Stats Table - -**Files:** -- Modify: `include/ProxySQL_Plugin.h` -- Modify: `lib/ProxySQL_PluginManager.cpp` -- Modify: `include/MySQL_HostGroups_Manager.h` -- Modify: `lib/MySQL_HostGroups_Manager.cpp` -- Modify: `plugins/aws/src/aws_plugin.cpp` -- Create: `test/tap/tests/unit/aws_locality_stats_unit-t.cpp` -- Modify: `test/tap/tests/unit/Makefile` -- Modify: `test/tap/groups/groups.json` - -**Interfaces:** -- Produces the plugin-owned `stats_mysql_aws_locality` schema and runtime-view registration. -- Produces ABI service `refresh_mysql_aws_locality_stats(SQLite3DB*)`, callable by the plugin's static refresh callback. -- Consumes `MySQLAwsLocalityManager::diagnostic_rows()`. - -- [ ] **Step 1: Write failing table lifecycle/projection tests** - - Bootstrap Admin with no AWS plugin and assert `SELECT * FROM stats_mysql_aws_locality` returns `no such table`. Bootstrap through the real AWS plugin schema-registration phase and assert the exact 17-column schema exists. Publish manager rows for pending/fresh/stale/expired/error/disabled, query through `ProxySQL_Admin` twice across a generation swap, and assert each result is a complete single-generation snapshot. - - Use a provider request counter to prove querying the table issues zero metadata requests. Assert writes are rejected and no disk/config table/checksum contains the name. - -- [ ] **Step 2: Run and verify RED** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_stats_unit-t` - - Expected: plugin-loaded table query fails because the schema/callback is absent. - -- [ ] **Step 3: Register schema and refresh callback** - - In `register_schemas`, register only a stats-db table: - - ```sql - CREATE TABLE stats_mysql_aws_locality ( - hostgroup_id INT NOT NULL, hostname VARCHAR NOT NULL, port INT NOT NULL, - endpoint_type VARCHAR NOT NULL, configured_weight INT NOT NULL, - effective_weight INT NOT NULL, local_region VARCHAR NOT NULL, - local_az VARCHAR NOT NULL, backend_region VARCHAR NOT NULL, - backend_az VARCHAR NOT NULL, account_match VARCHAR NOT NULL, - locality VARCHAR NOT NULL, active_multiplier REAL NOT NULL, - metadata_status VARCHAR NOT NULL, last_success_timestamp INT NOT NULL, - last_attempt_timestamp INT NOT NULL, last_error_category VARCHAR NOT NULL, - PRIMARY KEY(hostgroup_id, hostname, port)) - ``` - - The plugin callback invokes the core service. Core retains one diagnostics snapshot, executes `BEGIN; DELETE; INSERT...; COMMIT`, and never calls the provider. With the master off, preserve cached text but force multiplier 1.0/effective configured/status disabled. No valid policies means zero rows. - -- [ ] **Step 4: Run focused GREEN and lifecycle regressions** - - Run: `PROXYSQL40=1 make -C test/tap/tests/unit -j aws_locality_stats_unit-t aws_plugin_load_unit-t plugin_runtime_views_unit-t test_aws_iam_metrics-t` - -- [ ] **Step 5: Commit** - - ```bash - git add include/ProxySQL_Plugin.h lib/ProxySQL_PluginManager.cpp \ - include/MySQL_HostGroups_Manager.h lib/MySQL_HostGroups_Manager.cpp \ - plugins/aws/src/aws_plugin.cpp test/tap/tests/unit/aws_locality_stats_unit-t.cpp \ - test/tap/tests/unit/Makefile test/tap/groups/groups.json - git commit -m "feat(stats): expose AWS locality decisions" - ``` - ---- - -### Task 7: Operator Documentation and Final Verification - -**Files:** -- Create: `doc/aws-locality-awareness.md` -- Modify: `README.md` -- Modify: `.github/workflows/CI-aws.yml` -- Modify: `docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md` only if implementation review exposes an approved contract correction. - -**Interfaces:** -- Consumes all preceding production/test behavior. -- Produces operator documentation and CI gates; no new runtime API. - -- [ ] **Step 1: Write operator documentation** - - Document the variable, JSON example, numeric/timing bounds, integer truncation, non-cumulative tiers, instance-vs-cluster AZ behavior, account requirement, environment fallback, EC2/EKS credential delivery, exact read-only RDS IAM policy, stale/fail-neutral behavior, and plugin-conditional stats table. State explicitly that displayed/runtime configured weights never change. - -- [ ] **Step 2: Extend established container CI** - - Add the five locality tests to the existing AWS workflow's established ProxySQL build container. Preserve `actions/checkout` LFS hydration, pass `PROXYSQL40=1 make -j` from workflow invocations, and do not install compiler/development dependencies directly onto the GitHub runner. - -- [ ] **Step 3: Run the complete normal regression gate** - - Run: - - ```bash - PROXYSQL40=1 make -j clean - PROXYSQL40=1 make -j - PROXYSQL40=1 make -C test/tap/tests/unit -j - ``` - - Execute every generated unit binary and record TAP totals. Run all existing AWS IAM, connection-pool, hostgroup, cluster/checksum, plugin lifecycle/runtime-view, controlled TLS, and metrics targets explicitly. - -- [ ] **Step 4: Run sanitizer gates** - - Build and run policy, manager, config, selection, plugin, stats, and affected IAM/pool tests under ASan+LSan. Run manager, selection, provider, stats, and plugin lifecycle concurrency tests under TSan. Restore normal artifacts afterward using `PROXYSQL40=1 make -j clean && PROXYSQL40=1 make -j`. - -- [ ] **Step 5: Run linkage/security/final-diff gates** - - Capture `nm` and `ldd` output to files before grepping. Prove daemon/archive contain no `Aws::` symbols and daemon/plugin have no AWS/CRT DSOs; prove plugin contains expected static SDK/locality symbols. Scan the combined diff/test output for credentials, tokens, account IDs, raw AWS errors, and unredacted fixture markers. Run `git diff --check` and validate the vendored archive remains unmodified. - - If an externally provisioned AWS integration runner, credentials, and RDS endpoints are configured, run one instance-endpoint and one cluster/reader-endpoint locality test. Otherwise record the optional gate as `NOT RUN`; never substitute fake-provider coverage and label it real AWS verification. - -- [ ] **Step 6: Request independent review and fix all Critical/Important findings** - - Provide the reviewer the approved design, this plan, base SHA, head SHA, exact verification evidence, and explicit non-goals. For every valid finding, write a focused failing test before the production correction, then rerun affected and full gates. - -- [ ] **Step 7: Commit documentation/CI and prepare handoff** - - ```bash - git add doc/aws-locality-awareness.md README.md .github/workflows/CI-aws.yml - git commit -m "docs: document AWS locality awareness" - ``` - - Do not push or open/retarget a PR until the user requests publication. diff --git a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md b/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md deleted file mode 100644 index 9991ab28e4..0000000000 --- a/docs/superpowers/specs/2026-08-13-aws-locality-awareness-design.md +++ /dev/null @@ -1,679 +0,0 @@ -# AWS Locality-Aware MySQL Backend Selection - -**Status:** Approved design - -**Date:** 2026-08-13 - -## Summary - -ProxySQL 4.0 will optionally prefer MySQL backends that are in the same AWS -Region or Availability Zone as the ProxySQL process. The feature changes only -the temporary weights used by a server-selection attempt. It never modifies -`mysql_servers.weight`, `runtime_mysql_servers.weight`, the saved -configuration, or ProxySQL Cluster checksums. - -The MySQL module owns the feature's configuration and traffic policy. The -general AWS plugin provides asynchronous, normalized AWS metadata and makes no -traffic-routing decisions. - -The first version supports RDS and Aurora endpoints. It does not discover -arbitrary MySQL servers on EC2. - -## Goals - -- Preserve configured server weights while giving operators a bounded local - Region and local AZ preference. -- Keep all AWS, IMDS, DNS, and credential-provider work out of connection - selection. -- Keep locality configuration in the MySQL module because MySQL Hostgroup - Manager consumes it. -- Discover every ProxySQL process's location independently so ProxySQL Cluster - cannot propagate one process's Region or AZ to another. -- Degrade to ordinary configured weights whenever metadata is unavailable, - expired, invalid, or unsupported. -- Make every classification and active multiplier observable without changing - `runtime_mysql_servers`. -- Reuse the general AWS plugin and its statically linked vendored AWS SDK - runtime alongside the IAM database-authentication capability. - -## Non-goals - -- Mutating configured or runtime server weights. -- Synchronizing discovered metadata through ProxySQL Cluster. -- Moving, closing, or rebalancing existing backend connections. -- Discovering arbitrary EC2-hosted MySQL servers. -- Resolving custom CNAMEs to infer an RDS or Aurora target. -- Scanning every AWS Region to locate an endpoint. -- Assuming roles into other AWS accounts. -- Managing RDS/Aurora topology or hostgroup membership. -- Supporting RDS Proxy endpoints in the first version. -- Providing PostgreSQL locality awareness in the first version. -- Making real AWS infrastructure mandatory for ordinary CI. - -## User-facing configuration - -### Global master switch - -The MySQL module adds one dynamic global variable: - -```text -mysql-aws_locality_awareness = false -``` - -It is available in the ProxySQL 4.0 build and defaults to `false`. - -Changing it follows the normal MySQL-variable lifecycle: - -```sql -SET mysql-aws_locality_awareness = true; -LOAD MYSQL VARIABLES TO RUNTIME; -SAVE MYSQL VARIABLES TO DISK; -``` - -There are deliberately no configured Region, AZ, or AWS-account variables. -Such variables could be synchronized to ProxySQL processes in other locations -and would therefore be unsafe. - -When the switch is disabled: - -- selection uses configured weights exactly as it does today; -- no new locality metadata refreshes are scheduled; -- in-flight locality requests are cancelled or ignored by generation; -- when the AWS plugin is loaded, cached rows remain visible in its diagnostic - table with status `disabled`, but cannot affect selection. - -### Per-hostgroup policy - -The existing `mysql_hostgroup_attributes.hostgroup_settings` JSON is the -configuration extension point: - -```json -{ - "aws": { - "locality_awareness": { - "same_region_multiplier": 2.0, - "same_az_multiplier": 4.0, - "refresh_interval_seconds": 300, - "stale_ttl_seconds": 1800 - } - } -} -``` - -Presence of a valid `aws.locality_awareness` object enables locality awareness -for that hostgroup. There is no additional per-hostgroup `enabled` field. - -Both multipliers are required and must be finite JSON numbers satisfying: - -```text -1.0 <= same_region_multiplier <= same_az_multiplier <= 10.0 -``` - -The timing fields are optional. Their defaults and accepted bounds are: - -```text -refresh_interval_seconds = 300 -stale_ttl_seconds = 1800 - -30 <= refresh_interval_seconds <= 86400 -refresh_interval_seconds <= stale_ttl_seconds <= 604800 -``` - -An invalid locality object disables locality bias for that hostgroup after the -load. Diagnostics identify the rejected field and hostgroup but never log the -complete JSON document. - -The existing `aws_iam_region` key remains an IAM-authentication setting. It is -not required by, or treated as authoritative for, locality discovery. - -## Selection semantics - -Locality produces a temporary effective weight for an eligible server: - -```text -remote or unknown configured_weight -same Region, different AZ int(configured_weight * same_region_multiplier) -same AZ int(configured_weight * same_az_multiplier) -``` - -The Region and AZ tiers are mutually exclusive. A same-AZ server receives only -`same_az_multiplier`; the two multipliers are never multiplied together. - -Conversion to an integer truncates toward zero. Weight zero remains zero. -Arithmetic uses a wide intermediate and saturates safely before entering the -64-bit weighted-selection accumulator. The current MySQL weight bounds make -saturation unlikely, but the operation must still be defined for every input. - -Examples with configured weights `10`, `20`, and `30`, Region multiplier -`2.0`, and AZ multiplier `4.0`: - -- first server in the same AZ: effective weight `40`; -- second server in the same Region but another AZ: effective weight `40`; -- third server in another Region: effective weight `30`. - -The values `10`, `20`, and `30` remain stored and reported by -`mysql_servers` and `runtime_mysql_servers`. - -### Classification rules - -The selector classifies a backend from one immutable metadata snapshot: - -- `same_az`: local and backend Regions match, local and backend AZ names match, - and both sides have the same confirmed AWS account ID; -- `same_region`: Regions match, but AZ is different, unavailable, inapplicable, - or cannot be trusted because account identity is unavailable or different; -- `remote`: both Regions are known and differ; -- `unknown`: either Region required for comparison is unknown. - -AZ names can map to different physical zones in different AWS accounts. The -same-AZ multiplier is therefore never applied without a same-account check. -Same-Region preference does not require matching accounts. - -### Existing eligibility remains authoritative - -Locality changes only the weighted lottery among candidates that have already -passed the existing rules, including: - -- ONLINE status and shun recovery; -- `max_connections` capacity; -- latency bounds; -- GTID requirements; -- replication-lag and Aurora-lag requirements; -- session-tracking capability backoff; -- the existing Aurora writer/replica filtering. - -Locality never makes a backend healthy, eligible, or available. - -### Global and thread-local pool paths - -The global path in `MyHGC::get_random_MySrvC()` computes each final -candidate's effective weight and performs its existing weighted selection with -a 64-bit accumulator. - -The per-thread local connection cache must also honor locality. Otherwise, a -remote cached connection could repeatedly bypass Hostgroup Manager's weighted -lottery. For locality-enabled hostgroups only, the local-cache path: - -1. finds connections that pass all existing compatibility, GTID, lag, health, - and session-state checks; -2. groups them by parent server, so a server with more idle connections does - not gain more selection probability; -3. selects a parent server using the same effective server weight helper; -4. returns a compatible cached connection belonging to that parent. - -When locality is inactive, the current local-cache first-match fast path stays -unchanged. - -One snapshot is retained for an entire selection attempt, preventing a refresh -from mixing classifications within one lottery. Metadata changes affect only -future selections. Existing connections are not migrated or closed. - -## Ownership and architecture - -### Chosen approach - -Core owns locality state and consumes an asynchronous AWS metadata provider. - -This was selected over two alternatives: - -1. A synchronous cached lookup into the plugin from every selection would add - plugin ABI calls and lifecycle/synchronization risk to a hot path. -2. A plugin-owned server-selection hook would move MySQL traffic policy into - the capability provider and violate the intended ownership boundary. - -### MySQL core responsibilities - -A core `MySQLAwsLocalityManager` owns: - -- parsed hostgroup policies; -- registered backend endpoint identities; -- configuration generations; -- refresh scheduling and request coalescing; -- normalized results received from the plugin; -- last-attempt and last-success times; -- per-policy fresh/stale/expired evaluation; -- immutable snapshots used by selection; -- diagnostic snapshot data and redacted failure state. - -Core types contain only ProxySQL-owned strings, enums, timestamps, request -IDs, and result structures. They expose no AWS SDK types. - -The manager starts work only when the master switch is enabled and at least one -hostgroup has a valid locality policy. `LOAD MYSQL SERVERS TO RUNTIME` rebuilds -the endpoint registration set, advances its generation, and schedules the -necessary asynchronous refreshes. `LOAD MYSQL VARIABLES TO RUNTIME` activates -or bypasses the manager according to the master switch. - -### AWS plugin responsibilities - -The general `aws` plugin owns: - -- AWS SDK initialization and shutdown; -- the default AWS credential-provider chain; -- IMDSv2 access; -- regional RDS clients; -- paginated RDS API calls; -- bounded retries, timeouts, and background execution; -- cancellation and clean shutdown; -- normalization into the core-defined result contract. - -It makes no multiplier, eligibility, hostgroup, or traffic decision. - -The plugin extends its advertised capabilities beyond `aws_iam`, for example -with local-instance metadata and RDS-topology capabilities. The plugin reuses -the same SDK runtime already used by IAM authentication. - -The plugin also registers the `stats_mysql_aws_locality` schema and its -query-time refresh callback. The MySQL module remains the source of the rows; -the plugin registration only makes the AWS-specific diagnostic surface exist -when the AWS capability is actually present. - -### Generic asynchronous provider ABI - -ProxySQL's plugin services gain a generic AWS metadata-provider installation -contract. The first request kinds are: - -- discover the local ProxySQL process's AWS location; -- describe the RDS/Aurora endpoints in one candidate Region. - -Requests carry opaque IDs, deadlines, endpoint sets, and core configuration -generations. Results contain normalized endpoint type, Region, AZ where -applicable, account identity for comparison, timestamps, and a redacted status -category. - -The provider uses the same lease/drain principle as the IAM token source: - -- a plugin module cannot unload while requests or callbacks retain leases; -- shutdown stops accepting work, cancels queued work, and drains active work; -- callbacks target weak/core-owned completion sinks; -- callbacks run without holding plugin or Hostgroup Manager locks; -- core rejects completions from an obsolete configuration generation. - -No selection path invokes this ABI. - -### Immutable snapshot publication - -Core publishes immutable per-hostgroup locality snapshots. A snapshot maps the -current stable backend identity `(hostgroup_id, normalized hostname, port)` to -its classification inputs and metadata timestamps. Publication is atomic; a -selection retains one snapshot for its duration and performs no network calls -or mutable-cache locking. - -The optional feature may pay for immutable-map lookups. With the global switch -off or no hostgroup policy, the existing hot path bypasses those lookups. - -## Local ProxySQL location discovery - -Every ProxySQL process discovers its own location independently. Discovery is -node-local state and is never persisted or cluster-synchronized. - -Discovery order: - -1. Retrieve the EC2 instance identity document through IMDSv2. It provides - Region, Availability Zone, and account ID. -2. If IMDSv2 is unavailable, use process environment fallback: - - Region: `AWS_REGION`, then `AWS_DEFAULT_REGION`; - - AZ: `AWS_AVAILABILITY_ZONE`; - - account assertion: optional `AWS_ACCOUNT_ID`. -3. Leave any unavailable field unknown. - -An environment AZ is usable only with an environment Region. The same-AZ tier -also requires `AWS_ACCOUNT_ID`; without it, same-Region preference still -works. In Kubernetes, operators can inject the node topology AZ and Region as -pod environment values without adding synchronized ProxySQL settings. - -The EC2 instance identity document and fields are documented by AWS at: - - - -Local metadata follows the same refresh and stale policy as backend metadata. -If local Region expires, all locality classifications are neutral. If local -Region remains usable but AZ/account becomes unavailable, same-Region -classification remains possible while same-AZ does not. - -## Backend endpoint discovery - -### Candidate recognition - -Core recognizes official RDS/Aurora endpoint DNS forms only to extract a -candidate AWS Region and partition. This is routing for the API request, not -authoritative metadata. - -The implementation normalizes endpoint hostnames by lowercasing ASCII and -removing one trailing DNS dot. It does not resolve DNS or follow CNAMEs. -Supported official suffixes include the standard/GovCloud AWS suffix and the -China partition suffix. An unrecognized endpoint remains neutral. - -### Authoritative API matching - -Requests are coalesced by candidate Region. The plugin performs paginated: - -- `rds:DescribeDBInstances`; -- `rds:DescribeDBClusters`; -- `rds:DescribeDBClusterEndpoints`. - -Core accepts a result only when the normalized configured hostname exactly -matches an endpoint returned by the AWS APIs. Where the response supplies a -port, a configured port mismatch is rejected. A custom cluster endpoint that -does not expose a distinct port is matched by its exact authoritative -hostname. - -The endpoint mappings are: - -- RDS DB instance endpoint: `instance`, with Region, instance AZ, and account; -- Aurora DB instance endpoint: `instance`, with Region, instance AZ, and - account; -- Aurora or Multi-AZ cluster writer endpoint: `cluster`, with Region and - account but no stable endpoint AZ; -- Aurora reader endpoint: `reader`, with Region and account but no stable - endpoint AZ; -- Aurora custom endpoint: `custom`, with Region and account but no stable - endpoint AZ; -- unmatched or unsupported endpoint: `unknown`. - -Cluster, reader, and custom endpoints can route to instances in multiple AZs. -They can receive the same-Region multiplier but never the same-AZ multiplier. - -`DescribeDBInstances` exposes both an endpoint address and Availability Zone: - - - -`DescribeDBClusters` exposes cluster, reader, and member information: - - - -`DescribeDBClusterEndpoints` exposes custom and managed cluster endpoints: - - - -Custom CNAMEs, RDS Proxy endpoints, arbitrary hosts, and AWS-looking hostnames -that do not appear in an authoritative response stay `unknown`. - -## Refresh, sharing, and stale data - -Successful metadata records include a monotonic success time and a wall-clock -time for diagnostics. A failed refresh records an attempt time and redacted -error but does not immediately discard the last successful value. - -Each hostgroup evaluates freshness using its own policy: - -- `fresh`: age is no greater than `refresh_interval_seconds`; -- `stale`: a refresh is due or has failed, but age is no greater than - `stale_ttl_seconds`; the last successful metadata remains active; -- `expired`: age exceeds `stale_ttl_seconds`; metadata becomes unknown and the - configured weight is used; -- `error`: no usable successful metadata exists for the endpoint; -- `pending`: discovery has not completed yet; -- `disabled`: the global switch is off. - -If several hostgroups reference the same endpoint with different intervals, -the endpoint is refreshed at the shortest active interval. The shared result -retains its success timestamp; each hostgroup independently determines whether -that result is fresh, stale, or expired under its own TTL. - -Regional API scans are coalesced so one in-flight scan serves all registered -endpoints in that Region. Repeated load operations cancel or supersede older -generations. Late completions cannot attach to removed servers, removed -policies, or a newer generation. - -A successful scan that does not contain a configured endpoint records -`endpoint_not_found`. A prior match may remain active only through its bounded -stale TTL, after which the endpoint becomes neutral. - -## Failure behavior and security - -Every failure is fail-neutral, not fail-closed for database traffic: - -- missing AWS plugin; -- plugin unload or shutdown; -- missing credentials; -- IMDS disabled or unreachable; -- Kubernetes without injected location; -- access denied; -- throttling; -- timeout; -- malformed or unsupported endpoint; -- endpoint not found; -- callback cancellation; -- metadata expiration. - -In all cases, the backend remains subject to its ordinary eligibility and -configured weight. - -Logs are rate-limited by stable endpoint/Region/error-category keys. They never -include credentials, authorization headers, IMDS tokens, raw AWS errors, -account IDs, or the complete hostgroup JSON. Supported fixed categories -include: - -```text -access_denied -throttled -provider_unavailable -imds_unavailable -endpoint_not_found -timeout -cancelled -invalid_response -``` - -The plugin uses the normal AWS SDK credential provider chain. ProxySQL adds no -access-key or secret-key settings. Expected deployments include EC2 instance -profiles, EKS IRSA or Pod Identity, and externally provided standard AWS -credential sources. - -The read-only RDS policy required for backend discovery is: - -```json -{ - "Version": "2012-10-17", - "Statement": [ - { - "Effect": "Allow", - "Action": [ - "rds:DescribeDBInstances", - "rds:DescribeDBClusters", - "rds:DescribeDBClusterEndpoints" - ], - "Resource": "*" - } - ] -} -``` - -IMDS and environment discovery require no AWS API permission. - -## Observability - -The AWS plugin registers `stats_mysql_aws_locality` as a read-only table in the -stats database. Its existence follows the plugin lifecycle: - -- when the AWS plugin loads successfully, its schema-registration phase adds - the table before the Admin databases are materialized; -- when the AWS plugin is not configured or does not load successfully, the - table is not created, and querying it returns the normal SQLite - `no such table` error; -- ProxySQL does not currently support hot unloading configured plugins. If hot - unload is introduced, the unload contract must unregister and drop this - table rather than leave an empty or stale table behind. - -The table is a query-time projection of the MySQL locality manager's current -immutable in-memory snapshot. Before a query that references the table is -executed, Admin invokes the registered refresh callback. The callback replaces -the prior SQLite rows in one transaction from one retained manager snapshot, -so a result never mixes locality generations. This is the same materialized- -on-query model used by other runtime and stats views; the SQLite rows are not -the authoritative locality state. - -Refreshing the table never performs an IMDS or AWS API request and never waits -for metadata discovery. Network refresh remains bounded asynchronous plugin -work; a table query reports the most recently published state, including -`pending`, `stale`, `expired`, or `error` as applicable. - -The projection is non-persistent: it is not saved to disk, loaded to runtime, -included in ProxySQL Cluster checksums, or accepted as configuration. Writes -to it are unsupported. Each refresh emits one row for each backend in a -hostgroup that currently has a valid locality policy: - -```text -hostgroup_id -hostname -port -endpoint_type -configured_weight -effective_weight -local_region -local_az -backend_region -backend_az -account_match -locality -active_multiplier -metadata_status -last_success_timestamp -last_attempt_timestamp -last_error_category -``` - -Definitions: - -- `endpoint_type`: `instance`, `cluster`, `reader`, `custom`, or `unknown`; -- `account_match`: `same`, `different`, or `unknown`; -- `locality`: `same_az`, `same_region`, `remote`, or `unknown`; -- `active_multiplier`: the multiplier currently affecting selection, otherwise - `1.0`; -- `effective_weight`: a diagnostic calculation only; it is never written back - to a server table; -- timestamps: Unix epoch seconds, or zero if the event has never occurred; -- `metadata_status`: one of the lifecycle states defined above. - -Account IDs are never exposed. When the AWS plugin is loaded but the master -switch is disabled, the table remains present and its rows retain cached -location text for diagnosis. Every row has `active_multiplier` equal to `1.0`, -`effective_weight` equal to `configured_weight`, and `metadata_status` equal -to `disabled`. If no hostgroup currently has a valid locality policy, the -table exists but is empty. - -## Runtime sequence - -1. ProxySQL loads the AWS plugin and installs its generic metadata provider. -2. `LOAD MYSQL VARIABLES TO RUNTIME` enables the global feature. -3. `LOAD MYSQL SERVERS TO RUNTIME` parses locality policies, advances the core - registration generation, and schedules local and regional discovery. -4. The AWS plugin performs IMDS and RDS work on bounded background workers. -5. Completions return normalized metadata to the core manager. -6. Core validates request ID and generation, updates timestamps/error state, - and atomically publishes immutable hostgroup snapshots. -7. Global and local-cache selection attempts retain one snapshot and calculate - temporary effective weights for eligible server parents. -8. Periodic refreshes repeat at the shortest interval required by registered - hostgroups. Stale and expiration decisions remain per hostgroup. -9. Disabling the global variable immediately bypasses the snapshot and stops - scheduling new work. - -Until step 6 first succeeds, selection behaves exactly as it did before the -feature. - -## Verification strategy - -### Configuration and arithmetic unit tests - -- valid policy with defaults and explicit timing values; -- missing fields, wrong JSON types, NaN/infinity-equivalent rejection, - multiplier bounds and ordering; -- timing minimum, maximum, and `refresh <= stale` relationship; -- invalid reload removes prior locality influence; -- `1.0` and `10.0` multiplier boundaries; -- integer truncation, zero weight, non-cumulative tiers, and saturation; -- proof that configured/runtime table weights and checksums do not change. - -### Classification and discovery unit tests - -- IMDSv2 success and every failure phase; -- environment fallback precedence and partial values; -- missing account, matching account, and cross-account AZ-name collision; -- RDS instance and Aurora instance endpoints; -- cluster writer, reader, and custom endpoints; -- Multi-AZ/failover metadata refresh; -- remote Region and unknown local Region; -- custom CNAME, arbitrary host, unsupported RDS Proxy, and false AWS-looking - endpoint; -- exact normalized API endpoint match and port validation; -- paginated regional responses and duplicated endpoints across hostgroups; -- fixed/redacted errors without secrets or account IDs. - -### Cache and lifecycle unit tests - -- pending to fresh, fresh to stale, stale to expired, and recovery transitions - under a fake clock; -- different refresh/TTL policies sharing one endpoint; -- regional request coalescing and bounded queues; -- configuration reload, server removal, late completion, cancellation, and - generation rejection; -- provider replacement, plugin stop, core shutdown, and callback lifetime; -- enable, disable, and re-enable behavior; -- missing plugin and provider-unavailable neutral fallback; -- TSan coverage for publication, callbacks, reload, and shutdown. - -### Selection tests - -- deterministic effective-weight selection for the global Hostgroup Manager - path; -- deterministic parent-server weighting in the thread-local connection cache; -- proof that multiple cached connections do not amplify a server's weight; -- configured-weight relative ratios within each locality tier; -- no multiplier for unknown or expired metadata; -- cluster/reader/custom endpoints receive only same-Region preference; -- existing health, status, latency, lag, GTID, capacity, and backoff filters win - before locality; -- no locality-specific allocation, lock, plugin call, DNS, or network operation - in the hot path; -- current fast path remains in use when the feature is inactive. - -### Integration and regression tests - -- a fake asynchronous AWS provider drives the real MySQL Hostgroup Manager and - produces deterministic distributions; -- policy reload affects future selections only; -- `mysql_servers`, `runtime_mysql_servers`, saved configuration, and cluster - checksums remain byte-for-byte unchanged by discovered metadata; -- `stats_mysql_aws_locality` is absent without the AWS plugin and is registered - only when that plugin loads successfully; -- each table query projects one consistent in-memory snapshot without issuing - an IMDS or AWS API request, and projected rows are never persisted or - clustered; -- exact `stats_mysql_aws_locality` rows for fresh, stale, expired, error, and - disabled states; -- existing IAM database-authentication behavior continues through the shared - AWS plugin runtime; -- ASan, TSan, plugin lifecycle, static-linkage, SDK-free daemon, and existing - IAM/pool selection regressions remain green; -- optional externally provisioned AWS integration verifies one RDS/Aurora - instance endpoint and one cluster or reader endpoint without becoming a - normal-CI requirement. - -## Acceptance criteria - -The feature is complete when all of the following are true: - -1. The global switch defaults off and inactive builds preserve the existing - selection fast paths. -2. No locality operation mutates configured/runtime weights or cluster-visible - state. -3. Valid hostgroups apply bounded, non-cumulative Region/AZ multipliers only at - selection time. -4. The global pool and thread-local cache use identical server-level effective - weight semantics. -5. Instance endpoints can receive same-AZ preference; cluster, reader, and - custom endpoints cannot. -6. Same-AZ preference requires a confirmed same-account identity. -7. AWS/IMDS work is asynchronous, bounded, cancellable, and absent from the - hot path. -8. Refresh failure retains last-known metadata only through the configured - stale TTL, then returns to configured weighting. -9. Missing capability, credentials, permissions, or metadata never prevents a - database connection solely because locality awareness is enabled. -10. The plugin-conditional, query-refreshed diagnostic table explains every - active or neutral decision without performing network discovery or - exposing account IDs or sensitive AWS data. -11. Sanitizer, lifecycle, linkage, existing IAM, and selection regression gates - pass. diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp index 9253fed6f6..ee3f417736 100644 --- a/lib/Aws_Locality_Manager.cpp +++ b/lib/Aws_Locality_Manager.cpp @@ -1022,6 +1022,18 @@ class MySQLAwsLocalityManager::Impl { return entry; } + // A missing provider is different from a transient provider error: no + // concrete locality authority remains installed. Do not continue to + // apply a multiplier derived from a previous provider's cached result. + const bool provider_unavailable = local_.error == "provider_unavailable" || + (endpoint != endpoint_cache_.end() && + endpoint->second.error == "provider_unavailable"); + if (provider_unavailable) { + entry.status = AwsLocalityMetadataStatus::error; + entry.failure_category = "provider_unavailable"; + return entry; + } + const bool endpoint_pending = region_in_flight_.find( backend_config.endpoint.region) != region_in_flight_.end(); if (!local_.has_value || endpoint == endpoint_cache_.end() || diff --git a/lib/ProxySQL_Admin.cpp b/lib/ProxySQL_Admin.cpp index e851a3253d..8d519075a7 100644 --- a/lib/ProxySQL_Admin.cpp +++ b/lib/ProxySQL_Admin.cpp @@ -1613,11 +1613,11 @@ bool ProxySQL_Admin::GenericRefreshStatistics(const char *query_no_space, unsign } #ifdef PROXYSQL40 // Plugin-registered runtime views: if the query references any chassis- -// registered runtime view (e.g. runtime_mysqlx_users or an on-demand stats -// table), refresh it BEFORE the SELECT runs. Admin sessions can project all -// three DB kinds; stats sessions receive only the stats handle, so they cannot -// trigger an admin/config projection. The chassis decides whether to fire any -// plugin's refresh + // registered runtime view (e.g. runtime_mysqlx_users or an on-demand stats + // table), refresh it BEFORE the SELECT runs. Admin sessions can project all + // three DB kinds; stats sessions receive only the stats handle, so they cannot + // trigger an admin/config projection. The chassis decides whether to fire any + // plugin's refresh // callback by per-view substring match against query_no_space, so a // query that touches no registered view is a cheap no-op (one shared // lock + N substring scans, N == registered-view count). diff --git a/plugins/aws/src/aws_locality_provider.cpp b/plugins/aws/src/aws_locality_provider.cpp deleted file mode 100644 index 24649d7729..0000000000 --- a/plugins/aws/src/aws_locality_provider.cpp +++ /dev/null @@ -1,789 +0,0 @@ -#include "aws_locality_provider.h" - -#include "json.hpp" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef PROXYSQL_AWS_SDK_PROVIDER -#include "curl/curl.h" - -#include -#include -#include -#include -#include -#include -#include -#endif - -using nlohmann::json; - -namespace { - -const char* fixed_failure_category(AwsMetadataStatus status) { - switch (status) { - case AwsMetadataStatus::ok: return ""; - case AwsMetadataStatus::provider_unavailable: return "provider_unavailable"; - case AwsMetadataStatus::access_denied: return "access_denied"; - case AwsMetadataStatus::throttled: return "throttled"; - case AwsMetadataStatus::imds_unavailable: return "imds_unavailable"; - case AwsMetadataStatus::timeout: return "timeout"; - case AwsMetadataStatus::cancelled: return "cancelled"; - case AwsMetadataStatus::invalid_response: return "invalid_response"; - case AwsMetadataStatus::shutdown: return "cancelled"; - } - return "invalid_response"; -} - -AwsMetadataResult fixed_failure(AwsMetadataStatus status) { - AwsMetadataResult result; - result.status = status; - result.failure_category = fixed_failure_category(status); - return result; -} - -void normalize_failure(AwsMetadataResult& result) { - result.failure_category = fixed_failure_category(result.status); - if (result.status != AwsMetadataStatus::ok) { - result.local = {}; - result.endpoints.clear(); - } -} - -bool valid_location_value(const std::string& value, size_t maximum) { - if (value.empty() || value.size() > maximum) return false; - for (const unsigned char character : value) { - if (!(std::isalnum(character) || character == '-')) return false; - } - return true; -} - -bool valid_account_id(const std::string& value) { - return value.size() == 12 && std::all_of(value.begin(), value.end(), - [](unsigned char character) { return std::isdigit(character); }); -} - -std::string account_from_rds_arn( - const std::string& arn, - const std::string& expected_region) { - size_t begin = 0; - std::string fields[6]; - for (size_t field = 0; field < 5; ++field) { - const size_t end = arn.find(':', begin); - if (end == std::string::npos) return {}; - fields[field] = arn.substr(begin, end - begin); - begin = end + 1; - } - fields[5] = arn.substr(begin); - if (fields[0] != "arn" || fields[2] != "rds" || - fields[3] != expected_region || !valid_account_id(fields[4])) { - return {}; - } - return fields[4]; -} - -AwsEndpointType endpoint_type(const std::string& input) { - std::string value; - value.reserve(input.size()); - for (const unsigned char character : input) { - value.push_back(static_cast(std::toupper(character))); - } - if (value == "WRITER") return AwsEndpointType::cluster; - if (value == "READER") return AwsEndpointType::reader; - if (value == "CUSTOM") return AwsEndpointType::custom; - return AwsEndpointType::unknown; -} - -bool expired( - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled, - AwsMetadataResult& result) { - if (cancelled()) { - result = fixed_failure(AwsMetadataStatus::cancelled); - return true; - } - if (std::chrono::steady_clock::now() >= deadline) { - result = fixed_failure(AwsMetadataStatus::timeout); - return true; - } - return false; -} - -} // namespace - -class AwsSdkMetadataProvider::Impl { -public: - struct Job { - AwsMetadataRequestHandle handle; - AwsMetadataRequest request; - std::weak_ptr sink; - std::atomic cancelled { false }; - }; - - Impl( - std::shared_ptr backend, - AwsMetadataProviderConfig config) - : backend_(std::move(backend)), config_(std::move(config)) { - if (config_.worker_count == 0) config_.worker_count = 1; - if (config_.worker_count > 2) config_.worker_count = 2; - if (config_.max_pending < config_.worker_count) { - config_.max_pending = config_.worker_count; - } - workers_.reserve(config_.worker_count); - for (size_t i = 0; i < config_.worker_count; ++i) { - workers_.emplace_back([this] { worker_loop(); }); - } - } - - ~Impl() { shutdown(); } - - AwsMetadataRequestHandle request( - const AwsMetadataRequest& request, - std::weak_ptr sink) { - AwsMetadataResult immediate; - bool deliver = false; - { - std::lock_guard lock(mutex_); - if (stopping_ || backend_ == nullptr) return {}; - if (config_.steady_clock() >= request.deadline) { - immediate = fixed_failure(AwsMetadataStatus::timeout); - deliver = true; - } else if (jobs_.size() >= config_.max_pending) { - immediate = fixed_failure(AwsMetadataStatus::throttled); - deliver = true; - } else { - auto job = std::make_shared(); - job->handle.value = next_handle_++; - job->request = request; - job->sink = std::move(sink); - jobs_.emplace(job->handle.value, job); - queue_.push_back(job); - cv_.notify_one(); - return job->handle; - } - } - if (deliver) deliver_immediate(request, std::move(sink), std::move(immediate)); - return {}; - } - - void cancel(AwsMetadataRequestHandle handle) { - if (handle.value == 0) return; - std::lock_guard lock(mutex_); - const auto found = jobs_.find(handle.value); - if (found != jobs_.end()) found->second->cancelled.store(true); - cv_.notify_all(); - } - - void shutdown() { - { - std::unique_lock lock(mutex_); - if (shutdown_complete_) return; - if (shutdown_started_) { - cv_.wait(lock, [&] { return shutdown_complete_; }); - return; - } - shutdown_started_ = true; - stopping_ = true; - for (const auto& item : jobs_) item.second->cancelled.store(true); - cv_.notify_all(); - } - for (auto& worker : workers_) { - if (worker.joinable()) worker.join(); - } - { - std::unique_lock lock(mutex_); - cv_.wait(lock, [&] { return active_callbacks_ == 0; }); - jobs_.clear(); - queue_.clear(); - shutdown_complete_ = true; - cv_.notify_all(); - } - } - -private: - void deliver_immediate( - const AwsMetadataRequest& request, - std::weak_ptr weak_sink, - AwsMetadataResult result) { - auto sink = weak_sink.lock(); - if (!sink) return; - { - std::lock_guard lock(mutex_); - if (stopping_) return; - ++active_callbacks_; - } - AwsMetadataCompletion completion; - completion.opaque_id = request.opaque_id; - completion.generation = request.generation; - completion.result = std::move(result); - try { - sink->post(std::move(completion)); - } catch (...) { - // Plugin callbacks must not escape across the provider ABI boundary. - } - { - std::lock_guard lock(mutex_); - --active_callbacks_; - cv_.notify_all(); - } - } - - void worker_loop() { - for (;;) { - std::shared_ptr job; - { - std::unique_lock lock(mutex_); - cv_.wait(lock, [&] { return stopping_ || !queue_.empty(); }); - if (stopping_ && queue_.empty()) return; - job = queue_.front(); - queue_.pop_front(); - } - - AwsMetadataResult result; - if (job->cancelled.load()) { - result = fixed_failure(AwsMetadataStatus::cancelled); - } else if (config_.steady_clock() >= job->request.deadline) { - result = fixed_failure(AwsMetadataStatus::timeout); - } else { - try { - result = backend_->discover(job->request, - [job, this] { - return job->cancelled.load() || stopping_.load(); - }); - } catch (...) { - result = fixed_failure(AwsMetadataStatus::provider_unavailable); - } - if (config_.steady_clock() >= job->request.deadline) { - result = fixed_failure(AwsMetadataStatus::timeout); - } - } - normalize_failure(result); - - std::shared_ptr sink; - { - std::lock_guard lock(mutex_); - jobs_.erase(job->handle.value); - if (!stopping_ && !job->cancelled.load()) { - sink = job->sink.lock(); - if (sink) ++active_callbacks_; - } - } - if (sink) { - AwsMetadataCompletion completion; - completion.opaque_id = job->request.opaque_id; - completion.generation = job->request.generation; - completion.result = std::move(result); - try { - sink->post(std::move(completion)); - } catch (...) { - // Keep worker and shutdown bookkeeping intact on a bad consumer. - } - std::lock_guard lock(mutex_); - --active_callbacks_; - cv_.notify_all(); - } - } - } - - std::shared_ptr backend_; - AwsMetadataProviderConfig config_; - std::mutex mutex_; - std::condition_variable cv_; - std::deque> queue_; - std::unordered_map> jobs_; - std::vector workers_; - std::atomic stopping_ { false }; - bool shutdown_started_ { false }; - bool shutdown_complete_ { false }; - uint64_t next_handle_ { 1 }; - size_t active_callbacks_ { 0 }; -}; - -AwsSdkMetadataProvider::AwsSdkMetadataProvider( - std::shared_ptr backend, - AwsMetadataProviderConfig config) - : impl_(new Impl(std::move(backend), std::move(config))) {} - -AwsSdkMetadataProvider::~AwsSdkMetadataProvider() = default; - -AwsMetadataRequestHandle AwsSdkMetadataProvider::request( - const AwsMetadataRequest& request, - std::weak_ptr sink) { - return impl_->request(request, std::move(sink)); -} - -void AwsSdkMetadataProvider::cancel(AwsMetadataRequestHandle handle) { - impl_->cancel(handle); -} - -void AwsSdkMetadataProvider::shutdown() { - impl_->shutdown(); -} - -AwsLocalLocation aws_locality_environment_location( - const AwsLocalityEnvironmentGetter& getenv_value) { - AwsLocalLocation location; - if (!getenv_value) return location; - location.region = getenv_value("AWS_REGION"); - if (!valid_location_value(location.region, 64)) { - location.region = getenv_value("AWS_DEFAULT_REGION"); - } - if (!valid_location_value(location.region, 64)) { - location.region.clear(); - return location; - } - location.availability_zone = getenv_value("AWS_AVAILABILITY_ZONE"); - if (!valid_location_value(location.availability_zone, 64)) { - location.availability_zone.clear(); - } - location.account_id = getenv_value("AWS_ACCOUNT_ID"); - if (!valid_account_id(location.account_id)) location.account_id.clear(); - return location; -} - -AwsLocalityLocalDiscovery::AwsLocalityLocalDiscovery( - std::shared_ptr transport, - AwsLocalityEnvironmentGetter getenv_value) - : transport_(std::move(transport)), getenv_value_(std::move(getenv_value)) {} - -AwsMetadataResult AwsLocalityLocalDiscovery::discover( - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) const { - AwsMetadataResult result; - if (expired(deadline, cancelled, result)) return result; - - auto fallback = [&](AwsMetadataStatus failure_status) { - AwsMetadataResult fallback_result; - fallback_result.local = aws_locality_environment_location(getenv_value_); - if (!fallback_result.local.region.empty()) { - fallback_result.status = AwsMetadataStatus::ok; - return fallback_result; - } - return fixed_failure(failure_status); - }; - - if (!transport_) return fallback(AwsMetadataStatus::imds_unavailable); - AwsImdsResponse token = transport_->put_token(deadline, cancelled); - if (expired(deadline, cancelled, result)) return result; - if (!token.transport_ok || token.status_code != 200 || token.body.empty() || - token.body.size() > 4096) { - if (!token.body.empty()) OPENSSL_cleanse(&token.body[0], token.body.size()); - return fallback(AwsMetadataStatus::imds_unavailable); - } - - AwsImdsResponse document = transport_->get_identity_document( - token.body, deadline, cancelled); - OPENSSL_cleanse(&token.body[0], token.body.size()); - if (expired(deadline, cancelled, result)) return result; - if (!document.transport_ok || document.status_code != 200) { - return fallback(AwsMetadataStatus::imds_unavailable); - } - if (document.body.empty() || document.body.size() > 16384) { - return fallback(AwsMetadataStatus::invalid_response); - } - - try { - const json identity = json::parse(document.body); - if (!identity.is_object() || !identity.contains("region") || - !identity["region"].is_string()) { - return fallback(AwsMetadataStatus::invalid_response); - } - result.local.region = identity["region"].get(); - if (!valid_location_value(result.local.region, 64)) { - return fallback(AwsMetadataStatus::invalid_response); - } - if (identity.contains("availabilityZone") && - identity["availabilityZone"].is_string()) { - result.local.availability_zone = - identity["availabilityZone"].get(); - if (!valid_location_value(result.local.availability_zone, 64)) { - result.local.availability_zone.clear(); - } - } - if (identity.contains("accountId") && identity["accountId"].is_string()) { - result.local.account_id = identity["accountId"].get(); - if (!valid_account_id(result.local.account_id)) result.local.account_id.clear(); - } - result.status = AwsMetadataStatus::ok; - return result; - } catch (...) { - return fallback(AwsMetadataStatus::invalid_response); - } -} - -AwsLocalityRdsDiscovery::AwsLocalityRdsDiscovery( - std::shared_ptr api) - : api_(std::move(api)) {} - -AwsMetadataResult AwsLocalityRdsDiscovery::discover( - const AwsMetadataRequest& request, - const AwsLocalityCancelPredicate& cancelled) const { - AwsMetadataResult result; - if (request.region.empty() || !api_) { - return fixed_failure(AwsMetadataStatus::invalid_response); - } - if (expired(request.deadline, cancelled, result)) return result; - result.status = AwsMetadataStatus::ok; - - std::unordered_map endpoint_indices; - auto append = [&](const std::string& hostname_input, int port, - AwsEndpointType type, const std::string& az, const std::string& account) { - const std::string hostname = aws_locality_normalized_hostname(hostname_input); - if (hostname.empty() || type == AwsEndpointType::unknown || - port < 0 || port > 65535) return; - AwsMetadataEndpoint endpoint; - endpoint.hostname = hostname; - endpoint.port = static_cast(port); - endpoint.endpoint_type = type; - endpoint.region = request.region; - endpoint.availability_zone = az; - endpoint.account_id = account; - const std::string key = hostname + "\n" + std::to_string(port); - const auto found = endpoint_indices.find(key); - if (found == endpoint_indices.end()) { - endpoint_indices.emplace(key, result.endpoints.size()); - result.endpoints.push_back(std::move(endpoint)); - } else { - result.endpoints[found->second] = std::move(endpoint); - } - }; - - auto fail = [&](AwsMetadataStatus status) { - result = fixed_failure(status); - return result; - }; - auto check_page = [&](AwsMetadataStatus status) { - if (cancelled()) return AwsMetadataStatus::cancelled; - if (std::chrono::steady_clock::now() >= request.deadline) { - return AwsMetadataStatus::timeout; - } - return status; - }; - - std::set markers; - std::string marker; - for (;;) { - AwsRdsInstancesPage page = api_->describe_instances( - request.region, marker, request.deadline, cancelled); - const AwsMetadataStatus status = check_page(page.status); - if (status != AwsMetadataStatus::ok) return fail(status); - for (const auto& instance : page.instances) { - if (instance.port <= 0 || instance.port > 65535) continue; - append(instance.endpoint, instance.port, AwsEndpointType::instance, - instance.availability_zone, - account_from_rds_arn(instance.arn, request.region)); - } - if (page.next_marker.empty()) break; - if (!markers.insert(page.next_marker).second) { - return fail(AwsMetadataStatus::invalid_response); - } - marker = std::move(page.next_marker); - } - - std::unordered_map cluster_accounts; - markers.clear(); - marker.clear(); - for (;;) { - AwsRdsClustersPage page = api_->describe_clusters( - request.region, marker, request.deadline, cancelled); - const AwsMetadataStatus status = check_page(page.status); - if (status != AwsMetadataStatus::ok) return fail(status); - for (const auto& cluster : page.clusters) { - const std::string account = account_from_rds_arn(cluster.arn, request.region); - if (!cluster.identifier.empty()) cluster_accounts[cluster.identifier] = account; - if (cluster.port > 0 && cluster.port <= 65535) { - append(cluster.endpoint, cluster.port, AwsEndpointType::cluster, {}, account); - append(cluster.reader_endpoint, cluster.port, AwsEndpointType::reader, {}, account); - } - for (const auto& custom : cluster.custom_endpoints) { - append(custom, 0, AwsEndpointType::custom, {}, account); - } - } - if (page.next_marker.empty()) break; - if (!markers.insert(page.next_marker).second) { - return fail(AwsMetadataStatus::invalid_response); - } - marker = std::move(page.next_marker); - } - - markers.clear(); - marker.clear(); - for (;;) { - AwsRdsClusterEndpointsPage page = api_->describe_cluster_endpoints( - request.region, marker, request.deadline, cancelled); - const AwsMetadataStatus status = check_page(page.status); - if (status != AwsMetadataStatus::ok) return fail(status); - for (const auto& endpoint : page.endpoints) { - const auto account = cluster_accounts.find(endpoint.cluster_identifier); - append(endpoint.endpoint, 0, endpoint_type(endpoint.endpoint_type), {}, - account == cluster_accounts.end() ? std::string() : account->second); - } - if (page.next_marker.empty()) break; - if (!markers.insert(page.next_marker).second) { - return fail(AwsMetadataStatus::invalid_response); - } - marker = std::move(page.next_marker); - } - - return result; -} - -AwsLocalityCompositeDiscovery::AwsLocalityCompositeDiscovery( - std::shared_ptr imds, - AwsLocalityEnvironmentGetter getenv_value, - std::shared_ptr rds) - : local_(std::move(imds), std::move(getenv_value)), - rds_(std::move(rds)) {} - -AwsMetadataResult AwsLocalityCompositeDiscovery::discover( - const AwsMetadataRequest& request, - const AwsLocalityCancelPredicate& cancelled) { - if (request.kind == AwsMetadataRequestKind::local_location) { - return local_.discover(request.deadline, cancelled); - } - if (request.kind == AwsMetadataRequestKind::rds_region) { - return rds_.discover(request, cancelled); - } - return fixed_failure(AwsMetadataStatus::invalid_response); -} - -#ifdef PROXYSQL_AWS_SDK_PROVIDER -namespace { - -struct CurlResponseContext { - std::string* body; - size_t maximum; - bool overflow { false }; -}; - -size_t imds_write_callback(char* data, size_t size, size_t count, void* opaque) { - auto* context = static_cast(opaque); - if (size != 0 && count > context->maximum / size) { - context->overflow = true; - return 0; - } - const size_t bytes = size * count; - if (bytes > context->maximum - std::min(context->maximum, context->body->size())) { - context->overflow = true; - return 0; - } - context->body->append(data, bytes); - return bytes; -} - -struct CurlProgressContext { - std::chrono::steady_clock::time_point deadline; - const AwsLocalityCancelPredicate* cancelled; -}; - -int imds_progress_callback(void* opaque, curl_off_t, curl_off_t, curl_off_t, curl_off_t) { - auto* context = static_cast(opaque); - return (*(context->cancelled))() || - std::chrono::steady_clock::now() >= context->deadline; -} - -AwsImdsResponse imds_request( - const char* url, - const char* method, - struct curl_slist* headers, - size_t maximum, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) { - AwsImdsResponse response; - if (cancelled() || std::chrono::steady_clock::now() >= deadline) return response; - CURL* handle = curl_easy_init(); // NOSONAR(cpp:S4423): IMDSv2 is HTTP-only; this handle is restricted to HTTP below. - if (handle == nullptr) return response; - CurlResponseContext write_context {&response.body, maximum}; - CurlProgressContext progress_context {deadline, &cancelled}; - const auto remaining = std::chrono::duration_cast( - deadline - std::chrono::steady_clock::now()).count(); - const long timeout = static_cast(std::max(1, remaining)); - curl_easy_setopt(handle, CURLOPT_URL, url); - curl_easy_setopt(handle, CURLOPT_CUSTOMREQUEST, method); - curl_easy_setopt(handle, CURLOPT_HTTPHEADER, headers); - curl_easy_setopt(handle, CURLOPT_NOBODY, 0L); - curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 0L); - curl_easy_setopt(handle, CURLOPT_NOPROXY, "*"); - // IMDS is an HTTP-only link-local service. Restrict this handle to that - // protocol; TLS options are intentionally inapplicable to this transport. - curl_easy_setopt(handle, CURLOPT_PROTOCOLS_STR, "http"); - curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, std::min(timeout, 500L)); - curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, timeout); - curl_easy_setopt(handle, CURLOPT_NOSIGNAL, 1L); - curl_easy_setopt(handle, CURLOPT_WRITEFUNCTION, &imds_write_callback); - curl_easy_setopt(handle, CURLOPT_WRITEDATA, &write_context); - curl_easy_setopt(handle, CURLOPT_XFERINFOFUNCTION, &imds_progress_callback); - curl_easy_setopt(handle, CURLOPT_XFERINFODATA, &progress_context); - curl_easy_setopt(handle, CURLOPT_NOPROGRESS, 0L); - const CURLcode code = curl_easy_perform(handle); - curl_easy_getinfo(handle, CURLINFO_RESPONSE_CODE, &response.status_code); - curl_easy_cleanup(handle); - response.transport_ok = code == CURLE_OK && !write_context.overflow; - if (!response.transport_ok) response.body.clear(); - return response; -} - -template -AwsMetadataStatus map_sdk_error(const Error& error) { - const int type = static_cast(error.GetErrorType()); - const int code = static_cast(error.GetResponseCode()); - if (type == static_cast(Aws::Client::CoreErrors::ACCESS_DENIED) || code == 401 || code == 403) { - return AwsMetadataStatus::access_denied; - } - if (type == static_cast(Aws::Client::CoreErrors::THROTTLING) || - type == static_cast(Aws::Client::CoreErrors::SLOW_DOWN) || code == 429) { - return AwsMetadataStatus::throttled; - } - if (type == static_cast(Aws::Client::CoreErrors::REQUEST_TIMEOUT) || - type == static_cast(Aws::Client::CoreErrors::NETWORK_CONNECTION)) { - return AwsMetadataStatus::timeout; - } - return AwsMetadataStatus::provider_unavailable; -} - -} // namespace - -AwsImdsResponse AwsCurlImdsTransport::put_token( - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) { - struct curl_slist* headers = nullptr; - headers = curl_slist_append(headers, "X-aws-ec2-metadata-token-ttl-seconds: 21600"); - AwsImdsResponse response = imds_request( - "http://169.254.169.254/latest/api/token", "PUT", headers, - 4096, deadline, cancelled); - curl_slist_free_all(headers); - return response; -} - -AwsImdsResponse AwsCurlImdsTransport::get_identity_document( - const std::string& token, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) { - struct curl_slist* headers = nullptr; - std::string token_header = "X-aws-ec2-metadata-token: " + token; - headers = curl_slist_append(headers, token_header.c_str()); - AwsImdsResponse response = imds_request( - "http://169.254.169.254/latest/dynamic/instance-identity/document", - "GET", headers, 16384, deadline, cancelled); - curl_slist_free_all(headers); - if (!token_header.empty()) { - OPENSSL_cleanse(&token_header[0], token_header.size()); - } - return response; -} - -std::shared_ptr AwsSdkRdsDiscoveryApi::client_for_region( - const std::string& region) { - std::lock_guard lock(clients_mutex_); - const auto found = clients_.find(region); - if (found != clients_.end()) return found->second; - Aws::Client::ClientConfiguration config; - config.region = region.c_str(); - config.connectTimeoutMs = 500; - config.requestTimeoutMs = 4000; - config.httpRequestTimeoutMs = 4000; - config.retryStrategy = std::make_shared(2, 50); - auto client = std::make_shared(config); - clients_.emplace(region, client); - return client; -} - -AwsRdsInstancesPage AwsSdkRdsDiscoveryApi::describe_instances( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) { - AwsRdsInstancesPage page; - if (cancelled()) { page.status = AwsMetadataStatus::cancelled; return page; } - if (std::chrono::steady_clock::now() >= deadline) { - page.status = AwsMetadataStatus::timeout; - return page; - } - Aws::RDS::Model::DescribeDBInstancesRequest request; - if (!marker.empty()) request.SetMarker(marker.c_str()); - const auto outcome = client_for_region(region)->DescribeDBInstances(request); - if (!outcome.IsSuccess()) { - page.status = map_sdk_error(outcome.GetError()); - return page; - } - page.status = AwsMetadataStatus::ok; - page.next_marker = outcome.GetResult().GetMarker().c_str(); - for (const auto& instance : outcome.GetResult().GetDBInstances()) { - const auto& endpoint = instance.GetEndpoint(); - page.instances.push_back({endpoint.GetAddress().c_str(), endpoint.GetPort(), - instance.GetAvailabilityZone().c_str(), instance.GetDBInstanceArn().c_str()}); - } - return page; -} - -AwsRdsClustersPage AwsSdkRdsDiscoveryApi::describe_clusters( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) { - AwsRdsClustersPage page; - if (cancelled()) { page.status = AwsMetadataStatus::cancelled; return page; } - if (std::chrono::steady_clock::now() >= deadline) { - page.status = AwsMetadataStatus::timeout; - return page; - } - Aws::RDS::Model::DescribeDBClustersRequest request; - if (!marker.empty()) request.SetMarker(marker.c_str()); - const auto outcome = client_for_region(region)->DescribeDBClusters(request); - if (!outcome.IsSuccess()) { - page.status = map_sdk_error(outcome.GetError()); - return page; - } - page.status = AwsMetadataStatus::ok; - page.next_marker = outcome.GetResult().GetMarker().c_str(); - for (const auto& cluster : outcome.GetResult().GetDBClusters()) { - std::vector custom; - for (const auto& endpoint : cluster.GetCustomEndpoints()) { - custom.emplace_back(endpoint.c_str()); - } - page.clusters.push_back({cluster.GetDBClusterIdentifier().c_str(), - cluster.GetEndpoint().c_str(), cluster.GetReaderEndpoint().c_str(), - cluster.GetPort(), std::move(custom), cluster.GetDBClusterArn().c_str()}); - } - return page; -} - -AwsRdsClusterEndpointsPage AwsSdkRdsDiscoveryApi::describe_cluster_endpoints( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) { - AwsRdsClusterEndpointsPage page; - if (cancelled()) { page.status = AwsMetadataStatus::cancelled; return page; } - if (std::chrono::steady_clock::now() >= deadline) { - page.status = AwsMetadataStatus::timeout; - return page; - } - Aws::RDS::Model::DescribeDBClusterEndpointsRequest request; - if (!marker.empty()) request.SetMarker(marker.c_str()); - const auto outcome = client_for_region(region)->DescribeDBClusterEndpoints(request); - if (!outcome.IsSuccess()) { - page.status = map_sdk_error(outcome.GetError()); - return page; - } - page.status = AwsMetadataStatus::ok; - page.next_marker = outcome.GetResult().GetMarker().c_str(); - for (const auto& endpoint : outcome.GetResult().GetDBClusterEndpoints()) { - page.endpoints.push_back({endpoint.GetEndpoint().c_str(), - endpoint.GetEndpointType().c_str(), endpoint.GetDBClusterIdentifier().c_str()}); - } - return page; -} -#endif // PROXYSQL_AWS_SDK_PROVIDER diff --git a/plugins/aws/src/aws_locality_provider.h b/plugins/aws/src/aws_locality_provider.h deleted file mode 100644 index 8d6643c848..0000000000 --- a/plugins/aws/src/aws_locality_provider.h +++ /dev/null @@ -1,218 +0,0 @@ -#ifndef PROXYSQL_AWS_LOCALITY_PROVIDER_H -#define PROXYSQL_AWS_LOCALITY_PROVIDER_H - -#include "Aws_Locality_Types.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace Aws { namespace RDS { class RDSClient; } } - -using AwsLocalityCancelPredicate = std::function; -using AwsLocalityEnvironmentGetter = std::function; - -struct AwsMetadataProviderConfig { - using SteadyClock = std::function; - - size_t worker_count { 2 }; - size_t max_pending { 256 }; - SteadyClock steady_clock { - [] { return std::chrono::steady_clock::now(); } - }; -}; - -class AwsLocalityDiscoveryBackend { -public: - virtual AwsMetadataResult discover( - const AwsMetadataRequest& request, - const AwsLocalityCancelPredicate& cancelled) = 0; - virtual ~AwsLocalityDiscoveryBackend() = default; -}; - -class AwsSdkMetadataProvider final : public AwsMetadataProvider { -public: - explicit AwsSdkMetadataProvider( - std::shared_ptr backend, - AwsMetadataProviderConfig config = {}); - ~AwsSdkMetadataProvider() override; - - AwsMetadataRequestHandle request( - const AwsMetadataRequest& request, - std::weak_ptr sink) override; - void cancel(AwsMetadataRequestHandle handle) override; - void shutdown() override; - - AwsSdkMetadataProvider(const AwsSdkMetadataProvider&) = delete; - AwsSdkMetadataProvider& operator=(const AwsSdkMetadataProvider&) = delete; - -private: - class Impl; - std::unique_ptr impl_; -}; - -struct AwsImdsResponse { - bool transport_ok { false }; - long status_code { 0 }; - std::string body; -}; - -class AwsImdsTransport { -public: - virtual AwsImdsResponse put_token( - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) = 0; - virtual AwsImdsResponse get_identity_document( - const std::string& token, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) = 0; - virtual ~AwsImdsTransport() = default; -}; - -AwsLocalLocation aws_locality_environment_location( - const AwsLocalityEnvironmentGetter& getenv_value); - -class AwsLocalityLocalDiscovery { -public: - AwsLocalityLocalDiscovery( - std::shared_ptr transport, - AwsLocalityEnvironmentGetter getenv_value); - - AwsMetadataResult discover( - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) const; - -private: - std::shared_ptr transport_; - AwsLocalityEnvironmentGetter getenv_value_; -}; - -struct AwsRdsInstanceRecord { - std::string endpoint; - int port { 0 }; - std::string availability_zone; - std::string arn; -}; - -struct AwsRdsClusterRecord { - std::string identifier; - std::string endpoint; - std::string reader_endpoint; - int port { 0 }; - std::vector custom_endpoints; - std::string arn; -}; - -struct AwsRdsClusterEndpointRecord { - std::string endpoint; - std::string endpoint_type; - std::string cluster_identifier; -}; - -struct AwsRdsInstancesPage { - AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; - std::vector instances; - std::string next_marker; -}; - -struct AwsRdsClustersPage { - AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; - std::vector clusters; - std::string next_marker; -}; - -struct AwsRdsClusterEndpointsPage { - AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; - std::vector endpoints; - std::string next_marker; -}; - -class AwsRdsDiscoveryApi { -public: - virtual AwsRdsInstancesPage describe_instances( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) = 0; - virtual AwsRdsClustersPage describe_clusters( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) = 0; - virtual AwsRdsClusterEndpointsPage describe_cluster_endpoints( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) = 0; - virtual ~AwsRdsDiscoveryApi() = default; -}; - -class AwsLocalityRdsDiscovery { -public: - explicit AwsLocalityRdsDiscovery(std::shared_ptr api); - - AwsMetadataResult discover( - const AwsMetadataRequest& request, - const AwsLocalityCancelPredicate& cancelled) const; - -private: - std::shared_ptr api_; -}; - -class AwsCurlImdsTransport final : public AwsImdsTransport { -public: - AwsImdsResponse put_token( - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) override; - AwsImdsResponse get_identity_document( - const std::string& token, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) override; -}; - -class AwsSdkRdsDiscoveryApi final : public AwsRdsDiscoveryApi { -public: - AwsRdsInstancesPage describe_instances( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) override; - AwsRdsClustersPage describe_clusters( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) override; - AwsRdsClusterEndpointsPage describe_cluster_endpoints( - const std::string& region, - const std::string& marker, - std::chrono::steady_clock::time_point deadline, - const AwsLocalityCancelPredicate& cancelled) override; - -private: - std::shared_ptr client_for_region(const std::string& region); - std::mutex clients_mutex_; - std::unordered_map> clients_; -}; - -class AwsLocalityCompositeDiscovery final : public AwsLocalityDiscoveryBackend { -public: - AwsLocalityCompositeDiscovery( - std::shared_ptr imds, - AwsLocalityEnvironmentGetter getenv_value, - std::shared_ptr rds); - - AwsMetadataResult discover( - const AwsMetadataRequest& request, - const AwsLocalityCancelPredicate& cancelled) override; - -private: - AwsLocalityLocalDiscovery local_; - AwsLocalityRdsDiscovery rds_; -}; - -#endif // PROXYSQL_AWS_LOCALITY_PROVIDER_H diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index 95534ba10a..e313f3f2b5 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -16,14 +16,17 @@ # cannot terminate the producer early under pipefail. .PHONY: aws_public_provider_boundary-t aws_public_provider_boundary-t: + @test ! -d $(PROXYSQL_PATH)/plugins/aws || { \ + echo "FAIL: public tree owns the concrete AWS provider" >&2; exit 1; \ + } + @test ! -e $(PROXYSQL_PATH)/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp || { \ + echo "FAIL: public tree owns the real AWS locality provider test" >&2; exit 1; \ + } @test -x $(PROXYSQL_PATH)/src/proxysql @nm -C $(PROXYSQL_PATH)/src/proxysql > /tmp/proxysql-public.nm @if grep -Fq 'Aws::' /tmp/proxysql-public.nm; then \ echo "FAIL: src/proxysql contains AWS C++ SDK symbols" >&2; exit 1; \ fi - @test ! -d $(PROXYSQL_PATH)/plugins/aws || { \ - echo "FAIL: public tree owns the AWS plugin" >&2; exit 1; \ - } @test ! -e $(PROXYSQL_PATH)/plugins/aws/ProxySQL_Aws_Plugin.so || { \ echo "FAIL: public build produced the AWS plugin" >&2; exit 1; \ } @@ -950,13 +953,6 @@ aws_locality_manager_unit-t: aws_locality_manager_unit-t.cpp \ $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ $(IDIRS) $(OPT) -lpthread -ldl -o $@ -aws_locality_plugin_unit-t: aws_locality_plugin_unit-t.cpp \ - $(PROXYSQL_PATH)/plugins/aws/src/aws_locality_provider.cpp \ - $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o - $(CXX) $< $(PROXYSQL_PATH)/plugins/aws/src/aws_locality_provider.cpp \ - $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ - $(IDIRS) -I$(PROXYSQL_PATH)/plugins/aws/src $(OPT) -lpthread -lcrypto -o $@ - test_aws_iam_metrics-t: ../test_aws_iam_metrics-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) $(CXX) $< $(TEST_HELPERS_OBJ) -I$(TEST_HELPERS_DIR) \ $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ @@ -968,9 +964,8 @@ test_aws_iam_backend_auth-t: ../test_aws_iam_backend_auth-t.cpp $(TEST_HELPERS_O $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ -aws_locality_stats_unit-t: aws_locality_stats_unit-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) aws_plugin_build - $(CXX) -DPROXYSQL_AWS_PLUGIN_PATH=\"$(PROXYSQL_PATH)/plugins/aws/ProxySQL_Aws_Plugin.so\" \ - $< $(TEST_HELPERS_OBJ) $(IDIRS) $(LDIRS) $(OPT) \ +aws_locality_stats_unit-t: aws_locality_stats_unit-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) + $(CXX) $< $(TEST_HELPERS_OBJ) $(IDIRS) $(LDIRS) $(OPT) \ $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ ifeq ($(UNAME_S),Linux) diff --git a/test/tap/tests/unit/aws_locality_manager_unit-t.cpp b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp index 6ffc0d6176..f5fef8a6d8 100644 --- a/test/tap/tests/unit/aws_locality_manager_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp @@ -202,7 +202,7 @@ const AwsLocalitySnapshotEntry* lookup( } // namespace int main() { - plan(44); + plan(49); auto provider_state = std::make_shared(); ok(install_global_aws_metadata_provider( @@ -411,9 +411,10 @@ int main() { MySQLAwsLocalityManager absent_provider_manager(manager_config); const auto absent_endpoint = "db-absent.abcdefghijkl.us-east-1.rds.amazonaws.com"; - absent_provider_manager.configure({ - make_hostgroup(20, 2.0, 4.0, 300, 1800, {absent_endpoint}), - }); + auto absent_hostgroup = + make_hostgroup(20, 2.0, 4.0, 300, 1800, {absent_endpoint}); + absent_hostgroup.backends[0].configured_weight = 37; + absent_provider_manager.configure({std::move(absent_hostgroup)}); absent_provider_manager.set_enabled(true); ok(wait_until([&] { const auto current = absent_provider_manager.snapshot(); @@ -421,6 +422,12 @@ int main() { return entry && entry->status == AwsLocalityMetadataStatus::error && entry->failure_category == "provider_unavailable"; }), "missing plugin provider is reported with a fixed neutral category"); + const auto absent_snapshot = absent_provider_manager.snapshot(); + const auto* absent_entry = lookup(absent_snapshot, 20, absent_endpoint); + ok(absent_entry != nullptr && absent_entry->configured_weight == 37 && + absent_entry->multiplier == 1.0 && + absent_snapshot->effective_weight(20, absent_endpoint, 3306, 37) == 37, + "missing provider preserves configured and effective neutral weights"); auto replacement_state = std::make_shared(); ok(install_global_aws_metadata_provider( @@ -633,8 +640,38 @@ int main() { path_entry != nullptr && ipv6_entry->configured_weight == 7 && path_entry->configured_weight == 9, "rejected DNS spellings retain distinct snapshot identities"); + invalid_identity_manager.shutdown(); - shutdown_global_aws_metadata_provider(); + const size_t lease_drain_request_count = requests_copy(replacement_state).size(); + MySQLAwsLocalityManager lease_drain_manager(manager_config); + lease_drain_manager.configure({ + make_hostgroup(600, 2.0, 4.0, 300, 1800, + {"db-drain.abcdefghijkl.us-east-2.rds.amazonaws.com"}), + }); + lease_drain_manager.set_enabled(true); + ok(wait_for_request_count(replacement_state, lease_drain_request_count + 2), + "enabled manager retains the installed provider during discovery"); + std::atomic unload_done { false }; + std::thread unload_thread([&] { + shutdown_global_aws_metadata_provider(); + unload_done.store(true); + }); + ok(wait_until([&] { + return !acquire_global_aws_metadata_provider(); + }) && !unload_done.load(), + "provider unload rejects new leases while the manager lease is active"); + lease_drain_manager.set_enabled(false); + unload_thread.join(); + ok(unload_done.load() && replacement_state->shut_down && + replacement_state->destroyed, + "disabling locality drains the manager lease before provider destruction"); + const size_t requests_after_unload = requests_copy(replacement_state).size(); + lease_drain_manager.shutdown(); + lease_drain_manager.set_enabled(true); + lease_drain_manager.request_refresh(); + std::this_thread::sleep_for(20ms); + ok(requests_copy(replacement_state).size() == requests_after_unload, + "manager shutdown leaves no locality worker able to restart provider work"); return exit_status(); } diff --git a/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp b/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp deleted file mode 100644 index 18df1e6cc1..0000000000 --- a/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp +++ /dev/null @@ -1,409 +0,0 @@ -#include "tap.h" - -#include "aws_locality_provider.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace std::chrono_literals; - -namespace { - -class CapturingSink final : public AwsMetadataCompletionSink { -public: - void post(AwsMetadataCompletion&& completion) override { - std::lock_guard lock(mutex_); - completions_.push_back(std::move(completion)); - cv_.notify_all(); - } - - bool wait_for(size_t count) { - std::unique_lock lock(mutex_); - return cv_.wait_for(lock, 2s, [&] { return completions_.size() >= count; }); - } - - std::vector snapshot() const { - std::lock_guard lock(mutex_); - return completions_; - } - -private: - mutable std::mutex mutex_; - std::condition_variable cv_; - std::vector completions_; -}; - -class BlockingBackend final : public AwsLocalityDiscoveryBackend { -public: - AwsMetadataResult discover( - const AwsMetadataRequest& request, - const AwsLocalityCancelPredicate& cancelled) override { - const int active = active_.fetch_add(1) + 1; - int observed = max_active_.load(); - while (active > observed && - !max_active_.compare_exchange_weak(observed, active)) {} - { - std::unique_lock lock(mutex_); - started_.push_back(request.opaque_id); - cv_.notify_all(); - cv_.wait(lock, [&] { return released_ || cancelled(); }); - } - active_.fetch_sub(1); - AwsMetadataResult result; - result.status = cancelled() ? AwsMetadataStatus::cancelled : AwsMetadataStatus::ok; - result.failure_category = "FAKE_SECRET_RAW_ERROR"; - result.local.region = request.region; - return result; - } - - bool wait_started(size_t count) { - std::unique_lock lock(mutex_); - return cv_.wait_for(lock, 2s, [&] { return started_.size() >= count; }); - } - - void release() { - std::lock_guard lock(mutex_); - released_ = true; - cv_.notify_all(); - } - - int max_active() const { return max_active_.load(); } - -private: - std::atomic active_ { 0 }; - std::atomic max_active_ { 0 }; - std::mutex mutex_; - std::condition_variable cv_; - std::vector started_; - bool released_ { false }; -}; - -class FakeImdsTransport final : public AwsImdsTransport { -public: - AwsImdsResponse put_token( - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { - ++token_calls; - return token; - } - - AwsImdsResponse get_identity_document( - const std::string& supplied_token, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { - ++document_calls; - seen_token = supplied_token; - return document; - } - - AwsImdsResponse token { true, 200, "imds-token" }; - AwsImdsResponse document { true, 200, - R"({"region":"us-east-1","availabilityZone":"us-east-1b","accountId":"111122223333"})" }; - int token_calls { 0 }; - int document_calls { 0 }; - std::string seen_token; -}; - -class FakeRdsApi final : public AwsRdsDiscoveryApi { -public: - AwsRdsInstancesPage describe_instances( - const std::string&, const std::string& marker, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { - ++instance_calls; - if (marker.empty()) { - AwsRdsInstancesPage page; - page.status = AwsMetadataStatus::ok; - page.next_marker = "instances-2"; - page.instances.push_back({ - "DB-ONE.ABCDEFGHIJKL.US-EAST-1.RDS.AMAZONAWS.COM.", 3306, - "us-east-1a", "arn:aws:rds:us-east-1:111122223333:db:one"}); - page.instances.push_back({ - "missing-port.abcdefghijkl.us-east-1.rds.amazonaws.com", 0, - "us-east-1a", "arn:aws:rds:us-east-1:111122223333:db:missing"}); - return page; - } - AwsRdsInstancesPage page; - page.status = AwsMetadataStatus::ok; - page.instances.push_back({ - "db-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306, - "us-east-1d", "arn:aws:rds:us-east-1:111122223333:db:one"}); - page.instances.push_back({ - "db-two.abcdefghijkl.us-east-1.rds.amazonaws.com", 3307, - "us-east-1c", "arn:aws:rds:us-east-1:444455556666:db:two"}); - return page; - } - - AwsRdsClustersPage describe_clusters( - const std::string&, const std::string&, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { - ++cluster_calls; - AwsRdsClustersPage page; - page.status = AwsMetadataStatus::ok; - page.clusters.push_back({ - "cluster-one", "cluster-one.abcdefghijkl.us-east-1.rds.amazonaws.com", - "cluster-ro-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306, - {"custom-one.abcdefghijkl.us-east-1.rds.amazonaws.com"}, - "arn:aws:rds:us-east-1:111122223333:cluster:cluster-one"}); - return page; - } - - AwsRdsClusterEndpointsPage describe_cluster_endpoints( - const std::string&, const std::string&, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { - ++endpoint_calls; - AwsRdsClusterEndpointsPage page; - page.status = AwsMetadataStatus::ok; - page.endpoints.push_back({ - "custom-two.abcdefghijkl.us-east-1.rds.amazonaws.com", - "CUSTOM", "cluster-one"}); - return page; - } - - int instance_calls { 0 }; - int cluster_calls { 0 }; - int endpoint_calls { 0 }; -}; - -const AwsMetadataEndpoint* find_endpoint( - const AwsMetadataResult& result, - const std::string& hostname, - uint16_t port) { - for (const auto& endpoint : result.endpoints) { - if (endpoint.hostname == hostname && endpoint.port == port) return &endpoint; - } - return nullptr; -} - -} // namespace - -int main() { - plan(30); - - auto backend = std::make_shared(); - AwsSdkMetadataProvider provider(backend, AwsMetadataProviderConfig {2, 3}); - auto sink = std::make_shared(); - std::vector handles; - for (uint64_t id = 1; id <= 4; ++id) { - AwsMetadataRequest request; - request.kind = AwsMetadataRequestKind::rds_region; - request.opaque_id = id; - request.generation = 17; - request.region = "us-east-1"; - request.deadline = std::chrono::steady_clock::now() + 2s; - handles.push_back(provider.request(request, sink)); - } - ok(backend->wait_started(2), "provider starts its two bounded workers"); - ok(backend->max_active() == 2, "provider never exceeds two concurrent discoveries"); - ok(handles[0].value != 0 && handles[1].value != 0 && handles[2].value != 0, - "accepted work receives cancellable handles"); - ok(handles[3].value == 0 && sink->wait_for(1), - "bounded queue rejects excess work immediately"); - auto completions = sink->snapshot(); - ok(completions.size() >= 1 && completions[0].opaque_id == 4 && completions[0].generation == 17 && - completions[0].result.status == AwsMetadataStatus::throttled, - "queue rejection preserves request identity with a fixed category"); - ok(completions.size() >= 1 && completions[0].result.failure_category == "throttled" && - completions[0].result.failure_category.find("FAKE_SECRET") == std::string::npos, - "provider never forwards a backend's raw failure text"); - provider.cancel(handles[2]); - backend->release(); - ok(sink->wait_for(3), "accepted non-cancelled work completes"); - completions = sink->snapshot(); - bool saw_one = false; - bool saw_two = false; - bool saw_three = false; - for (const auto& completion : completions) { - saw_one = saw_one || (completion.opaque_id == 1 && completion.generation == 17); - saw_two = saw_two || (completion.opaque_id == 2 && completion.generation == 17); - saw_three = saw_three || completion.opaque_id == 3; - } - ok(saw_one && saw_two && !saw_three, - "queued cancellation suppresses its callback without affecting other work"); - - AwsMetadataRequest expired; - expired.opaque_id = 9; - expired.generation = 18; - expired.deadline = std::chrono::steady_clock::now() - 1ms; - ok(provider.request(expired, sink).value == 0 && sink->wait_for(4), - "already-expired work is rejected without entering the backend"); - completions = sink->snapshot(); - ok(!completions.empty() && completions.back().opaque_id == 9 && - completions.back().result.status == AwsMetadataStatus::timeout && - completions.back().result.failure_category == "timeout", - "deadline rejection is normalized and preserves identity"); - provider.shutdown(); - ok(provider.request(expired, sink).value == 0, - "shutdown permanently rejects new metadata work"); - - const std::unordered_map env { - {"AWS_REGION", "us-west-2"}, - {"AWS_DEFAULT_REGION", "eu-west-1"}, - {"AWS_AVAILABILITY_ZONE", "us-west-2b"}, - {"AWS_ACCOUNT_ID", "999900001111"}, - }; - const auto env_getter = [&](const char* name) { - const auto found = env.find(name); - return found == env.end() ? std::string() : found->second; - }; - AwsLocalLocation environment = aws_locality_environment_location(env_getter); - ok(environment.region == "us-west-2", - "AWS_REGION takes precedence over AWS_DEFAULT_REGION"); - ok(environment.availability_zone == "us-west-2b" && - environment.account_id == "999900001111", - "environment fallback retains optional AZ and account assertion"); - const auto partial_getter = [](const char* name) { - return std::string(name) == "AWS_DEFAULT_REGION" ? "eu-central-1" : ""; - }; - environment = aws_locality_environment_location(partial_getter); - ok(environment.region == "eu-central-1" && - environment.availability_zone.empty() && environment.account_id.empty(), - "partial environment fallback keeps Region while leaving AZ/account unknown"); - - auto imds = std::make_shared(); - AwsLocalityLocalDiscovery local_discovery(imds, env_getter); - const auto never_cancelled = [] { return false; }; - AwsMetadataResult local_result = local_discovery.discover( - std::chrono::steady_clock::now() + 1s, never_cancelled); - ok(local_result.status == AwsMetadataStatus::ok && - local_result.local.region == "us-east-1" && - local_result.local.availability_zone == "us-east-1b" && - local_result.local.account_id == "111122223333", - "IMDSv2 identity document supplies Region, AZ, and account"); - ok(imds->token_calls == 1 && imds->document_calls == 1 && - imds->seen_token == "imds-token", - "IMDSv2 token is required for the identity-document request"); - imds->token = {false, 0, "FAKE_SECRET_TRANSPORT_ERROR"}; - local_result = local_discovery.discover( - std::chrono::steady_clock::now() + 1s, never_cancelled); - ok(local_result.status == AwsMetadataStatus::ok && - local_result.local.region == "us-west-2", - "IMDS unavailability falls back to process environment"); - ok(local_result.failure_category.empty(), - "successful environment fallback exposes no IMDS transport detail"); - imds->token = {true, 200, "imds-token"}; - imds->document = {true, 200, "{not-json-FAKE_SECRET}"}; - AwsLocalityLocalDiscovery invalid_local(imds, - [](const char*) { return std::string(); }); - local_result = invalid_local.discover( - std::chrono::steady_clock::now() + 1s, never_cancelled); - ok(local_result.status == AwsMetadataStatus::invalid_response && - local_result.failure_category == "invalid_response", - "malformed IMDS data without fallback returns only a fixed category"); - - auto rds_api = std::make_shared(); - AwsLocalityRdsDiscovery rds_discovery(rds_api); - AwsMetadataRequest rds_request; - rds_request.kind = AwsMetadataRequestKind::rds_region; - rds_request.region = "us-east-1"; - rds_request.deadline = std::chrono::steady_clock::now() + 2s; - AwsMetadataResult rds_result = rds_discovery.discover(rds_request, never_cancelled); - ok(rds_result.status == AwsMetadataStatus::ok, - "paginated RDS discovery completes successfully"); - ok(rds_api->instance_calls == 2 && rds_api->cluster_calls == 1 && - rds_api->endpoint_calls == 1, - "RDS instances, clusters, and cluster endpoints are all paginated"); - const auto* instance_one = find_endpoint(rds_result, - "db-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); - ok(instance_one != nullptr && instance_one->endpoint_type == AwsEndpointType::instance && - instance_one->availability_zone == "us-east-1d" && - instance_one->account_id == "111122223333", - "later duplicate instance metadata replaces the earlier page after failover"); - size_t instance_one_count = 0; - for (const auto& endpoint : rds_result.endpoints) { - if (endpoint.hostname == "db-one.abcdefghijkl.us-east-1.rds.amazonaws.com" && - endpoint.port == 3306) ++instance_one_count; - } - ok(instance_one_count == 1, - "duplicate paginated endpoint metadata is emitted exactly once"); - const auto* instance_two = find_endpoint(rds_result, - "db-two.abcdefghijkl.us-east-1.rds.amazonaws.com", 3307); - ok(instance_two != nullptr && instance_two->account_id == "444455556666", - "later instance pages and cross-account identities are retained"); - ok(find_endpoint(rds_result, - "missing-port.abcdefghijkl.us-east-1.rds.amazonaws.com", 0) == nullptr, - "instance metadata without a valid port is rejected"); - const auto* cluster = find_endpoint(rds_result, - "cluster-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); - const auto* reader = find_endpoint(rds_result, - "cluster-ro-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); - ok(cluster != nullptr && cluster->endpoint_type == AwsEndpointType::cluster && - reader != nullptr && reader->endpoint_type == AwsEndpointType::reader && - cluster->availability_zone.empty() && reader->availability_zone.empty(), - "cluster writer and reader endpoints have Region/account but no stable AZ"); - const auto* custom_one = find_endpoint(rds_result, - "custom-one.abcdefghijkl.us-east-1.rds.amazonaws.com", 0); - const auto* custom_two = find_endpoint(rds_result, - "custom-two.abcdefghijkl.us-east-1.rds.amazonaws.com", 0); - ok(custom_one != nullptr && custom_two != nullptr && - custom_one->endpoint_type == AwsEndpointType::custom && - custom_two->endpoint_type == AwsEndpointType::custom && - custom_two->account_id == "111122223333", - "custom endpoints from both cluster APIs inherit cluster account and no port"); - - class RepeatingApi final : public AwsRdsDiscoveryApi { - public: - AwsRdsInstancesPage describe_instances(const std::string&, const std::string&, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { - AwsRdsInstancesPage page; - page.status = AwsMetadataStatus::ok; - page.next_marker = "same"; - return page; - } - AwsRdsClustersPage describe_clusters(const std::string&, const std::string&, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { return {}; } - AwsRdsClusterEndpointsPage describe_cluster_endpoints( - const std::string&, const std::string&, - std::chrono::steady_clock::time_point, - const AwsLocalityCancelPredicate&) override { return {}; } - }; - AwsLocalityRdsDiscovery repeating(std::make_shared()); - rds_result = repeating.discover(rds_request, never_cancelled); - ok(rds_result.status == AwsMetadataStatus::invalid_response && - rds_result.failure_category == "invalid_response", - "repeated pagination markers fail with a fixed category"); - - class ShutdownBackend final : public AwsLocalityDiscoveryBackend { - public: - AwsMetadataResult discover(const AwsMetadataRequest&, - const AwsLocalityCancelPredicate& cancelled) override { - started.store(true); - while (!cancelled()) std::this_thread::yield(); - AwsMetadataResult result; - result.status = AwsMetadataStatus::cancelled; - return result; - } - std::atomic started { false }; - }; - auto shutdown_backend = std::make_shared(); - AwsSdkMetadataProvider shutdown_provider( - shutdown_backend, AwsMetadataProviderConfig {2, 8}); - auto shutdown_sink = std::make_shared(); - AwsMetadataRequest shutdown_request; - shutdown_request.opaque_id = 44; - shutdown_request.deadline = std::chrono::steady_clock::now() + 2s; - shutdown_provider.request(shutdown_request, shutdown_sink); - while (!shutdown_backend->started.load()) std::this_thread::yield(); - std::thread shutdown_one([&] { shutdown_provider.shutdown(); }); - std::thread shutdown_two([&] { shutdown_provider.shutdown(); }); - shutdown_one.join(); - shutdown_two.join(); - ok(shutdown_sink->snapshot().empty(), - "concurrent shutdown is idempotent, cancels work, and publishes no late callback"); - ok(shutdown_provider.request(shutdown_request, shutdown_sink).value == 0, - "post-shutdown requests remain rejected"); - - return exit_status(); -} diff --git a/test/tap/tests/unit/aws_locality_selection_unit-t.cpp b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp index 70440d0f79..4af69f2889 100644 --- a/test/tap/tests/unit/aws_locality_selection_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp @@ -214,7 +214,7 @@ MySQL_Connection* make_connection(MySrvC* server, int fd) { } // namespace int main() { - plan(28); + plan(30); ok(aws_locality_saturating_add( std::numeric_limits::max() - 2, 5) == std::numeric_limits::max(), @@ -441,8 +441,43 @@ int main() { local_connection->options.client_flag &= ~CLIENT_FOUND_ROWS; } MyHGM->set_aws_locality_awareness_enabled(false); - test_cleanup_hostgroups(); shutdown_global_aws_metadata_provider(); + MyHGM->set_aws_locality_awareness_enabled(true); + ok(wait_until([&] { + const auto unavailable = MyHGM->aws_locality_manager()->snapshot(); + const auto* entry = unavailable->find(kHostgroup, kLocal, 3306); + return entry != nullptr && + entry->status == AwsLocalityMetadataStatus::error && + entry->failure_category == "provider_unavailable" && + unavailable->effective_weight(kHostgroup, kLocal, 3306, 10) == 10; + }), "missing provider publishes fixed provider_unavailable with neutral weight"); + { + MySQL_Thread neutral_worker; + if (!neutral_worker.init()) BAIL_OUT("neutral worker init failed"); + SessionFixture neutral_session(neutral_worker); + int neutral_counts[3] = {0, 0, 0}; + for (int i = 0; i < 12000; ++i) { + MySrvC* selected = hostgroup->get_random_MySrvC( + nullptr, 0, -1, neutral_session.session); + if (selected == local) ++neutral_counts[0]; + else if (selected == regional) ++neutral_counts[1]; + else if (selected == remote) ++neutral_counts[2]; + } + const double neutral_local_share = + static_cast(neutral_counts[0]) / 12000.0; + const double neutral_regional_share = + static_cast(neutral_counts[1]) / 12000.0; + const double neutral_remote_share = + static_cast(neutral_counts[2]) / 12000.0; + ok(neutral_local_share > 0.13 && neutral_local_share < 0.20 && + neutral_regional_share > 0.29 && neutral_regional_share < 0.38 && + neutral_remote_share > 0.45 && neutral_remote_share < 0.55 && + local->weight == 10 && regional->weight == 20 && remote->weight == 30, + "provider absence selects by unchanged configured 10:20:30 weights (%.3f/%.3f/%.3f)", + neutral_local_share, neutral_regional_share, neutral_remote_share); + } + MyHGM->set_aws_locality_awareness_enabled(false); + test_cleanup_hostgroups(); delete GloMyLogger; GloMyLogger = nullptr; test_cleanup_query_processor(); diff --git a/test/tap/tests/unit/aws_locality_stats_unit-t.cpp b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp index b842d3c3d4..5a1514ef27 100644 --- a/test/tap/tests/unit/aws_locality_stats_unit-t.cpp +++ b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp @@ -11,14 +11,22 @@ #include #include -#ifndef PROXYSQL_AWS_PLUGIN_PATH -#error "PROXYSQL_AWS_PLUGIN_PATH must be defined" -#endif - extern MySQL_HostGroups_Manager* MyHGM; namespace { +constexpr const char* kLocalityStatsProjectionFixture = + "CREATE TABLE stats_mysql_aws_locality (" + "hostgroup_id INT NOT NULL, hostname VARCHAR NOT NULL, port INT NOT NULL, " + "endpoint_type VARCHAR NOT NULL, configured_weight INT NOT NULL, " + "effective_weight INT NOT NULL, local_region VARCHAR NOT NULL, " + "local_az VARCHAR NOT NULL, backend_region VARCHAR NOT NULL, " + "backend_az VARCHAR NOT NULL, account_match VARCHAR NOT NULL, " + "locality VARCHAR NOT NULL, active_multiplier REAL NOT NULL, " + "metadata_status VARCHAR NOT NULL, last_success_timestamp INT NOT NULL, " + "last_attempt_timestamp INT NOT NULL, last_error_category VARCHAR NOT NULL, " + "PRIMARY KEY(hostgroup_id, hostname, port))"; + class CountingProvider final : public AwsMetadataProvider { public: AwsMetadataRequestHandle request( @@ -95,7 +103,7 @@ AwsLocalitySnapshotEntry diagnostic_row( } // namespace int main() { - plan(23); + plan(22); if (test_globals_init() != 0) { BAIL_OUT("test global initialization failed"); } @@ -105,26 +113,19 @@ int main() { SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX); ok(statsdb.return_one_int( "SELECT count(*) FROM sqlite_master WHERE name='stats_mysql_aws_locality'") == 0, - "AWS locality stats table is absent without the AWS plugin"); + "public core does not register an AWS locality stats table"); std::unique_ptr manager; std::string error; - ok(proxysql_load_configured_plugins(manager, {PROXYSQL_AWS_PLUGIN_PATH}, error), - "real AWS plugin completes schema-registration phase"); + ok(proxysql_load_configured_plugins(manager, {}, error) && manager == nullptr, + "provider-neutral plugin services expose no always-present locality schema"); if (!error.empty()) diag("plugin error: %s", error.c_str()); - const auto& tables = manager->tables(ProxySQL_PluginDBKind::stats_db); - ok(tables.size() == 1 && - std::string(tables[0].table_name) == "stats_mysql_aws_locality", - "AWS plugin registers exactly its locality table in stats DB"); - ok(manager->tables(ProxySQL_PluginDBKind::admin_db).empty() && - manager->tables(ProxySQL_PluginDBKind::config_db).empty(), - "locality diagnostics add no admin/config persistence surface"); - - if (!tables.empty()) statsdb.execute(tables[0].table_def); + ok(statsdb.execute(kLocalityStatsProjectionFixture), + "test-owned schema fixture accepts the public projection callback"); ok(statsdb.return_one_int( "SELECT count(*) FROM pragma_table_info('stats_mysql_aws_locality')") == 17, - "plugin-owned locality table has the exact 17-column schema"); + "projection fixture has the external-provider contract's 17 columns"); ok(statsdb.return_one_int( "SELECT count(*) FROM pragma_table_info('stats_mysql_aws_locality') " "WHERE name IN ('hostgroup_id','hostname','port','endpoint_type'," @@ -132,7 +133,7 @@ int main() { "'backend_region','backend_az','account_match','locality'," "'active_multiplier','metadata_status','last_success_timestamp'," "'last_attempt_timestamp','last_error_category')") == 17, - "locality stats schema exposes the documented column names"); + "projection callback targets the documented external schema columns"); std::vector rows; rows.push_back(diagnostic_row(1, AwsLocalityMetadataStatus::pending, 4.0, 10)); @@ -219,28 +220,25 @@ int main() { hostgroups.aws_locality_manager()->configure({disabled_hostgroup( 101, "first.abcdefghijkl.us-east-1.rds.amazonaws.com", 7)}); - proxysql_refresh_configured_plugin_runtime_views( - "SELECT * FROM stats_mysql_aws_locality", nullptr, nullptr, &statsdb); + hostgroups.refresh_aws_locality_stats(&statsdb); ok(statsdb.return_one_int( "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=101 " "AND metadata_status='disabled' AND configured_weight=7 " "AND effective_weight=7") == 1, - "real plugin callback projects the MySQL manager's current snapshot"); + "public callback projects the MySQL manager's current snapshot"); ok(counting_provider->requests.load() == 0, "query-time refresh issues no metadata-provider request"); hostgroups.aws_locality_manager()->configure({disabled_hostgroup( 202, "second.abcdefghijkl.us-east-1.rds.amazonaws.com", 9)}); - proxysql_refresh_configured_plugin_runtime_views( - "SELECT * FROM stats_mysql_aws_locality", nullptr, nullptr, &statsdb); + hostgroups.refresh_aws_locality_stats(&statsdb); ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 1 && statsdb.return_one_int( "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=202") == 1, "generation swap replaces the prior projection without mixed rows"); hostgroups.aws_locality_manager()->configure({}); - proxysql_refresh_configured_plugin_runtime_views( - "SELECT * FROM stats_mysql_aws_locality", nullptr, nullptr, &statsdb); + hostgroups.refresh_aws_locality_stats(&statsdb); ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 0, "no valid locality policy produces zero rows"); ok(counting_provider->requests.load() == 0,