PS-11220 Add MySQL instrumentation to memory

Author: jankowsk

mysql-test/suite/auth_openid_connect/r/auth.result

@@ -58,6 +58,9 @@ auth_openid_connect_users_authenticated	5
 auth_openid_connect_users_denied	16
 auth_openid_connect_users_proxied	0
 auth_openid_connect_users_with_roles	0
+SELECT COUNT_ALLOC>0, SUM_NUMBER_OF_BYTES_ALLOC>0, CURRENT_NUMBER_OF_BYTES_USED>0 FROM performance_schema.memory_summary_global_by_event_name WHERE EVENT_NAME LIKE 'memory/auth_openid_connect/%';
+COUNT_ALLOC>0	SUM_NUMBER_OF_BYTES_ALLOC>0	CURRENT_NUMBER_OF_BYTES_USED>0
+1	1	1
 
 ### CLEANUP
 DROP USER 'mysql_oidc_user'@'%';

mysql-test/suite/auth_openid_connect/t/auth.test

@@ -194,6 +194,8 @@ DROP USER invalid_user;
 
 SHOW GLOBAL STATUS LIKE 'auth_openid_connect%';
 
+# ensure memory instrumentation works
+SELECT COUNT_ALLOC>0, SUM_NUMBER_OF_BYTES_ALLOC>0, CURRENT_NUMBER_OF_BYTES_USED>0 FROM performance_schema.memory_summary_global_by_event_name WHERE EVENT_NAME LIKE 'memory/auth_openid_connect/%';
 ###################### CLEANUP #######################
 --echo
 --echo ### CLEANUP

plugin/auth_openid_connect/CMakeLists.txt

@@ -20,6 +20,7 @@ IF(WITH_AUTH_OPENID_CONNECT)
             src/config.cc
             src/jwk.cc
             src/jwks.cc
+            src/psi_openid_connect.cc
             src/udf.cc
             src/id_token.cc
     )

plugin/auth_openid_connect/src/config.cc

@@ -106,7 +106,7 @@ void Idp_configs::update(MYSQL_THD thd [[maybe_unused]],
 long long Idp_configs::update_keys() noexcept {
   long long no_updated_keys{0};
   try {
-    std::vector<std::pair<std::string, Idp_config>> configs;
+    std::vector<std::pair<Config_string, Idp_config>> configs;
     create_tmp_configs(configs);
     for (auto &[key, val] : configs) {
       if (!val.load_keys()) ++no_updated_keys;
@@ -129,7 +129,7 @@ long long Idp_configs::update_keys(const char *idp_name) noexcept {
     // effectively the keys will be removed from IDP. Note 2: during 2 the IDP
     // config may be removed, then 3 fails and the function returns an error.
     // That is not effective, but it is an edge case.
-    const std::string jwks_url{get_safe_jwks_url(idp_name)};
+    const Config_string jwks_url{get_safe_jwks_url(idp_name)};
     if (jwks_url.empty()) return 0;
     Idp_config config("", jwks_url, "", {}, {});
     const long long result = config.load_keys() ? 0 : 1;
@@ -177,11 +177,11 @@ void Idp_configs::load(const std::string &config_json) {
 #pragma GCC diagnostic ignored "-Wdangling-reference"
 #endif
 
-std::map<std::string, std::string> Idp_configs::load_group_roles(
+Config_string_map Idp_configs::load_group_roles(
     const picojson::object &idp_object, const std::string &idp_name) {
   static constexpr std::string_view key_group_role{"group-role"};
 
-  std::map<std::string, std::string> roles;
+  Config_string_map roles;
 
   const picojson::array &roles_array{json_get<picojson::array>(
       idp_object, std::string(key_group_role), idp_name, false)};
@@ -201,14 +201,14 @@ std::map<std::string, std::string> Idp_configs::load_group_roles(
   return roles;
 }
 
-std::unordered_set<std::string> Idp_configs::load_audiences(
+Config_string_set Idp_configs::load_audiences(
     const picojson::object &idp_object, const std::string &idp_name) {
   static constexpr std::string_view key_audiences{"audiences"};
-  std::unordered_set<std::string> audiences{};
+  Config_string_set audiences{};
   const picojson::array &audience_array{json_get<picojson::array>(
       idp_object, std::string(key_audiences), idp_name, false)};
   for (const auto &audience : audience_array) {
-    audiences.insert(audience.get<std::string>());
+    audiences.insert(Config_string(audience.get<std::string>()));
   }
   return audiences;
 }
@@ -323,7 +323,7 @@ void Idp_config::load_keys(const picojson::array &key_array,
 
     const picojson::object &key_object{key_value.get<picojson::object>()};
     const std::string &kty{json_get<std::string>(key_object, key_kty, from)};
-    const std::string &kid{json_get<std::string>(key_object, key_kid, from)};
+    const Config_string &kid{json_get<std::string>(key_object, key_kid, from)};
 
     std::string pem_key;
     if (kty == key_kty_rsa) {
@@ -478,5 +478,6 @@ std::string Idp_configs::verify_token(const Id_token &token,
   const std::shared_lock lock(mutex(), lock_timeout);
   if (!lock.owns_lock())
     throw std::runtime_error("failed to acquire shared lock on configuration");
-  return token.verify(ext_user, ext_group, current()->get_idp(idp_name), roles);
+  return token.verify(ext_user, ext_group,
+                      current()->get_idp(Config_string(idp_name)), roles);
 }

plugin/auth_openid_connect/src/config.h

@@ -18,6 +18,7 @@ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
 #ifndef AUTH_OIDC_CONFIG_H
 #define AUTH_OIDC_CONFIG_H
 
+#include <picojson/picojson.h>
 #include <chrono>
 #include <map>
 #include <memory>
@@ -29,16 +30,38 @@ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
 
 #include <mysql/components/services/bits/system_variables_bits.h>
 #include <mysql/plugin.h>
+#include <mysql/psi/psi_memory.h>
 #include <mysql/service_thd_alloc.h>
-#include <picojson/picojson.h>
 
 #include <shared_mutex>
 #include <vector>
 
 #include "jwks.h"
+#include "psi_openid_connect.h"
 
 class Id_token;
 
+template <typename T>
+using Config_allocator =
+    Psi_openid_connect::Allocator<T, Psi_openid_connect::config_memory_key>;
+
+using Config_string =
+    std::basic_string<char, std::char_traits<char>,
+                      Psi_openid_connect::Config_allocator<char>>;
+using Config_string_set =
+    std::unordered_set<Config_string, std::hash<Config_string>,
+                       std::equal_to<Config_string>,
+                       Psi_openid_connect::Config_allocator<Config_string>>;
+using Config_string_map =
+    std::map<Config_string, Config_string, std::less<Config_string>,
+             Psi_openid_connect::Config_allocator<
+                 std::pair<const Config_string, Config_string>>>;
+class Idp_config;
+using Config_idp_map =
+    std::map<Config_string, Idp_config, std::less<Config_string>,
+             Psi_openid_connect::Config_allocator<
+                 std::pair<const Config_string, Idp_config>>>;
+
 /**
  * @brief Statistics for Identity Provider (IDP) configurations.
  */
@@ -56,18 +79,16 @@ struct Idp_config_statistics {
  */
 class Idp_config {
  private:
-  std::string issuer_name;  ///< The token's issuer name.
-  Jwks jwks;                ///< URL for the JSON Web Key Set.
-  std::string group_claim;  ///< Name of the claim in the JWT that contains
-                            ///< group information.
-  std::unordered_set<std::string>
-      audiences;  ///< Set of allowed audiences for the token.
-  std::map<std::string, std::string>
-      roles;  ///< Map of IDP groups to database roles.
-  std::map<std::string, std::string>
-      keys;  ///< Map of Key ID (kid) to public key in PEM format.
+  Config_string issuer_name;    ///< The token's issuer name.
+  Jwks jwks;                    ///< Object managing the JSON Web Key Set.
+  Config_string group_claim;    ///< Name of the claim in the JWT that contains
+                                ///< group information.
+  Config_string_set audiences;  ///< Set of allowed audiences for the token.
+  Config_string_map roles;      ///< Map of IDP groups to database roles.
+  Config_string_map keys;  ///< Map of Key ID (kid) to public key in PEM format.
 
  public:
+  static PSI_memory_key config_memory_key;
   /**
    * @brief Constructs an Idp_config object.
    * @param issuer_name The token's issuer name.
@@ -76,10 +97,10 @@ class Idp_config {
    * @param audiences A set of allowed audiences.
    * @param roles A map of group-to-role mappings.
    */
-  Idp_config(const std::string &issuer_name, const std::string &jwks_url,
-             const std::string &group_claim,
-             std::unordered_set<std::string> &&audiences,
-             std::map<std::string, std::string> &&roles)
+  Idp_config(const std::string_view &issuer_name,
+             const std::string_view &jwks_url,
+             const std::string_view &group_claim, Config_string_set &&audiences,
+             Config_string_map &&roles)
       : issuer_name(issuer_name),
         jwks(jwks_url),
         group_claim(group_claim),
@@ -108,7 +129,9 @@ class Idp_config {
    * @brief Gets the issuer name.
    * @return The issuer name string.
    */
-  const std::string &get_issuer_name() const noexcept { return issuer_name; }
+  std::string get_issuer_name() const noexcept {
+    return std::string(issuer_name);
+  }
 
   /**
    * @brief Gets the JWKS URL.
@@ -120,7 +143,7 @@ class Idp_config {
    * @brief Gets the name of the group claim.
    * @return The group claim name.
    */
-  const std::string &get_group_claim() const noexcept { return group_claim; }
+  std::string_view get_group_claim() const noexcept { return group_claim; }
 
   /**
    * @brief Gets the only public key.
@@ -129,7 +152,7 @@ class Idp_config {
    * @return The only public key in PEM format.
    * @throws std::runtime_error if no keys are available.
    */
-  const std::string &get_the_only_pub_key() const {
+  std::string_view get_the_only_pub_key() const {
     if (keys.size() != 1) throw std::runtime_error("incorrect number of keys");
     const auto key{keys.cbegin()};
     return key->second;
@@ -141,8 +164,8 @@ class Idp_config {
    * @return The public key in PEM format.
    * @throws std::out_of_range if the Key ID is not found.
    */
-  const std::string &get_pub_key(const std::string &kid) const {
-    return keys.at(kid);
+  std::string_view get_pub_key(const std::string_view &kid) const {
+    return keys.at(Config_string(kid));
   }
 
   /**
@@ -156,18 +179,18 @@ class Idp_config {
    * @param audience The audience string from the token.
    * @return true if the audience is allowed
    */
-  bool is_audience_allowed(const std::string &audience) const noexcept {
-    return audiences.contains(audience);
+  bool is_audience_allowed(const std::string_view &audience) const noexcept {
+    return audiences.contains(Config_string(audience));
   }
 
   /**
    * @brief Gets the mapped database role for an IDP group.
    * @param group The IDP group name.
    * @return The mapped role name, or an empty string if no mapping exists.
    */
-  const std::string &get_role(const std::string &group) const noexcept {
-    static std::string no_role;
-    const auto it = roles.find(group);
+  std::string_view get_role(const std::string_view &group) const noexcept {
+    static Config_string no_role;
+    const auto it = roles.find(Config_string(group));
     return it == roles.end() ? no_role : it->second;
   }
 
@@ -190,9 +213,8 @@ class Idp_config {
  */
 class Idp_configs {
  private:
-  std::string sysvar_str{};  ///< Value of the configuration system variable.
-  std::map<std::string, Idp_config>
-      idp_configs{};  ///< Map of IDP names to Idp_config objects.
+  Config_string sysvar_str{};  ///< Value of the configuration system variable.
+  Config_idp_map idp_configs{};  ///< Map of IDP names to Idp_config objects.
   Idp_configs() = delete;
 
   /**
@@ -255,8 +277,8 @@ class Idp_configs {
    * @return map of IDP group names to database role names. If no mapping is
    * present the returned map will be empty.
    */
-  static std::map<std::string, std::string> load_group_roles(
-      const picojson::object &idp_object, const std::string &idp_name);
+  static Config_string_map load_group_roles(const picojson::object &idp_object,
+                                            const std::string &idp_name);
 
   /**
    * @brief Load the set of allowed audiences from an IDP JSON object.
@@ -268,8 +290,8 @@ class Idp_configs {
    * @return An unordered set containing the allowed audience
    *         strings. If no audiences are configured the set will be empty.
    */
-  static std::unordered_set<std::string> load_audiences(
-      const picojson::object &idp_object, const std::string &idp_name);
+  static Config_string_set load_audiences(const picojson::object &idp_object,
+                                          const std::string &idp_name);
 
   /**
    * @brief Parses the prefix of the configuration system variable.
@@ -301,10 +323,11 @@ class Idp_configs {
    * @return reference to the Idp_config object, or nullptr if not found.
    * @throws std::runtime_error if the IDP is not found in the configuration.
    */
-  Idp_config &get_idp(const std::string &idp_name) {
+  Idp_config &get_idp(const Config_string &idp_name) {
     const auto it = idp_configs.find(idp_name);
     if (it == idp_configs.end())
-      throw std::runtime_error("IDP not found: " + idp_name);
+      throw std::runtime_error(std::string("IDP not found: ") +
+                               idp_name.c_str());
     return it->second;
   }
 
@@ -315,7 +338,7 @@ class Idp_configs {
    * @param other_idp The second IDP config.
    * @throws std::runtime_error if the IDP is not found in the configuration.
    */
-  static void swap_idp_keys(const std::string &idp_name,
+  static void swap_idp_keys(const Config_string &idp_name,
                             Idp_config &other_idp) {
     std::unique_lock lock(mutex());
     current()->get_idp(idp_name).swap_keys(other_idp);
@@ -326,7 +349,7 @@ class Idp_configs {
    * @param idp_name The name of the IDP.
    * @return The JWKS URL for the specified IDP.
    */
-  static const std::string &get_safe_jwks_url(const std::string &idp_name) {
+  static const std::string &get_safe_jwks_url(const Config_string &idp_name) {
     std::shared_lock lock(mutex(), lock_timeout);
     if (!lock.owns_lock())
       throw std::runtime_error("failed to acquire shared lock");
@@ -337,7 +360,7 @@ class Idp_configs {
    * @brief Gets the JWKS URL for a specific IDP in a thread-safe manner.
    */
   static void swap_idp_keys(
-      std::vector<std::pair<std::string, Idp_config>> &configs) {
+      std::vector<std::pair<Config_string, Idp_config>> &configs) {
     std::unique_lock lock(mutex(), lock_timeout);
     if (!lock.owns_lock())
       throw std::runtime_error("failed to acquire unique lock");
@@ -353,7 +376,7 @@ class Idp_configs {
    * corresponding temporary Idp_config objects
    */
   static void create_tmp_configs(
-      std::vector<std::pair<std::string, Idp_config>> &configs) {
+      std::vector<std::pair<Config_string, Idp_config>> &configs) {
     std::shared_lock lock(mutex(), lock_timeout);
     if (!lock.owns_lock())
       throw std::runtime_error("failed to acquire shared lock");

plugin/auth_openid_connect/src/plugin_openid_connect.cc

@@ -20,7 +20,6 @@ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
 #include <picojson/picojson.h>
 #include <stddef.h>
 #include <cassert>
-#include <cstring>
 #include <exception>
 #include <map>
 #include <memory>
@@ -38,6 +37,7 @@ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
 
 #include "config.h"
 #include "id_token.h"
+#include "psi_openid_connect.h"
 
 static SERVICE_TYPE(registry) * reg_srv(nullptr);
 SERVICE_TYPE(log_builtins) * log_bi(nullptr);
@@ -90,6 +90,7 @@ long long Connection_status::users_with_roles{0};
  * @return 0 for success, 1 for error.
  */
 static int auth_oidc_init(MYSQL_PLUGIN plugin_info [[maybe_unused]]) {
+  Psi_openid_connect::init();
   if (init_logging_service_for_plugin(&reg_srv, &log_bi, &log_bs)) return 1;
   if (curl_global_init(CURL_GLOBAL_DEFAULT) != CURLE_OK) {
     LogPluginErr(ERROR_LEVEL, ER_LOG_PRINTF_MSG, "curl_global_init failed");

plugin/auth_openid_connect/src/psi_openid_connect.cc

@@ -0,0 +1,40 @@
+/*
+(C) 2026 Percona LLC and/or its affiliates
+
+This program is free software; you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation; version 2 of the License.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program; if not, write to the Free Software
+Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
+*/
+
+#include "psi_openid_connect.h"
+
+#include <mysql/psi/mysql_memory.h>
+#include <template_utils.h>
+
+namespace Psi_openid_connect {
+
+PSI_memory_key config_memory_key;
+
+static constexpr auto config_memory{"config_memory"};
+
+static PSI_memory_info all_memory[] = {{&config_memory_key, config_memory, 0,
+                                        PSI_VOLATILITY_UNKNOWN,
+                                        PSI_DOCUMENT_ME}};
+
+void init() {
+  static constexpr auto category{"auth_openid_connect"};
+
+  int count = array_elements(all_memory);
+  mysql_memory_register(category, all_memory, count);
+}
+
+}  // namespace Psi_openid_connect