in-memory-index-rbtree.c

#include "headers.h"
#include "indexes/rbtree.h"

static uint64_t get_prefix_for_item(char *item) {
  struct item_metadata *meta = (struct item_metadata *)item;
  char *item_key = &item[sizeof(*meta)];
  return *(uint64_t *)item_key;
}

/* In memory RB-Tree */
static rbtree *items_locations;
static pthread_spinlock_t *items_location_locks;
index_entry_t *rbtree_worker_lookup(int worker_id, void *item) {
  return rbtree_lookup(items_locations[worker_id],
                       (void *)get_prefix_for_item(item), pointer_cmp);
}
void rbtree_worker_insert(int worker_id, void *item, index_entry_t *e) {
  pthread_spin_lock(&items_location_locks[worker_id]);
  rbtree_insert(items_locations[worker_id], (void *)get_prefix_for_item(item),
                e, pointer_cmp);
  pthread_spin_unlock(&items_location_locks[worker_id]);
}
void rbtree_worker_delete(int worker_id, void *item) {
  pthread_spin_lock(&items_location_locks[worker_id]);
  rbtree_delete(items_locations[worker_id], (void *)get_prefix_for_item(item),
                pointer_cmp);
  pthread_spin_unlock(&items_location_locks[worker_id]);
}

void rbtree_index_add(struct slab_callback *cb, void *item) {
  index_entry_t e = {
      .slab = cb->slab,
      .slab_idx = cb->slab_idx,
  };
  rbtree_worker_insert(get_worker(e.slab), item, &e);
}

/*
 * Returns up to scan_size keys >= item.key.
 * If item is not in the database, this will still return up to scan_size keys >
 * item.key.
 */
struct index_scan rbtree_init_scan(void *item, size_t scan_size) {
  size_t nb_workers = get_nb_workers();

  struct rbtree_scan_tmp *res = malloc(nb_workers * sizeof(*res));
  for (size_t w = 0; w < nb_workers; w++) {
    pthread_spin_lock(&items_location_locks[w]);
    res[w] =
        rbtree_lookup_n(items_locations[w], (void *)get_prefix_for_item(item),
                        scan_size, pointer_cmp);
    pthread_spin_unlock(&items_location_locks[w]);
  }

  struct index_scan scan_res;
  scan_res.entries = malloc(scan_size * sizeof(*scan_res.entries));
  scan_res.hashes = malloc(scan_size * sizeof(*scan_res.hashes));
  scan_res.nb_entries = 0;

  size_t *positions = calloc(nb_workers, sizeof(*positions));
  while (scan_res.nb_entries < scan_size) {
    size_t min_worker = nb_workers;
    struct rbtree_node_t *min = NULL;
    for (size_t w = 0; w < nb_workers; w++) {
      if (res[w].nb_entries <= positions[w]) {
        continue; // no more item to read in that rbtree
      } else {
        struct rbtree_node_t *current = &res[w].entries[positions[w]];
        if (!min || pointer_cmp(min->key, current->key) > 0) {
          min = current;
          min_worker = w;
        }
      }
    }
    if (min_worker == nb_workers)
      break; // no worker has any scannable item left
    positions[min_worker]++;
    scan_res.entries[scan_res.nb_entries] = min->value;
    scan_res.hashes[scan_res.nb_entries] = (uint64_t)min->key;
    scan_res.nb_entries++;
  }
  for (size_t w = 0; w < nb_workers; w++) {
    free(res[w].entries);
  }
  free(res);
  free(positions);
  return scan_res;
}

void rbtree_init(void) {
  items_locations = malloc(get_nb_workers() * sizeof(*items_locations));
  items_location_locks =
      malloc(get_nb_workers() * sizeof(*items_location_locks));
  for (size_t w = 0; w < get_nb_workers(); w++) {
    items_locations[w] = rbtree_create();
    pthread_spin_init(&items_location_locks[w], PTHREAD_PROCESS_PRIVATE);
  }
}