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/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 new file mode 100644 index 0000000000..fa16f1fc50 --- /dev/null +++ b/doc/aws-locality-awareness.md @@ -0,0 +1,114 @@ +# AWS locality-aware MySQL backend selection + +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. + +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. + +## Hostgroup policy + +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, + '{ + "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 +``` + +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. + +## Master switch + +The process-wide MySQL variable defaults to `false`: + +```sql +SET mysql-aws_locality_awareness = true; +LOAD MYSQL VARIABLES TO RUNTIME; +SAVE MYSQL VARIABLES TO DISK; +``` + +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. + +## Selection contract + +For each selection attempt, after the normal health, capacity, lag, GTID, +backoff, and session-compatibility checks, 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, +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. + +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 + +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. + +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. + +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. + +## Rollback + +Set `mysql-aws_locality_awareness` to `false` and load MySQL variables to +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/include/Aws_Locality_Manager.h b/include/Aws_Locality_Manager.h new file mode 100644 index 0000000000..cf7b0f18e2 --- /dev/null +++ b/include/Aws_Locality_Manager.h @@ -0,0 +1,141 @@ +#ifndef __CLASS_AWS_LOCALITY_MANAGER_H +#define __CLASS_AWS_LOCALITY_MANAGER_H + +#include "Aws_Locality_Types.h" +#include "json_fwd.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#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); + +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 { +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_multimap entries; + std::unordered_set hostgroups; + + 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; + bool has_hostgroup(uint32_t hostgroup_id) const { + return hostgroups.find(hostgroup_id) != hostgroups.end(); + } +}; + +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::chrono::milliseconds disable_wait_timeout { std::chrono::milliseconds(250) }; + 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 // __CLASS_AWS_LOCALITY_MANAGER_H diff --git a/include/Aws_Locality_Types.h b/include/Aws_Locality_Types.h new file mode 100644 index 0000000000..47ff5a91f0 --- /dev/null +++ b/include/Aws_Locality_Types.h @@ -0,0 +1,169 @@ +#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, + 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; +}; + +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 // __CLASS_AWS_LOCALITY_TYPES_H 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..de85449b96 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,17 @@ 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); + 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(); + } +#endif #if 0 void wrlock(); void wrunlock(); diff --git a/include/MySQL_Thread.h b/include/MySQL_Thread.h index 25de6f075a..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]; @@ -645,6 +654,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_Plugin.h b/include/ProxySQL_Plugin.h index 5226a0ba1c..1b6cd36408 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,14 @@ 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. +// ABI 7: ProxySQL_PluginServices gains the MySQL-owned AWS-locality stats +// projection callback used by the AWS plugin's runtime view. +// 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, @@ -240,8 +247,17 @@ 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); + +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. @@ -308,6 +324,13 @@ 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; + // 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/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/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp new file mode 100644 index 0000000000..ee3f417736 --- /dev/null +++ b/lib/Aws_Locality_Manager.cpp @@ -0,0 +1,1149 @@ +#include "Aws_Locality_Manager.h" + +#include "json.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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 value >= minimum && value <= maximum; + } + 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; +} + +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 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 + +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 = aws_locality_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 (is_rds_proxy_endpoint_prefix(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); +} + +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; +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; + +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, + uint16_t 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) { + 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) { + 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 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( + 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() 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) { + 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(); + if (force_refresh_) { + ensure_worker_locked(); + } + 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_for(lock, config_.disable_wait_timeout, [&] { + 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, [&] { + return stopping_ || cancel_requested_ || force_refresh_ || + (enabled_ && !hostgroups_.empty()); + }); + 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; + if (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 = 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[endpoint_key(hostname, 0)] = &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(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) { + 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; + } + + // 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() || + !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_) { + next->hostgroups.insert(hostgroup.hostgroup_id); + for (const auto& backend : hostgroup.backends) { + auto entry = build_entry_locked(hostgroup, backend, now); + next->entries.emplace(identity_hash(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/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/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/lib/MyHGC.cpp b/lib/MyHGC.cpp index 7f3e1d4dbb..df54889d7e 100644 --- a/lib/MyHGC.cpp +++ b/lib/MyHGC.cpp @@ -31,9 +31,32 @@ 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) { + // 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; + }; if (l) { //int j=0; for (j=0; j32) { free(mysrvcCandidates); @@ -275,7 +298,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 +350,59 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ } } +#ifdef PROXYSQL40 + if (use_aws_locality) { + uint64_t total_weight = 0; + for (j = 0; j < num_candidates; ++j) { + mysrvc = mysrvcCandidates[j]; + 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()); + 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]; + 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 no AWS locality candidate is eligible\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_HostGroups_Manager.cpp b/lib/MySQL_HostGroups_Manager.cpp index a9c7216489..29fc23d865 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,140 @@ 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); + } +} + +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; +} + +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) { + 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) { p_hg_dyn_counter::metric metric = p_hg_dyn_counter::mysql_error; if (err_type == p_mysql_error_type::proxysql) { @@ -1641,6 +1788,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 +6392,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..fcc8aca8d8 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); @@ -6827,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 = @@ -6917,6 +6941,153 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local( ++i; } return NULL; + } + +#ifdef PROXYSQL40 + // 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 = + 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; + } + ++i; + } + + 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 (check_session_track_backoff && + candidate->parent->session_track_backoff_until.load( + std::memory_order_relaxed) > curtime) { + return false; + } + 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(candidate->userinfo->schemaname, schema) != 0) { + return false; + } + unsigned int not_match = 0; + candidate->number_of_matching_session_variables(client_conn, not_match); + if (not_match != 0) { + return false; + } + if (gtid_uuid == nullptr && max_lag_ms >= 0 && + static_cast(max_lag_ms) < + (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; + }; + + 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; + } + 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}); + } + } + 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()); + if (num_candidates == 0) { + return NULL; + } + 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(candidate.cached_index)); + } + i = next; + } +#endif + return NULL; } @@ -6951,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_Admin.cpp b/lib/ProxySQL_Admin.cpp index b7a9ce5562..8d519075a7 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 2abb4c2a8b..b581ffb217 100644 --- a/lib/ProxySQL_PluginManager.cpp +++ b/lib/ProxySQL_PluginManager.cpp @@ -6,6 +6,8 @@ #include "ProxySQL_PluginManager.h" #include "Aws_Iam_Provider.h" +#include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" #include "MySQL_Thread.h" #include @@ -189,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) @@ -196,6 +206,26 @@ 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); +} + +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() { @@ -324,6 +354,9 @@ 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; + 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 @@ -349,6 +382,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/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/groups/groups.json b/test/tap/groups/groups.json index 9393bdbc88..38071c7273 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -21,6 +21,11 @@ "aws_iam_pool_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "aws_iam_provider_boundary_unit-t" : [ "unit-tests-g1","mysqlx-tsan-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" ], "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..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; \ } @@ -419,7 +422,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 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 +498,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 \ @@ -935,6 +941,18 @@ 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 -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 + $(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) \ @@ -946,6 +964,10 @@ 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) + $(CXX) $< $(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_config_unit-t.cpp b/test/tap/tests/unit/aws_locality_config_unit-t.cpp new file mode 100644 index 0000000000..c31ac9f254 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_config_unit-t.cpp @@ -0,0 +1,195 @@ +#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(20); + + 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(); + + 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; + 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"); + MyHGM = previous_hgm; + } + GloVars.prometheus_registry.reset(); + + return exit_status(); +} 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..f5fef8a6d8 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp @@ -0,0 +1,677 @@ +#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 }; + bool block_requests { false }; + bool request_entered { false }; + bool release_requests { 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::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; + } + + 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; +} + +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; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) { + return true; + } + std::this_thread::sleep_for(1ms); + } + 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(49); + + 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(); + const auto* failed_refresh_entry = lookup(failed_refresh_snapshot, 11, east_one); + ok(failed_refreshes_delivered && + failed_refresh_entry != nullptr && + failed_refresh_entry->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"); + + 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"}), + }); + 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"; + 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(); + 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"); + 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( + 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(); + + 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; + 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_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, 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_after(replacement_state, + AwsMetadataRequestKind::local_location, "", std::move(result), + blocking_request_count); + callback_returned.store(true); + }); + { + std::unique_lock lock(completion_hook_mutex); + 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; + 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"); + + 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"); + invalid_identity_manager.shutdown(); + + 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_policy_unit-t.cpp b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp new file mode 100644 index 0000000000..152406e07c --- /dev/null +++ b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp @@ -0,0 +1,190 @@ +#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::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"); +} + +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"); + 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, + "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(34); + test_policy_validation(); + test_endpoint_recognition(); + test_classification(); + test_effective_weight(); + 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 new file mode 100644 index 0000000000..4af69f2889 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp @@ -0,0 +1,487 @@ +#include "tap.h" +#include "test_globals.h" +#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" +#include "MySQL_Thread.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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; + +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(30); + 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, + "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"); + 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"); + 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) { + 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); + 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( + 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); + 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(); + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +} 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..5a1514ef27 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp @@ -0,0 +1,269 @@ +#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 + +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( + 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); + if (test_globals_init() != 0) { + BAIL_OUT("test global initialization failed"); + } + + 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, + "public core does not register an AWS locality stats table"); + + std::unique_ptr manager; + std::string error; + 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()); + + 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, + "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'," + "'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, + "projection callback targets the documented external schema columns"); + + 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"); + + 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(); + { + 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)}); + 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, + "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)}); + 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({}); + 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, + "repeated generation queries remain network-free"); + MyHGM = nullptr; + } + 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); + test_globals_cleanup(); + return exit_status(); +}