ibkvcli.c

/*
 * Copyright (c) 2016--2021  Wu, Xingbo <wuxb45@gmail.com>
 *
 * All rights reserved. No warranty, explicit or implicit, provided.
 */
#define _GNU_SOURCE

#include "lib.h"
#include "kv.h"
#include "ib.h"

#define XSA ((0))
#define XDA ((1))
#define XGA ((2))
#define XPA ((3))
#define XSS ((4))
#define XDS ((5))
#define XGS ((6))
#define XPS ((7))
#define XGO ((8)) // value not received: buffer of out space
#define NR_WARGS ((4))

  static bool
ibkv_analyze(void * const passdata[2], const u64 dt, const struct vctr * const va, struct damp * const d, char * const out)
{
  (void)passdata;
  u64 v[9];
  for (u64 i = 0; i < 9; i++)
    v[i] = vctr_get(va, i);

  const u64 nrop = v[XSA] + v[XDA] + v[XGA] + v[XPA];
  const double mops = ((double)nrop) * 1e3 / ((double)dt);
  const bool done = damp_add_test(d, mops);
  const char * const pat = " set %lu %lu del %lu %lu get %lu %lu (%lu) pro %lu %lu mops %.4lf avg %.4lf avg %.4lf\n";
  sprintf(out, pat, v[XSA], v[XSS], v[XDA], v[XDS], v[XGA], v[XGS], v[XGO], v[XPA], v[XPS], mops, damp_avg(d), damp_ravg(d));
  return done;
}

struct worker_info {
  struct ib_port * ibport;
  const char * host;
  const char * port;
  u64 depth;
  u64 batch;
};

struct thread_info {
  u64 cpro;
  u64 cget;
  u64 cset;
  u32 vlen;
};

struct kvbuf {
  struct kv * kv;
  u32 kvsz;
  u32 asz;
  u8 op;
};

  static u32
fill_requests(u8 * const rbuf, const u32 maxsz, u32 batch, const struct forker_worker_info * const mi)
{
  rgen_next_func next = mi->rgen_next;
  const struct thread_info * const ti = (typeof(ti))mi->priv;
  const u64 cpro = ti->cpro;
  const u64 cget = ti->cget;
  const u64 cset = ti->cset;
  struct kvbuf kvbuf[batch];
  u64 totsz = sizeof(u32);

  u64 idp = next(mi->gen); // 1st
  struct kv * kvp = kvss[idp];
  __builtin_prefetch(kvp, 0, 0);

  idp = next(mi->gen);
  __builtin_prefetch(&(kvss[idp]), 0, 0);

  for (u32 i = 0; i < batch; i++) {
    struct kv * const kv = kvp;
    kvp = kvss[idp];
    __builtin_prefetch(kvp, 0, 0);
    if ((i + 2lu) < batch) {
      idp = next(mi->gen);
      __builtin_prefetch(&(kvss[idp]), 0, 0);
    }
    //const u64 x = next(mi->gen);
    const u64 y = random_u64() % 100lu;
    u8 op;
    if (y < cpro) {
      op = 'P';
    } else if (y < cget) {
      op = 'G';
    } else if (y < cset) {
      op = 'S';
    } else {
      op = 'D';
    }
    kv->vlen = (op == 'S') ? ti->vlen : 0;
    const u32 kvsz = kv_size(kv);
    const u32 asz = bits_round_up(kvsz, 2);
    if ((totsz + asz + sizeof(u32)) > maxsz) {
      batch = i;
      break;
    }
    kvbuf[i].kv = kv;
    kvbuf[i].kvsz = kvsz;
    kvbuf[i].asz = asz;
    kvbuf[i].op = op;
    totsz += (asz + sizeof(u32));
  }
  u32 off = (batch+1) * sizeof(u32);
  u32 * const pm = (typeof(pm))rbuf;
  pm[0] = batch;
  for (u32 i = 0; i < batch; i++) {
    pm[i+1] = (off << 8) | kvbuf[i].op;
    kvbuf[i].kv->vlen = (kvbuf[i].op == 'S') ? ti->vlen : 0;
    memcpy(rbuf+off, kvbuf[i].kv, kvbuf[i].kvsz);
    off += kvbuf[i].asz;
  }
  off = bits_round_up(off, 6); // x64
  return off;
}

  static void
process_results(u8 * const buf_res, struct vctr * const vctr)
{
  const u32 * const hdrs = (typeof(hdrs))buf_res;
  const u32 n = hdrs[0];
  //for (u32 i = 0; i <= n; i++)
  //  printf("%u %u 0x%x\n", i, hdrs[i], hdrs[i]);
  for (u32 i = 1; i <= n; i++) {
    const u8 op = (u8)hdrs[i];
    if (op == 'P') {
      vctr_add1(vctr, XPA);
      if (hdrs[i] >> 8)
        vctr_add1(vctr, XPS);
    } else if (op == 'G') {
      vctr_add1(vctr, XGA);
      // TODO: check kv
      const u32 magic = hdrs[i] >> 8;
      if (magic) {
        vctr_add1(vctr, XGS);
        if (magic == 1)
          vctr_add1(vctr, XGO);
      }
    } else if (op == 'S') {
      vctr_add1(vctr, XSA);
      if (hdrs[i] >> 8)
        vctr_add1(vctr, XSS);
    } else if (op == 'D') {
      vctr_add1(vctr, XDA);
      if (hdrs[i] >> 8)
        vctr_add1(vctr, XDS);
    }
  }
}

  static void
ibkv_req_helper(struct ib_socket * const s, const u64 wr_id, struct ibv_mr * const mr,
    struct ib_remote_ptr * const rptr, const u32 off, const u32 len, const u32 reqsz)
{
  struct ibv_sge sg;
  struct ibv_send_wr swr;
  ib_sge_fill_off(&sg, mr, off, len);
  const u32 imm = ((off) >> 6) | (reqsz >> 6 << 16);
  ib_wr_fill_write_imm(&swr, wr_id, &sg, rptr, off, imm);
  if (!ib_socket_post_send(s, &swr, true))
    debug_die();
  // the server will send a imm back
  if (!ib_socket_post_recv_sge(s, wr_id, NULL))
    debug_die();
}

#define BUFSIZE ((1lu << 21))
  static void *
ibkv_worker(void * const ptr)
{
  struct forker_worker_info * const mi = (typeof(mi))ptr;
  struct worker_info * const wi = (typeof(wi))mi->passdata[0];
  const u64 depth = wi->depth;
  const u64 batch = wi->batch;

  const u64 pset = a2u64(mi->argv[0]);
  const u64 pdel = a2u64(mi->argv[1]);
  const u64 pget = a2u64(mi->argv[2]);
  const u64 cset = 100lu - pdel;
  const u64 cget = cset - pset;
  const u64 cpro = cget - pget;
  const u64 vlen = a2u64(mi->argv[3]);
  struct thread_info ti = {.cpro = cpro, .cget = cget, .cset = cset, .vlen = vlen};
  mi->priv = (void *)&ti;

  struct ib_socket * const s = ib_socket_create(wi->ibport, depth);
  if (!s)
    return NULL;
  const int fd = client_connect(wi->host, wi->port);
  if (fd < 0) {
    ib_socket_destroy(s);
    return NULL;
  }
  const bool rc = ib_socket_connect_fd(s, fd);
  if (!rc) {
    close(fd);
    ib_socket_destroy(s);
    return NULL;
  }

  // local buffer
  u64 memsz = 0;
  u8 * const mem = pages_alloc_best(BUFSIZE, true, &memsz);
  struct ibv_mr * const mr = ib_mr_helper(wi->ibport, mem, memsz);

  // remote memory
  struct ib_remote_ptr rptr_req;
  struct ib_remote_ptr rptr_res;
  if (!ib_mr_recv_rptr_fd(fd, &rptr_req) || !ib_mr_recv_rptr_fd(fd, &rptr_res)) {
    close(fd);
    ib_socket_destroy(s);
    return NULL;
  }
  debug_assert(rptr_req.len && (rptr_req.len == rptr_res.len));

  bool do_more = true;
  u64 nsend = 0;
  u64 ndone = 0;
  const u32 reqsz = bits_round_down(rptr_req.len / depth, 6) - 64; // aligned at 64B
  struct ibv_sge sg;
  struct ibv_send_wr swr;
  // fill the queue
  for (u64 i = 0; i < depth; i++) {
    // send request
    const u32 off_req = i*reqsz;
    const u32 len_req = fill_requests(mem+off_req, reqsz, batch, mi);
    ibkv_req_helper(s, i, mr, &rptr_req, off_req, len_req, reqsz);
    nsend++;
    if ((mi->end_type == FORKER_END_COUNT) && (nsend >= mi->end_magic)) {
      do_more = false;
      break;
    }
  }

  // the loop
  do {
    int n = ib_socket_poll(s);
    if (n < 0) {
      perror("ib_socket_poll");
      break;
    } else if (n == 0) {
      cpu_pause();
    }

    for (int i = 0; i < n; i++) {
      struct ibv_wc * const wc = &(s->wcs[i]);
      if (wc->status != IBV_WC_SUCCESS) {
        do_more = false;
        break;
      }
      switch (wc->opcode) {
      // receive imm: can use write-imm or send-imm
      case IBV_WC_RECV_RDMA_WITH_IMM:
        {
          // imm: where to read results
          const u32 imm = wc->imm_data;
          const u32 res_off = (imm & 0xffffu) << 6;
          const u32 res_len = (imm >> 16) << 6;
          ib_sge_fill_off(&sg, mr, res_off, res_len);
          ib_wr_fill_read(&swr, wc->wr_id, &sg, &rptr_res, res_off);
          if (!ib_socket_post_send(s, &swr, true))
            debug_die();
        }
        break;
      case IBV_WC_RDMA_READ:
        // check results
        process_results(mem + (wc->wr_id * reqsz), mi->vctr);
        ndone++;
        if (do_more) { // next req
          const u32 off_req = wc->wr_id*reqsz;
          const u32 len_req = fill_requests(mem+off_req, reqsz, batch, mi);
          ibkv_req_helper(s, wc->wr_id, mr, &rptr_req, off_req, len_req, reqsz);
          nsend++;
          if ((mi->end_type == FORKER_END_COUNT) && (nsend >= mi->end_magic))
            do_more = false;
        }
        break;
      default:
        break;
      }
    }
    if ((mi->end_type == FORKER_END_TIME) && (time_nsec() >= mi->end_magic))
      do_more = false;
  } while (do_more || (ndone < nsend));
  ib_sge_fill_off(&sg, mr, 0, 0);
  ib_wr_fill_send_imm(&swr, depth, NULL, UINT32_MAX);
  if (!ib_socket_post_send(s, &swr, true))
    debug_die();

  bool wait_last = true;
  do {
    int n = ib_socket_poll(s);
    if (n < 0) {
      perror("ib_socket_poll");
      break;
    }
    for (int i = 0; i < n; i++) {
      struct ibv_wc * const wc = &(s->wcs[i]);
      if (wc->status != IBV_WC_SUCCESS)
        continue;
      if (wc->opcode == IBV_WC_SEND && wc->wr_id == depth) {
        wait_last = false;
        break;
      }
    }
  } while (wait_last);

  ib_socket_destroy(s);
  ibv_dereg_mr(mr);
  pages_unmap(mem, memsz);
  return NULL;
}

  static void
ibkv_help_message(void)
{
  fprintf(stderr, "Usage: <keysfile> api <ib-port> <host> <port> <depth> <batch> {rgen ... {pass ... }}\n");
  forker_passes_message();
  fprintf(stderr, "%s wargs[%d]: <pset> <pdel> <pget> <vlen>\n", __func__,  NR_WARGS);
}

  int
test_ibkv(const int argc, char ** const argv)
{
  if (argc < 6 || (strcmp(argv[0], "api") != 0))
    return -1;

  char *pref[64];
  for (int i = 0; i < 6; i++)
    pref[i] = argv[i];
  pref[6] = NULL;

  struct worker_info wi;
  wi.ibport = ib_port_create(a2u64(argv[1]));
  wi.host = argv[2];
  wi.port = argv[3];
  wi.depth = a2u64(argv[4]);
  wi.batch = a2u64(argv[5]);

  struct pass_info pi = {};
  pi.passdata[0] = &wi;
  pi.vctr_size = 9;
  pi.wf = ibkv_worker;
  pi.af = ibkv_analyze;
  const int n2 = forker_passes(argc - 6, argv + 6, pref, &pi, NR_WARGS);
  ib_port_destroy(wi.ibport);
  if (n2 < 0)
    return n2;

  return 6 + n2;
}

  int
main(int argc, char ** argv)
{
  if (argc < 2) {
    ibkv_help_message();
    exit(0);
  }
  if (kv_load(argv[1]) == false)
    return 0;

  const bool r = forker_main(argc - 2, argv + 2, test_ibkv);
  if (r == false)
    ibkv_help_message();
  kvloader_destroy();
  return 0;
}