Enable arbitrary commands to be run on cluster mode

Author: filipecosta90

cluster_client.cpp

@@ -358,9 +358,54 @@ bool cluster_client::get_key_for_conn(unsigned int conn_id, int iter, unsigned l
     }
 }
 
+
+void cluster_client::create_arbitrary_request(const arbitrary_command* cmd, struct timeval& timestamp, unsigned int conn_id) {
+    int cmd_size = 0;
+
+    benchmark_debug_log("%s [%s]:\n", cmd->command_name.c_str(), cmd->command.c_str());
+
+    for (unsigned int i = 0; i < cmd->command_args.size(); i++) {
+        const command_arg* arg = &cmd->command_args[i];
+
+        if (arg->type == const_type) {
+            cmd_size += m_connections[conn_id]->send_arbitrary_command(arg);
+        } else if (arg->type == key_type) {
+            unsigned long long key_index;
+
+            // get key
+            if (!get_key_for_conn(conn_id,  get_arbitrary_obj_iter_type(cmd, m_executed_command_index), &key_index)) {
+                return;
+            }
+
+            assert(key_index >= 0);
+            assert(m_key_len > 0);
+
+            cmd_size += m_connections[conn_id]->send_arbitrary_command(arg, m_key_buffer, m_key_len);
+        } else if (arg->type == data_type) {
+            unsigned int value_len;
+            const char *value = m_obj_gen->get_value(0, &value_len);
+
+            assert(value != NULL);
+            assert(value_len > 0);
+
+            cmd_size += m_connections[conn_id]->send_arbitrary_command(arg, value, value_len);
+        }
+    }
+
+    m_connections[conn_id]->send_arbitrary_command_end(m_executed_command_index, &timestamp, cmd_size);
+}
+
 // This function could use some urgent TLC -- but we need to do it without altering the behavior
 void cluster_client::create_request(struct timeval timestamp, unsigned int conn_id)
 {
+    // are we using arbitrary command?
+    if (m_config->arbitrary_commands->is_defined()) {
+        const arbitrary_command* executed_command = m_config->arbitrary_commands->get_next_executed_command(m_arbitrary_command_ratio_count,
+                                                                                                      m_executed_command_index);
+        create_arbitrary_request(executed_command, timestamp, conn_id);
+        return;
+    }
+
     // If the Set:Wait ratio is not 0, start off with WAITs
     if (m_config->wait_ratio.b &&
         (m_tot_wait_ops == 0 ||
@@ -416,16 +461,28 @@ void cluster_client::create_request(struct timeval timestamp, unsigned int conn_
 void cluster_client::handle_moved(unsigned int conn_id, struct timeval timestamp,
                                   request *request, protocol_response *response) {
     // update stats
-    if (request->m_type == rt_get) {
-        m_stats.update_moved_get_op(&timestamp,
+    switch (request->m_type) {
+        case rt_get:
+            m_stats.update_moved_get_op(&timestamp,
                                     request->m_size + response->get_total_len(),
                                     ts_diff(request->m_sent_time, timestamp));
-    } else if (request->m_type == rt_set) {
-        m_stats.update_moved_set_op(&timestamp,
+            break;
+        case rt_set:
+            m_stats.update_moved_set_op(&timestamp,
                                     request->m_size + response->get_total_len(),
                                     ts_diff(request->m_sent_time, timestamp));
-    } else {
-        assert(0);
+            break;
+        case rt_arbitrary: {
+            arbitrary_request *ar = static_cast<arbitrary_request *>(request);
+            m_stats.update_moved_arbitrary_op(&timestamp,
+                                        request->m_size + response->get_total_len(),
+                                        ts_diff(request->m_sent_time, timestamp),
+                                        ar->index);
+            break;
+        }
+        default:
+            assert(0);
+            break;
     }
 
     // connection already issued 'cluster slots' command, wait for slots mapping to be updated
@@ -444,16 +501,28 @@ void cluster_client::handle_moved(unsigned int conn_id, struct timeval timestamp
 void cluster_client::handle_ask(unsigned int conn_id, struct timeval timestamp,
                                 request *request, protocol_response *response) {
     // update stats
-    if (request->m_type == rt_get) {
-        m_stats.update_ask_get_op(&timestamp,
-                                    request->m_size + response->get_total_len(),
-                                    ts_diff(request->m_sent_time, timestamp));
-    } else if (request->m_type == rt_set) {
-        m_stats.update_ask_set_op(&timestamp,
-                                    request->m_size + response->get_total_len(),
-                                    ts_diff(request->m_sent_time, timestamp));
-    } else {
-        assert(0);
+    switch (request->m_type) {
+        case rt_get:
+            m_stats.update_ask_get_op(&timestamp,
+                                      request->m_size + response->get_total_len(),
+                                      ts_diff(request->m_sent_time, timestamp));
+            break;
+        case rt_set:
+            m_stats.update_ask_set_op(&timestamp,
+                                      request->m_size + response->get_total_len(),
+                                      ts_diff(request->m_sent_time, timestamp));
+            break;
+        case rt_arbitrary: {
+            arbitrary_request *ar = static_cast<arbitrary_request *>(request);
+            m_stats.update_ask_arbitrary_op(&timestamp,
+                                              request->m_size + response->get_total_len(),
+                                              ts_diff(request->m_sent_time, timestamp),
+                                              ar->index);
+            break;
+        }
+        default:
+            assert(0);
+            break;
     }
 }
 

cluster_client.h

@@ -52,6 +52,7 @@ class cluster_client : public client {
 
     // client manager api's
     virtual void handle_cluster_slots(protocol_response *r);
+    virtual void create_arbitrary_request(const arbitrary_command* cmd, struct timeval& timestamp, unsigned int conn_id);
     virtual void create_request(struct timeval timestamp, unsigned int conn_id);
     virtual bool hold_pipeline(unsigned int conn_id);
     virtual void handle_response(unsigned int conn_id, struct timeval timestamp,

memtier_benchmark.cpp

@@ -313,9 +313,6 @@ static bool verify_cluster_option(struct benchmark_config *cfg) {
     } else if (cfg->unix_socket) {
         fprintf(stderr, "error: cluster mode dose not support unix-socket option.\n");
         return false;
-    } else if (cfg->arbitrary_commands->is_defined()) {
-        fprintf(stderr, "error: cluster mode dose not support arbitrary command option.\n");
-        return false;
     }
 
     return true;

run_stats.cpp

@@ -196,6 +196,17 @@ void run_stats::update_moved_set_op(struct timeval* ts, unsigned int bytes, unsi
     hdr_record_value(m_set_latency_histogram,latency);
 }
 
+void run_stats::update_moved_arbitrary_op(struct timeval *ts, unsigned int bytes,
+                                    unsigned int latency, size_t request_index) {
+    roll_cur_stats(ts);
+
+    m_cur_stats.m_ar_commands.at(request_index).update_moved_op(bytes, latency);
+    m_totals.update_op(bytes, latency);
+
+    struct hdr_histogram* hist = m_ar_commands_latency_histograms.at(request_index);
+    hdr_record_value(hist,latency);
+}
+
 void run_stats::update_ask_get_op(struct timeval* ts, unsigned int bytes, unsigned int latency)
 {
     roll_cur_stats(ts);
@@ -214,6 +225,17 @@ void run_stats::update_ask_set_op(struct timeval* ts, unsigned int bytes, unsign
     hdr_record_value(m_set_latency_histogram,latency);
 }
 
+void run_stats::update_ask_arbitrary_op(struct timeval *ts, unsigned int bytes,
+                                          unsigned int latency, size_t request_index) {
+    roll_cur_stats(ts);
+
+    m_cur_stats.m_ar_commands.at(request_index).update_ask_op(bytes, latency);
+    m_totals.update_op(bytes, latency);
+
+    struct hdr_histogram* hist = m_ar_commands_latency_histograms.at(request_index);
+    hdr_record_value(hist,latency);
+}
+
 void run_stats::update_wait_op(struct timeval *ts, unsigned int latency)
 {
     roll_cur_stats(ts);
@@ -975,11 +997,18 @@ void run_stats::print_moved_sec_column(output_table &table) {
 
     column.elements.push_back(*el.init_str("%12s ", "MOVED/sec"));
     column.elements.push_back(*el.init_str("%s", "-------------"));
-    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_set_cmd.m_moved_sec));
-    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_get_cmd.m_moved_sec));
-    column.elements.push_back(*el.init_str("%12s ", "---"));
-    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_moved_sec));
 
+    if (print_arbitrary_commands_results()) {
+        for (unsigned int i=0; i<m_totals.m_ar_commands.size(); i++) {
+            column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_ar_commands[i].m_moved_sec));
+        }
+    } else {
+        column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_set_cmd.m_moved_sec));
+        column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_get_cmd.m_moved_sec));
+        column.elements.push_back(*el.init_str("%12s ", "---"));
+
+    }
+    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_moved_sec));
     table.add_column(column);
 }
 
@@ -989,11 +1018,17 @@ void run_stats::print_ask_sec_column(output_table &table) {
 
     column.elements.push_back(*el.init_str("%12s ", "ASK/sec"));
     column.elements.push_back(*el.init_str("%s", "-------------"));
-    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_set_cmd.m_ask_sec));
+    if (print_arbitrary_commands_results()) {
+        for (unsigned int i=0; i<m_totals.m_ar_commands.size(); i++) {
+            column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_ar_commands[i].m_ask_sec));
+        }
+    } else {
+        column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_set_cmd.m_ask_sec));
     column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_get_cmd.m_ask_sec));
     column.elements.push_back(*el.init_str("%12s ", "---"));
-    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_ask_sec));
 
+    }
+    column.elements.push_back(*el.init_double("%12.2f ", m_totals.m_ask_sec));
     table.add_column(column);
 }
 

run_stats.h

@@ -119,9 +119,13 @@ class run_stats {
 
     void update_moved_get_op(struct timeval* ts, unsigned int bytes, unsigned int latency);
     void update_moved_set_op(struct timeval* ts, unsigned int bytes, unsigned int latency);
+    void update_moved_arbitrary_op(struct timeval *ts, unsigned int bytes,
+                             unsigned int latency, size_t arbitrary_index);
 
     void update_ask_get_op(struct timeval* ts, unsigned int bytes, unsigned int latency);
     void update_ask_set_op(struct timeval* ts, unsigned int bytes, unsigned int latency);
+    void update_ask_arbitrary_op(struct timeval *ts, unsigned int bytes,
+                                   unsigned int latency, size_t arbitrary_index);
 
     void update_wait_op(struct timeval* ts, unsigned int latency);
     void update_arbitrary_op(struct timeval *ts, unsigned int bytes,