Performance test for batch regression

batch_perf.c

/*
This is free and unencumbered software released into the public domain.

Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.

In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

For more information, please refer to <http://unlicense.org/>
*/

#include <assert.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include <uv.h>

#ifdef WIN32
#include <Windows.h>
#else
#include <unistd.h>
#endif

#include "cassandra.h"

/*
* Use this example with caution. It's just used as a scratch example for debugging and
* roughly analyzing performance.
*/

/*
 * NOTE: To run this example you'll need to increase these numbers in your cassandra.yaml file:
 * batch_size_fail_threshold_in_kb: 9999
 * batch_size_warn_threshold_in_kb: 9999
 */

#define USE_PREPARED 1
#define USE_SSL 0

#define NUM_THREADS 1
#define NUM_IO_WORKER_THREADS 1
#define NUM_ITERATIONS 450
#define NUM_CONCURRENT_REQUESTS 1
#define BATCH_SIZE 1000

#define MAX_NUM_THREADS 16
#define MAX_NUM_CONCURRENT_REQUESTS 15000


const char* one_byte = "a";

const char* one_kb = "0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567"
"0123456701234567012345670123456701234567012345670123456701234567";

const char* one_sixteenth_kb = "0123456701234567012345670123456701234567012345670123456701234567";

#if CASS_VERSION_MAJOR >= 2
#define STRING_PARAM(s) s
#else
#define STRING_PARAM(s) cass_string_init(s)
#endif

CassUuidGen* uuid_gen;

int num_threads = NUM_THREADS;
int num_io_threads = NUM_IO_WORKER_THREADS;
int num_iterations = NUM_ITERATIONS;
int num_concurrent_requests = NUM_CONCURRENT_REQUESTS;
int batch_size = BATCH_SIZE;
int protocol_version = 4;
int use_token_aware = 0;

typedef struct Status_ {
  uv_mutex_t mutex;
  uv_cond_t cond;
  int count;
} Status;

Status status;

int status_init(Status* status, int initial_count) {
  int rc;
  rc = uv_mutex_init(&status->mutex);
  if (rc != 0) return rc;
  rc = uv_cond_init(&status->cond);
  if (rc != 0) return rc;
  status->count = initial_count;
  return rc;
}

void status_destroy(Status* status) {
  uv_mutex_destroy(&status->mutex);
  uv_cond_destroy(&status->cond);
}

void status_notify(Status* status) {
  uv_mutex_lock(&status->mutex);
  status->count--;
  uv_cond_signal(&status->cond);
  uv_mutex_unlock(&status->mutex);
}

int status_wait(Status* status, uint64_t timeout_secs) {
  int count;
  uv_mutex_lock(&status->mutex);
  uv_cond_timedwait(&status->cond, &status->mutex, timeout_secs * 1000 * 1000 * 1000);
  count = status->count;
  uv_mutex_unlock(&status->mutex);
  return count;
}

#if USE_SSL
int load_trusted_cert_file(const char* file, CassSsl* ssl) {
  CassError rc;
  char* cert;
  long cert_size;
  size_t bytes_read;

  FILE *in = fopen(file, "rb");
  if (in == NULL) {
    fprintf(stderr, "Error loading certificate file '%s'\n", file);
    return 0;
  }

  fseek(in, 0, SEEK_END);
  cert_size = ftell(in);
  rewind(in);

  cert = (char*)malloc(cert_size);
  bytes_read = fread(cert, 1, cert_size, in);
  fclose(in);

  if (bytes_read == (size_t)cert_size) {
    rc = cass_ssl_add_trusted_cert_n(ssl, cert, cert_size);
    if (rc != CASS_OK) {
      fprintf(stderr, "Error loading SSL certificate: %s\n", cass_error_desc(rc));
      free(cert);
      return 0;
    }
  }

  free(cert);
  return 1;
}
#endif

void print_error(CassFuture* future) {
#if CASS_VERSION_MAJOR >= 2
  const char* message;
  size_t message_length;
  cass_future_error_message(future, &message, &message_length);
  fprintf(stderr, "Error: %.*s\n", (int)message_length, message);
#else
  CassString message = cass_future_error_message(future);
  fprintf(stderr, "Error: %.*s\n", (int)message.length, message.data);
#endif
}

CassCluster* create_cluster(const char* hosts, const char* trusted_cert_file) {
  CassCluster* cluster = cass_cluster_new();

  cass_cluster_set_contact_points(cluster, hosts);

#if 0
  CassRetryPolicy* retry_policy = cass_retry_policy_fallthrough_new();
  cass_cluster_set_retry_policy(cluster, retry_policy);
  cass_retry_policy_free(retry_policy);

  // This is unnecessary, the optimial protocol version is negotiated
  //cass_cluster_set_protocol_version(cluster, 2);

  cass_cluster_set_latency_aware_routing(cluster, cass_true);
#endif


  cass_cluster_set_connect_timeout(cluster, 10000);
  cass_cluster_set_reconnect_wait_time(cluster, 5000);
  cass_cluster_set_tcp_keepalive(cluster, cass_true, 15);

#if USE_SSL
  CassSsl* ssl = cass_ssl_new();
  cass_ssl_set_verify_flags(ssl, CASS_SSL_VERIFY_PEER_CERT);
  if (!load_trusted_cert_file(trusted_cert_file, ssl)) {
    fprintf(stderr, "Failed to load certificate '%s' disabling peer verification\n", trusted_cert_file);
    cass_ssl_set_verify_flags(ssl, CASS_SSL_VERIFY_NONE);
  }
  cass_cluster_set_ssl(cluster, ssl);
  cass_ssl_free(ssl);
#endif

  if (cass_cluster_set_protocol_version(cluster, protocol_version) != CASS_OK) {
    fprintf(stderr, "protocol version %d not supported using default\n", protocol_version);
  }
  cass_cluster_set_token_aware_routing(cluster, use_token_aware ? cass_true : cass_false);
  cass_cluster_set_num_threads_io(cluster, num_io_threads);
  cass_cluster_set_queue_size_io(cluster, NUM_CONCURRENT_REQUESTS * NUM_THREADS);
  cass_cluster_set_pending_requests_high_water_mark(cluster, NUM_CONCURRENT_REQUESTS * NUM_THREADS);
  cass_cluster_set_write_bytes_high_water_mark(cluster, 16 * 1048576);

  return cluster;
}

CassError connect_session(CassSession* session, const CassCluster* cluster) {
  CassError rc = CASS_OK;
  CassFuture* future = cass_session_connect(session, cluster);

  cass_future_wait(future);
  rc = cass_future_error_code(future);
  if (rc != CASS_OK) {
    print_error(future);
  }
  cass_future_free(future);

  return rc;
}

CassError execute_query(CassSession* session, const char* query) {
  CassError rc = CASS_OK;
  CassFuture* future = NULL;
  CassStatement* statement = cass_statement_new(STRING_PARAM(query), 0);

  future = cass_session_execute(session, statement);
  cass_future_wait(future);

  rc = cass_future_error_code(future);
  if (rc != CASS_OK) {
    print_error(future);
  }

  cass_future_free(future);
  cass_statement_free(statement);

  return rc;
}

CassError prepare_query(CassSession* session, const char* query, const CassPrepared** prepared) {
  CassError rc = CASS_OK;
  CassFuture* future = NULL;

  future = cass_session_prepare(session, STRING_PARAM(query));
  cass_future_wait(future);

  rc = cass_future_error_code(future);
  if (rc != CASS_OK) {
    print_error(future);
  }
  else {
    *prepared = cass_future_get_prepared(future);
  }

  cass_future_free(future);

  return rc;
}

int compare_dbl(const void* d1, const void* d2) {
  if (*((double*)d1) < *((double*)d2)) {
    return -1;
  }
  else if (*((double*)d1) > *((double*)d2)) {
    return 1;
  }
  else {
    return 0;
  }
}

void insert_into_perf(CassSession* session, const char* query, const CassPrepared* prepared) {
  int i;
  CassFuture* futures[MAX_NUM_CONCURRENT_REQUESTS];

  for (i = 0; i < num_concurrent_requests; ++i) {
    int j, c;
    CassUuid id;
    CassBatch* batch = cass_batch_new(CASS_BATCH_TYPE_UNLOGGED);

    for (j = 0; j < batch_size; ++j) {
      CassStatement* statement;

      if (prepared != NULL) {
        statement = cass_prepared_bind(prepared);
      }
      else {
        statement = cass_statement_new(STRING_PARAM(query), 17);
      }

      cass_uuid_gen_time(uuid_gen, &id);
      cass_statement_bind_uuid(statement, 0, id);
      for (c = 1; c <= 16; ++c) {
        cass_statement_bind_string(statement, c, STRING_PARAM(one_sixteenth_kb));
      }

      cass_batch_add_statement(batch, statement);

      cass_statement_free(statement);
    }

    futures[i] = cass_session_execute_batch(session, batch);

    cass_batch_free(batch);
  }

  for (i = 0; i < num_concurrent_requests; ++i) {
    CassFuture* future = futures[i];
    CassError rc = cass_future_error_code(future);
    if (rc != CASS_OK) {
      print_error(future);
    }
    cass_future_free(future);
  }
}

void run_insert_queries(void* data) {
  int i;
  CassSession* session = (CassSession*)data;

  const CassPrepared* insert_prepared = NULL;
  const char* insert_query = "INSERT INTO stress.stress (k, "
    "c1, c2,  c3,  c4,  c5,  c6,  c7,  c8, "
    "c9, c10, c11, c12, c13, c14, c15, c16)  "
    "VALUES (?, "
    "?, ?, ?, ?, ?, ?, ?, ?, "
    "?, ?, ?, ?, ?, ?, ?, ?)";

#if USE_PREPARED
  if (prepare_query(session, insert_query, &insert_prepared) == CASS_OK) {
#endif
    for (i = 0; i < num_iterations; ++i) {
      insert_into_perf(session, insert_query, insert_prepared);
    }
#if USE_PREPARED
    cass_prepared_free(insert_prepared);
  }
#endif

  status_notify(&status);
}

#define CHECK_ARG(name) do { \
if (i + 1 > argc) {\
  fprintf(stderr, #name " expects an argument\n"); \
  exit(-1); \
} \
} while (0)

int main(int argc, char* argv[]) {
  int i;
  uv_thread_t threads[MAX_NUM_THREADS];
  CassCluster* cluster = NULL;
  CassSession* session = NULL;
  const char* hosts = "127.0.0.1";
  const char* trusted_cert_file = "trusted_cert.pem"; // Provide either the cert or the CA chain

  for (i = 1; i < argc; ++i) {
    char* arg = argv[i];
    if (strcmp(arg, "--hosts") == 0) {
      CHECK_ARG("--hosts");
      hosts = argv[i + 1];
      i++;
    }
    else if (strcmp(arg, "--num-threads") == 0) {
      CHECK_ARG("--num-threads");
      num_threads = atoi(argv[i + 1]);
      if (num_threads > MAX_NUM_THREADS || num_threads <= 0) {
        fprintf(stderr, "--num-threads has the invalid value %d\n", num_threads);
        exit(-1);
      }
      i++;
    }
    else if (strcmp(arg, "--num-io-threads") == 0) {
      CHECK_ARG("--num-io-threads");
      num_io_threads = atoi(argv[i + 1]);
      if (num_io_threads <= 0) {
        fprintf(stderr, "--num-io-threads has the invalid value %d\n", num_io_threads);
        exit(-1);
      }
      i++;
    }
    else if (strcmp(arg, "--num-iterations") == 0) {
      CHECK_ARG("--num-iterations");
      num_iterations = atoi(argv[i + 1]);
      if (num_iterations <= 0) {
        fprintf(stderr, "--num-iterations has the invalid value %d\n", num_iterations);
        exit(-1);
      }
      i++;
    }
    else if (strcmp(arg, "--num-concurrent-requests") == 0) {
      CHECK_ARG("--num-concurrent-requests");
      num_concurrent_requests = atoi(argv[i + 1]);
      if (num_concurrent_requests > MAX_NUM_CONCURRENT_REQUESTS || num_concurrent_requests <= 0) {
        fprintf(stderr, "--num-concurrent-requests has the invalid value %d\n", num_concurrent_requests);
        exit(-1);
      }
      i++;
    }
    else if (strcmp(arg, "--batch-size") == 0) {
      CHECK_ARG("--batch-size");
      batch_size = atoi(argv[i + 1]);
      if (batch_size <= 0) {
        fprintf(stderr, "--batch-size has the invalid value %d\n", batch_size);
        exit(-1);
      }
      i++;
    }
    else if (strcmp(arg, "--protocol-version") == 0) {
      CHECK_ARG("--protocol-version");
      protocol_version = atoi(argv[i + 1]);
      if (protocol_version <= 0) {
        fprintf(stderr, "--protocol-version has the invalid value %d\n", protocol_version);
        exit(-1);
      }
      i++;
    }
    else if (strcmp(arg, "--use-token-aware") == 0) {
      CHECK_ARG("--use-token-aware");
      use_token_aware = atoi(argv[i + 1]);
      i++;
    }
    else if (strcmp(arg, "--trust-cert-file") == 0) {
      CHECK_ARG("--trust-cert-file");
      trusted_cert_file = argv[i + 1];
      i++;
    }
    else {
      fprintf(stderr, "%s is an invalid flags\n", arg);
      exit(-1);
    }
  }

  printf("running with driver version %d.%d.%d, hosts \"%s\", protocol version %d, use_token_aware %s, %d threads, %d iterations, %d concurrent requests, %d batch size, and %d io workers threads\n",
    CASS_VERSION_MAJOR, CASS_VERSION_MINOR, CASS_VERSION_PATCH, hosts, protocol_version, use_token_aware ? "true" : "false",
    num_threads, num_iterations, num_concurrent_requests, batch_size, num_io_threads);

  status_init(&status, num_threads);

  cass_log_set_level(CASS_LOG_WARN);

  cluster = create_cluster(hosts, trusted_cert_file);
  uuid_gen = cass_uuid_gen_new();
  session = cass_session_new();

  if (connect_session(session, cluster) != CASS_OK) {
    cass_cluster_free(cluster);
    cass_session_free(session);
    return -1;
  }

  execute_query(session, "CREATE KEYSPACE IF NOT EXISTS stress "
    "WITH replication = { 'class': 'SimpleStrategy', 'replication_factor': '1' }");

  execute_query(session, "CREATE TABLE IF NOT EXISTS stress.stress "
    "(k uuid PRIMARY KEY, "
    "c1  text, c2  text, c3  text, c4  text, "
    "c5  text, c6  text, c7  text, c8  text, "
    "c9  text, c10 text, c11 text, c12 text, "
    "c13 text, c14 text, c15 text, c16 text)");

  execute_query(session, "TRUNCATE stress.stress");

#ifdef WIN32
  Sleep(3000);
#else
  sleep(3);
#endif

  uint64_t start = uv_hrtime();

  for (i = 0; i < num_threads; ++i) {
    uv_thread_create(&threads[i], run_insert_queries, (void*)session);
  }

  while (status_wait(&status, 5) > 0) {
#if CASS_VERSION_MAJOR >= 2
    CassMetrics metrics;
    cass_session_get_metrics(session, &metrics);
    printf("rate stats (requests/second): mean %f 1m %f 5m %f 10m %f\n",
      metrics.requests.mean_rate,
      metrics.requests.one_minute_rate,
      metrics.requests.five_minute_rate,
      metrics.requests.fifteen_minute_rate);
    printf("stats (microseconds): min %llu max %llu median %llu 75th %llu 95th %llu 98th %llu 99th %llu 99.9th %llu\n",
      (unsigned long long int)metrics.requests.min, (unsigned long long int)metrics.requests.max,
      (unsigned long long int)metrics.requests.median, (unsigned long long int)metrics.requests.percentile_75th,
      (unsigned long long int)metrics.requests.percentile_95th, (unsigned long long int)metrics.requests.percentile_98th,
      (unsigned long long int)metrics.requests.percentile_99th, (unsigned long long int)metrics.requests.percentile_999th);
#endif
  }

  printf("inserted %d rows in %f seconds\n",
    num_threads * num_iterations *  num_concurrent_requests * batch_size,
    (uv_hrtime() - start) / (1000.0 * 1000.0 * 1000.0));

  for (i = 0; i < num_threads; ++i) {
    uv_thread_join(&threads[i]);
  }

  cass_cluster_free(cluster);
  cass_session_free(session);
  cass_uuid_gen_free(uuid_gen);

  status_destroy(&status);

  return 0;
}