workload-common.c

#include "headers.h"

/*
 * Create a workload item for the database
 */
char *create_unique_item(size_t item_size, uint64_t uid) {
   char *item = malloc(item_size);
   struct item_metadata *meta = (struct item_metadata *)item;
   meta->key_size = 8;
   meta->value_size = item_size - 8 - sizeof(*meta);

   char *item_key = &item[sizeof(*meta)];
   char *item_value = &item[sizeof(*meta) + meta->key_size];
   *(uint64_t*)item_key = uid;
   *(uint64_t*)item_value = uid;
   return item;
}

/* We also store an item in the database that says if the database has been populated for YCSB, PRODUCTION, or another workload. */
char *create_workload_item(struct workload *w) {
   const uint64_t key = -10;
   const char *name = w->api->api_name(); // YCSB or PRODUCTION?
   size_t key_size = 16;
   size_t value_size = strlen(name) + 1;

   struct item_metadata *meta;
   char *item = malloc(sizeof(*meta) + key_size + value_size);
   meta = (struct item_metadata *)item;
   meta->key_size = key_size;
   meta->value_size = value_size;

   char *item_key = &item[sizeof(*meta)];
   char *item_value = &item[sizeof(*meta) + meta->key_size];
   *(uint64_t*)item_key = key;
   strcpy(item_value, name);
   return item;
}

/*
 * Fill the DB with missing items
 */
static void add_in_tree(struct slab_callback *cb, void *item) {
   memory_index_add(cb, item);
   free(cb->item);
   free(cb);
}

struct rebuild_pdata {
   size_t id;
   size_t *pos;
   size_t start;
   size_t end;
   struct workload *w;
};
void *repopulate_db_worker(void *pdata) {
   declare_periodic_count;
   struct rebuild_pdata *data = pdata;

   pin_me_on(get_nb_workers() + data->id);

   size_t *pos = data->pos;
   struct workload *w = data->w;
   struct workload_api *api = w->api;
   size_t start = data->start;
   size_t end = data->end;
   for(size_t i = start; i < end; i++) {
      struct slab_callback *cb = malloc(sizeof(*cb));
      cb->cb = add_in_tree;
      cb->payload = NULL;
      cb->item = api->create_unique_item(pos[i], w->nb_items_in_db);
      kv_add_async(cb);
      periodic_count(1000, "Repopulating database (%lu%%)", 100LU-(end-i)*100LU/(end - start));
   }

   return NULL;
}

void repopulate_db(struct workload *w) {
   declare_timer;
   void *workload_item = create_workload_item(w);
   int64_t nb_inserts = (get_database_size() > w->nb_items_in_db)?0:(w->nb_items_in_db - get_database_size());


   if(nb_inserts != w->nb_items_in_db) { // Database at least partially populated
      // Check that the items correspond to the workload
      char *db_item = kv_read_sync(workload_item);
      if(!db_item)
         die("Running a benchmark on a pre-populated DB, but couldn't determine if items in the DB correspond to the benchmark --- please wipe DB before benching!\n");
      struct item_metadata *meta = (struct item_metadata *)db_item;
      char *item_value = &db_item[sizeof(*meta) + meta->key_size];
      if(strcmp(w->api->api_name(), item_value))
         die("Running %s benchmark, but the database contains elements from %s benchmark -- please wipe DB before benching!\n", w->api->api_name(), item_value);
   }

   if(nb_inserts == 0) {
      free(workload_item);
      return;
   }

   uint64_t nb_items_already_in_db = get_database_size();

   // Say that this database is for that workload
   if(nb_items_already_in_db == 0) {
      struct slab_callback *cb = malloc(sizeof(*cb));
      cb->cb = add_in_tree;
      cb->payload = NULL;
      cb->item = workload_item;
      kv_add_async(cb);
   } else {
      nb_items_already_in_db--; // do not count the workload_item
   }
   if(nb_items_already_in_db != 0 && nb_items_already_in_db != w->nb_items_in_db) {
      /*
       * Because we shuffle elements, we don't really want to start with a small database and have all the higher order elements at the end, that would be cheating.
       * Plus, we insert database items at random positions (see shuffle below) and I am too lazy to implement the logic of doing the shuffle minus existing elements.
       */
      die("The database contains %lu elements but the benchmark is configured to use %lu. Please delete the DB first.\n", nb_items_already_in_db, w->nb_items_in_db);
   }


   size_t *pos = NULL;
   start_timer {
      printf("Initializing big array to insert elements in random order... This might take a while. (Feel free to comment but then the database will be sorted and scans much faster -- unfair vs other systems)\n");
      pos = malloc(w->nb_items_in_db * sizeof(*pos));
      for(size_t i = 0; i < w->nb_items_in_db; i++)
         pos[i] = i;
      shuffle(pos, nb_inserts); // To be fair to other systems, we shuffle items in the DB so that the DB is not fully sorted by luck
   } stop_timer("Big array of random positions");

   start_timer {
      struct rebuild_pdata *pdata = malloc(w->nb_load_injectors*sizeof(*pdata));
      pthread_t *threads = malloc(w->nb_load_injectors*sizeof(*threads));
      for(size_t i = 0; i < w->nb_load_injectors; i++) {
         pdata[i].id = i;
         pdata[i].start = (w->nb_items_in_db / w->nb_load_injectors)*i;
         pdata[i].end = (w->nb_items_in_db / w->nb_load_injectors)*(i+1);
         if(i == w->nb_load_injectors - 1)
            pdata[i].end = w->nb_items_in_db;
         pdata[i].w = w;
         pdata[i].pos = pos;
         if(i)
            pthread_create(&threads[i], NULL, repopulate_db_worker, &pdata[i]);
      }
      repopulate_db_worker(&pdata[0]);
      for(size_t i = 1; i < w->nb_load_injectors; i++)
         pthread_join(threads[i], NULL);
      free(threads);
      free(pdata);
   } stop_timer("Repopulating %lu elements (%lu req/s)", nb_inserts, nb_inserts*1000000/elapsed);

   free(pos);
}

/*
 *  Print an item stored on disk
 */
void print_item(size_t idx, void* _item) {
   char *item = _item;
   struct item_metadata *meta = (struct item_metadata *)item;
   char *item_key = &item[sizeof(*meta)];
   if(meta->key_size == 0)
      printf("[%lu] Non existant?\n", idx);
   else if(meta->key_size == -1)
      printf("[%lu] Removed\n", idx);
   else
      printf("[%lu] K=%lu V=%s\n", idx, *(uint64_t*)item_key, &item[sizeof(*meta) + meta->key_size]);
}

/*
 * Various callbacks that are called once an item has been read / written
 */
void show_item(struct slab_callback *cb, void *item) {
   print_item(cb->slab_idx, item);
   free(cb->item);
   free(cb);
}

void free_callback(struct slab_callback *cb, void *item) {
   free(cb->item);
   free(cb);
}

void compute_stats(struct slab_callback *cb, void *item) {
   uint64_t start, end;
   declare_debug_timer;
   start_debug_timer {
      start = get_time_from_payload(cb, 0);
      rdtscll(end);
      add_timing_stat(end - start);
      if(DEBUG && cycles_to_us(end-start) > 10000) { // request took more than 10ms
         printf("Request [%lu: %lu] [%lu: %lu] [%lu: %lu] [%lu: %lu] [%lu: %lu] [%lu: %lu] [%lu: %lu] [%lu]\n",
               get_origin_from_payload(cb, 1), get_time_from_payload(cb, 1) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 1) - start),
               get_origin_from_payload(cb, 2), get_time_from_payload(cb, 2) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 2) - start),
               get_origin_from_payload(cb, 3), get_time_from_payload(cb, 3) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 3) - start),
               get_origin_from_payload(cb, 4), get_time_from_payload(cb, 4) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 4) - start),
               get_origin_from_payload(cb, 5), get_time_from_payload(cb, 5) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 5) - start),
               get_origin_from_payload(cb, 6), get_time_from_payload(cb, 6) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 6) - start),
               get_origin_from_payload(cb, 7), get_time_from_payload(cb, 7) < start ? 0 : cycles_to_us(get_time_from_payload(cb, 7) - start),
               cycles_to_us(end  - start));
      }
      free(cb->item);
      if(DEBUG)
         free_payload(cb);
      free(cb);
   } stop_debug_timer(5000, "Callback took more than 5ms???");
}

struct slab_callback *bench_cb(void) {
   struct slab_callback *cb = malloc(sizeof(*cb));
   cb->cb = compute_stats;
   cb->payload = allocate_payload();
   return cb;
}


/*
 * Generic worklad API.
 */
struct thread_data {
   size_t id;
   struct workload *workload;
   bench_t benchmark;
};

struct workload_api *get_api(bench_t b) {
   if(YCSB.handles(b))
      return &YCSB;
   if(PRODUCTION.handles(b))
      return &PRODUCTION;
   die("Unknown workload for benchmark!\n");
}

static pthread_barrier_t barrier;
void* do_workload_thread(void *pdata) {
   struct thread_data *d = pdata;

   init_seed();
   pin_me_on(get_nb_workers() + d->id);
   pthread_barrier_wait(&barrier);

   d->workload->api->launch(d->workload, d->benchmark);

   return NULL;
}

void run_workload(struct workload *w, bench_t b) {
   struct thread_data *pdata = malloc(w->nb_load_injectors*sizeof(*pdata));

   w->nb_requests_per_thread = w->nb_requests / w->nb_load_injectors;
   pthread_barrier_init(&barrier, NULL, w->nb_load_injectors);

   if(!w->api->handles(b))
      die("The database has not been configured to run this benchmark! (Are you trying to run a production benchmark on a database configured for YCSB?)");

   declare_timer;
   start_timer {
      pthread_t *threads = malloc(w->nb_load_injectors*sizeof(*threads));
      for(int i = 0; i < w->nb_load_injectors; i++) {
         pdata[i].id = i;
         pdata[i].workload = w;
         pdata[i].benchmark = b;
         if(i)
            pthread_create(&threads[i], NULL, do_workload_thread, &pdata[i]);
      }
      do_workload_thread(&pdata[0]);
      for(int i = 1; i < w->nb_load_injectors; i++)
         pthread_join(threads[i], NULL);
      free(threads);
   } stop_timer("%s - %lu requests (%lu req/s)", w->api->name(b), w->nb_requests, w->nb_requests*1000000/elapsed);
   print_stats();

   free(pdata);
}