Enable key placeholder to be combined with other data within the same command argument (#154)

* Enable key placeholder to be combined with other data within the same command argument

* Ensured tests cover key format with prefix/suffix

* Fixes per PR review: added KEY_BUFFER_STACK_SIZE instead of constexpr

---------

Co-authored-by: fcostaoliveira <[email protected]>

Author: filipecosta90

client.cpp

@@ -46,10 +46,12 @@
 
 #include <math.h>
 #include <algorithm>
+#include <sstream>
 #include <arpa/inet.h>
 
 #include "client.h"
 #include "cluster_client.h"
+#include "config_types.h"
 
 
 bool client::setup_client(benchmark_config *config, abstract_protocol *protocol, object_generator *objgen)
@@ -286,7 +288,46 @@ bool client::create_arbitrary_request(unsigned int command_index, struct timeval
             get_key_response res = get_key_for_conn(command_index, conn_id, &key_index);
             /* If key not available for this connection, we have a bug of sending partial request */
             assert(res == available_for_conn);
-            cmd_size += m_connections[conn_id]->send_arbitrary_command(arg, m_obj_gen->get_key(), m_obj_gen->get_key_len());
+
+            //when we have static data mixed with the key placeholder
+            if (arg->has_key_affixes) {
+                // Pre-calculate total length to avoid reallocations
+                const char* key = m_obj_gen->get_key();
+                unsigned int key_len = m_obj_gen->get_key_len();
+                size_t prefix_len = arg->data_prefix.length();
+                size_t suffix_len = arg->data_suffix.length();
+                size_t total_len = prefix_len + key_len + suffix_len;
+
+                // Optimization: use stack buffer for small keys to avoid heap allocation
+                if (total_len < KEY_BUFFER_STACK_SIZE) {
+                    char stack_buffer[KEY_BUFFER_STACK_SIZE];
+                    char* pos = stack_buffer;
+
+                    // Manual copy for better performance
+                    if (prefix_len > 0) {
+                        memcpy(pos, arg->data_prefix.data(), prefix_len);
+                        pos += prefix_len;
+                    }
+                    memcpy(pos, key, key_len);
+                    pos += key_len;
+                    if (suffix_len > 0) {
+                        memcpy(pos, arg->data_suffix.data(), suffix_len);
+                    }
+
+                    cmd_size += m_connections[conn_id]->send_arbitrary_command(arg, stack_buffer, total_len);
+                } else {
+                    // Fallback to string for large keys
+                    std::string combined_key;
+                    combined_key.reserve(total_len);
+                    combined_key.append(arg->data_prefix);
+                    combined_key.append(key, key_len);
+                    combined_key.append(arg->data_suffix);
+
+                    cmd_size += m_connections[conn_id]->send_arbitrary_command(arg, combined_key.c_str(), combined_key.length());
+                }
+            } else{
+                cmd_size += m_connections[conn_id]->send_arbitrary_command(arg, m_obj_gen->get_key(), m_obj_gen->get_key_len());
+            }
         } else if (arg->type == data_type) {
             unsigned int value_len;
             const char *value = m_obj_gen->get_value(0, &value_len);

client.h

@@ -49,6 +49,9 @@ class object_generator;
 #define SET_CMD_IDX 0
 #define GET_CMD_IDX 2
 
+// Stack buffer size for key operations to avoid heap allocation
+#define KEY_BUFFER_STACK_SIZE 512
+
 enum get_key_response { not_available, available_for_conn, available_for_other_conn };
 
 class client : public connections_manager {

config_types.h

@@ -115,9 +115,14 @@ enum command_arg_type {
 };
 
 struct command_arg {
-    command_arg(const char* arg, unsigned int arg_len) : type(undefined_type), data(arg, arg_len) {;}
+    command_arg(const char* arg, unsigned int arg_len) : type(undefined_type), data(arg, arg_len), has_key_affixes(false) {;}
     command_arg_type type;
     std::string data;
+    // the prefix and suffix strings are used for mixed key placeholder storing of substrings
+    std::string data_prefix;
+    std::string data_suffix;
+    // optimization flag to avoid runtime checks
+    bool has_key_affixes;
 };
 
 struct arbitrary_command {

protocol.cpp

@@ -713,8 +713,9 @@ bool redis_protocol::format_arbitrary_command(arbitrary_command &cmd) {
         // check arg type
         if (current_arg->data.find(KEY_PLACEHOLDER) != std::string::npos) {
             if (current_arg->data.length() != strlen(KEY_PLACEHOLDER)) {
-                benchmark_error_log("error: key placeholder can't combined with other data\n");
-                return false;
+                current_arg->has_key_affixes = true;
+                current_arg->data_prefix = current_arg->data.substr(0, current_arg->data.find(KEY_PLACEHOLDER));
+                current_arg->data_suffix = current_arg->data.substr(current_arg->data.find(KEY_PLACEHOLDER) + strlen(KEY_PLACEHOLDER));
             }
             cmd.keys_count++;
             current_arg->type = key_type;

tests/tests_oss_simple_flow.py

@@ -615,6 +615,76 @@ def test_default_arbitrary_command_hset_multi_data_placeholders(env):
     assert_minimum_memtier_outcomes(config, env, memtier_ok, overall_expected_request_count,
                                     overall_request_count)
 
+
+def test_key_placeholder(env):
+    env.skipOnCluster()
+    run_count = 1
+    benchmark_specs = {"name": env.testName, "args": ['--command=HSET __key__ f __data__']}
+    addTLSArgs(benchmark_specs, env)
+    config = get_default_memtier_config()
+    master_nodes_list = env.getMasterNodesList()
+    overall_expected_request_count = get_expected_request_count(config) * run_count
+
+    add_required_env_arguments(benchmark_specs, config, env, master_nodes_list)
+
+    # Create a temporary directory
+    test_dir = tempfile.mkdtemp()
+
+    config = RunConfig(test_dir, env.testName, config, {})
+    ensure_clean_benchmark_folder(config.results_dir)
+
+    benchmark = Benchmark.from_json(config, benchmark_specs)
+
+    # benchmark.run() returns True if the return code of memtier_benchmark was 0
+    memtier_ok = benchmark.run()
+    debugPrintMemtierOnError(config, env)
+
+    master_nodes_connections = env.getOSSMasterNodesConnectionList()
+    merged_command_stats = {'cmdstat_hset': {'calls': 0}}
+    overall_request_count = agg_info_commandstats(master_nodes_connections, merged_command_stats)
+    assert_minimum_memtier_outcomes(config, env, memtier_ok, overall_expected_request_count,
+                                    overall_request_count)
+
+
+# key placeholder combined with other data
+def test_key_placeholder_togetherwithdata(env):
+    env.skipOnCluster()
+    run_count = 1
+    benchmark_specs = {"name": env.testName, "args": ['--command=SET \"prefix:__key__:suffix\" \"__data__\"']}
+    addTLSArgs(benchmark_specs, env)
+    config = get_default_memtier_config(threads=4, clients=1,requests=50)
+    master_nodes_list = env.getMasterNodesList()
+    overall_expected_request_count = get_expected_request_count(config) * run_count
+
+    add_required_env_arguments(benchmark_specs, config, env, master_nodes_list)
+
+    # Create a temporary directory
+    test_dir = tempfile.mkdtemp()
+
+    config = RunConfig(test_dir, env.testName, config, {})
+    ensure_clean_benchmark_folder(config.results_dir)
+
+    benchmark = Benchmark.from_json(config, benchmark_specs)
+
+    # benchmark.run() returns True if the return code of memtier_benchmark was 0
+    memtier_ok = benchmark.run()
+    debugPrintMemtierOnError(config, env)
+
+    master_nodes_connections = env.getOSSMasterNodesConnectionList()
+    merged_command_stats = {'cmdstat_set': {'calls': 0}}
+    overall_request_count = agg_info_commandstats(master_nodes_connections, merged_command_stats)
+    assert_minimum_memtier_outcomes(config, env, memtier_ok, overall_expected_request_count,
+                                    overall_request_count)
+
+    # Ensure all keys have the correct prefix and suffix
+    for conn in master_nodes_connections:
+        for key in conn.scan_iter("*"):
+            decoded_key = key.decode().split(":")
+            env.assertEqual(decoded_key[0], "prefix")
+            env.assertEqual(decoded_key[1].split("-")[0], "memtier")
+            env.assertEqual(decoded_key[2], "suffix")
+
+
 def test_default_set_get_rate_limited(env):
     env.skipOnCluster()
     master_nodes_list = env.getMasterNodesList()