jax_fingerprint: Address user arg slot exhaustion

plugins/experimental/jax_fingerprint/context_map.h

@@ -0,0 +1,143 @@
+/** @file
+ *
+ * Shared context map for jax_fingerprint plugin instances.
+ *
+ * When multiple jax_fingerprint.so instances are loaded (e.g., one per
+ * fingerprinting method), they share a single user arg slot. This map
+ * stores the JAxContext for each method, keyed by method name.
+ *
+ * @section license License
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "config.h"
+#include "context.h"
+
+#include <string>
+#include <string_view>
+#include <unordered_map>
+#include <version>
+
+/**
+ * @brief Container holding JAxContext instances for multiple methods.
+ *
+ * ATS has a limited number of user arg slots (~4 per type). When loading
+ * many jax_fingerprint instances, we share a single slot and store all
+ * contexts in this map, keyed by method name.
+ */
+class ContextMap
+{
+public:
+  ~ContextMap()
+  {
+    for (auto &pair : m_contexts) {
+      delete pair.second;
+    }
+  }
+
+  /**
+   * @brief Store a context for a method.
+   * @param[in] method_name The method name (e.g., "JA3", "JA4").
+   * @param[in] ctx The context to store. Ownership is transferred.
+   */
+  void
+  set(std::string_view method_name, JAxContext *ctx)
+  {
+    auto it = find_context(method_name);
+    if (it != m_contexts.end()) {
+      delete it->second;
+      it->second = ctx;
+    } else {
+      m_contexts.emplace(std::string{method_name}, ctx);
+    }
+  }
+
+  /**
+   * @brief Retrieve a context for a method.
+   * @param[in] method_name The method name.
+   * @return The context, or nullptr if not found.
+   */
+  JAxContext *
+  get(std::string_view method_name)
+  {
+    auto it = find_context(method_name);
+    return it != m_contexts.end() ? it->second : nullptr;
+  }
+
+  /**
+   * @brief Remove a context for a method.
+   * @param[in] method_name The method name.
+   */
+  void
+  remove(std::string_view method_name)
+  {
+    auto it = find_context(method_name);
+    if (it != m_contexts.end()) {
+      delete it->second;
+      m_contexts.erase(it);
+    }
+  }
+
+  /**
+   * @brief Check if the map is empty.
+   * @return True if no contexts are stored.
+   */
+  bool
+  empty() const
+  {
+    return m_contexts.empty();
+  }
+
+private:
+  using ContextStorage = std::unordered_map<std::string, JAxContext *, StringHash, std::equal_to<>>;
+
+  /** Find context by method name with C++20 generic lookup fallback.
+   *
+   * C++20 generic unordered lookup allows finding with std::string_view in a
+   * std::unordered_map<std::string, ...> without creating a temporary string.
+   * For standard libraries without this feature, we fall back to constructing
+   * a std::string for the lookup.
+   *
+   * @param[in] method_name The method name to look up.
+   * @return Iterator to the found element, or end() if not found.
+   */
+  ContextStorage::iterator
+  find_context(std::string_view method_name)
+  {
+#ifdef __cpp_lib_generic_unordered_lookup
+    return m_contexts.find(method_name);
+#else
+    return m_contexts.find(std::string{method_name});
+#endif
+  }
+
+  /** const_iterator @overload */
+  ContextStorage::const_iterator
+  find_context(std::string_view method_name) const
+  {
+#ifdef __cpp_lib_generic_unordered_lookup
+    return m_contexts.find(method_name);
+#else
+    return m_contexts.find(std::string{method_name});
+#endif
+  }
+
+  ContextStorage m_contexts;
+};

plugins/experimental/jax_fingerprint/plugin.cc

@@ -282,8 +282,7 @@ handle_http_txn_close(void *edata, PluginConfig &config)
 {
   TSHttpTxn txnp = static_cast<TSHttpTxn>(edata);
 
-  delete get_user_arg(txnp, config);
-  set_user_arg(txnp, config, nullptr);
+  cleanup_user_arg(txnp, config);
 
   TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
   return TS_SUCCESS;
@@ -294,8 +293,7 @@ handle_vconn_close(void *edata, PluginConfig &config)
 {
   TSVConn vconn = static_cast<TSVConn>(edata);
 
-  delete get_user_arg(vconn, config);
-  set_user_arg(vconn, config, nullptr);
+  cleanup_user_arg(vconn, config);
 
   TSVConnReenable(vconn);
   return TS_SUCCESS;

plugins/experimental/jax_fingerprint/userarg.cc

@@ -22,32 +22,101 @@
 #include "plugin.h"
 #include "config.h"
 #include "userarg.h"
+#include "context_map.h"
+
+namespace
+{
+
+char const *
+user_arg_type_name(TSUserArgType type)
+{
+  switch (type) {
+  case TS_USER_ARGS_VCONN:
+    return "vconn";
+  case TS_USER_ARGS_TXN:
+    return "txn";
+  default:
+    return "unknown";
+  }
+}
+
+// Shared user arg indices for all jax_fingerprint instances. ATS has limited
+// slots (~4 per type), so we share one slot per type and use a ContextMap.
+int  vconn_user_arg_index = -1;
+int  txn_user_arg_index   = -1;
+bool vconn_slot_reserved  = false;
+bool txn_slot_reserved    = false;
+
+} // anonymous namespace
 
 int
 reserve_user_arg(PluginConfig &config)
 {
-  std::string name  = PLUGIN_NAME;
-  name             += config.method.name;
-
   TSUserArgType type;
+  int          *shared_index;
+  bool         *reserved_flag;
+
   if (config.method.type == Method::Type::CONNECTION_BASED) {
-    type = TS_USER_ARGS_VCONN;
+    type          = TS_USER_ARGS_VCONN;
+    shared_index  = &vconn_user_arg_index;
+    reserved_flag = &vconn_slot_reserved;
   } else {
-    type = TS_USER_ARGS_TXN;
+    type          = TS_USER_ARGS_TXN;
+    shared_index  = &txn_user_arg_index;
+    reserved_flag = &txn_slot_reserved;
   }
-  int ret = TSUserArgIndexReserve(type, name.c_str(), "used to pass JAx context between hooks", &config.user_arg_index);
-  Dbg(dbg_ctl, "user_arg_name: %s, user_arg_index: %d", name.c_str(), config.user_arg_index);
-  return ret;
+
+  // Only reserve the slot once per type; subsequent calls reuse it.
+  if (!*reserved_flag) {
+    int ret = TSUserArgIndexReserve(type, PLUGIN_NAME, "shared JAx context map", shared_index);
+    if (ret == TS_SUCCESS) {
+      *reserved_flag = true;
+      Dbg(dbg_ctl, "Reserved shared user_arg slot: type=%s, index=%d", user_arg_type_name(type), *shared_index);
+    } else {
+      Dbg(dbg_ctl, "Failed to reserve shared user_arg slot: type=%s", user_arg_type_name(type));
+      return ret;
+    }
+  }
+
+  config.user_arg_index = *shared_index;
+  Dbg(dbg_ctl, "Using shared user_arg: type=%s, method=%.*s, index=%d", user_arg_type_name(type),
+      static_cast<int>(config.method.name.size()), config.method.name.data(), config.user_arg_index);
+  return TS_SUCCESS;
 }
 
 void
 set_user_arg(void *container, PluginConfig &config, JAxContext *ctx)
 {
-  TSUserArgSet(container, config.user_arg_index, static_cast<void *>(ctx));
+  ContextMap *map = static_cast<ContextMap *>(TSUserArgGet(container, config.user_arg_index));
+  if (map == nullptr) {
+    map = new ContextMap();
+    TSUserArgSet(container, config.user_arg_index, static_cast<void *>(map));
+  }
+  map->set(config.method.name, ctx);
 }
 
 JAxContext *
 get_user_arg(void *container, PluginConfig &config)
 {
-  return static_cast<JAxContext *>(TSUserArgGet(container, config.user_arg_index));
+  ContextMap *map = static_cast<ContextMap *>(TSUserArgGet(container, config.user_arg_index));
+  if (map == nullptr) {
+    return nullptr;
+  }
+  return map->get(config.method.name);
+}
+
+void
+cleanup_user_arg(void *container, PluginConfig &config)
+{
+  ContextMap *map = static_cast<ContextMap *>(TSUserArgGet(container, config.user_arg_index));
+  if (map != nullptr) {
+    // Remove this plugin's context from the map.
+    map->remove(config.method.name);
+
+    // If the map is now empty, delete it and clear the user arg.
+    if (map->empty()) {
+      delete map;
+      TSUserArgSet(container, config.user_arg_index, nullptr);
+    }
+  }
 }

plugins/experimental/jax_fingerprint/userarg.h

@@ -30,3 +30,4 @@
 int         reserve_user_arg(PluginConfig &config);
 void        set_user_arg(void *container, PluginConfig &config, JAxContext *ctx);
 JAxContext *get_user_arg(void *container, PluginConfig &config);
+void        cleanup_user_arg(void *container, PluginConfig &config);