latency_spinlock.cpp

#include <stdio.h>
#include <stdint.h>
#include <math.h>
#include <unistd.h>

#include <sys/time.h>

#include <algorithm>

#include <boost/atomic.hpp>
#include <uv.h>

#define NUM_THREADS 4
#define NUM_HOSTS 100

int64_t get_time_in_nanos() {
  struct timeval time;
  gettimeofday(&time, NULL);
  int64_t usecs = (int64_t)time.tv_sec * (1000 * 1000) + (int64_t)time.tv_usec;
  return 1000 * usecs;
}

class Spinlock {
public:
  typedef enum { LOCKED, UNLOCKED } LockState;

  Spinlock()
    : state_(UNLOCKED) {}

  void lock() {
    while (state_.exchange(LOCKED, boost::memory_order_acquire) == LOCKED) { }
  }

  void unlock() {
    state_.store(UNLOCKED, boost::memory_order_release);
  }

private:
  boost::atomic<LockState> state_;

private:
  char cache_padding__[64];
};

struct TimestampedAverage {
  TimestampedAverage()
    : timestamp(-1)
    , average(-1)
    , count(0) {}

  int64_t timestamp;
  int64_t average;
  int64_t count;
};

class HostLatencyTracker {
public:
  HostLatencyTracker(int64_t scale = 50, int64_t threshold_to_account = 15)
    : scale_(scale)
    , threshold_to_account_(threshold_to_account) {}

  TimestampedAverage update(int64_t latency_ns) {
      TimestampedAverage next;

      int64_t now = get_time_in_nanos();

      Spinlock* spinlock = get_spinlock(this);

      spinlock->lock();
      TimestampedAverage previous = current_;

      if (previous.count < threshold_to_account_) {
        next.count = previous.count + 1;
        next.average = -1;
        next.timestamp = now;
      } else if (previous.average < 0) {
        next.count = previous.count + 1;
        next.average = latency_ns;
        next.timestamp = now;
      } else {
        int64_t delay = now - previous.timestamp;
        if (delay <= 0) {
          spinlock->unlock();
          return next;
        }

        double scaled_delay = static_cast<double>(delay) / scale_;
        double weight = log(scaled_delay + 1) / scaled_delay;

        next.count = previous.count + 1;
        next.average = static_cast<int64_t>((1.0 - weight) * latency_ns + weight * previous.average);
        next.timestamp = now;
      }

      current_ = next;
      spinlock->unlock();

      return next;
  }

  TimestampedAverage get() {
    Spinlock* spinlock = get_spinlock(this);
    spinlock->lock();
    TimestampedAverage copy = current_;
    spinlock->unlock();
    return copy;
  }

private:
  int64_t scale_;
  int64_t threshold_to_account_;
  TimestampedAverage current_;

private:
  Spinlock* get_spinlock(void* p) {
    return &spinlock_pool_[reinterpret_cast<size_t>(p) % 41];
  }
  static Spinlock spinlock_pool_[41];
};

Spinlock HostLatencyTracker::spinlock_pool_[41];

struct ThreadArg {
  ThreadArg(unsigned seed, HostLatencyTracker* hosts)
    : seed(seed)
    , hosts(hosts) {}

  unsigned seed;
  HostLatencyTracker* hosts;
};

void thread(void* arg) {
  unsigned seed = static_cast<ThreadArg*>(arg)->seed;
  HostLatencyTracker* hosts = static_cast<ThreadArg*>(arg)->hosts;
  delete static_cast<ThreadArg*>(arg);

  int r = rand_r(&seed);

  int count = 0;
  int64_t start = get_time_in_nanos(), elapsed = 0;
  while (true) {
    if (elapsed != 0) {
      HostLatencyTracker* host = &hosts[r % NUM_HOSTS];

      TimestampedAverage ts = host->get();
      host->update(elapsed);

      if (++count == 1000000) {
        printf("current average %lld for host %d\n", ts.average, r % NUM_HOSTS);
        count = 0;
      }

      r++;
    }

    elapsed = get_time_in_nanos() - start;
    start = get_time_in_nanos();
  }
}

void latency_thread(void* arg) {
  HostLatencyTracker* hosts = static_cast<ThreadArg*>(arg)->hosts;
  delete static_cast<ThreadArg*>(arg);

  while (true) {
    int64_t start = get_time_in_nanos();

    int64_t min_average = INT64_MAX;
    int64_t max_average = INT64_MIN;
    int64_t total = 0;
    int count = 0;
    for (int i = 0; i < NUM_HOSTS; ++i) {
      HostLatencyTracker* host = &hosts[i];

      TimestampedAverage ts = host->get();
      if (ts.count > 50 && ts.average > 0) {
        min_average = std::min(ts.average, min_average);
        max_average = std::max(ts.average, max_average);
        total += ts.average;
        count++;
      }
    }

    int64_t elapsed = get_time_in_nanos() - start;

    double average = 0.0;
    if (count > 0) {
      average = (double)total / count;
    }

    printf("found min %lld (max: %lld, avg: %f) in %lld\n", min_average, max_average, average, elapsed);
    usleep(100 * 1000);
  }
}

int main() {
  uv_thread_t thread_ids[NUM_THREADS];
  uv_thread_t latency_thread_id;

  HostLatencyTracker hosts[NUM_HOSTS];

  for (int i = 0; i < NUM_THREADS; ++i) {
    uv_thread_create(&thread_ids[i], thread, new ThreadArg(i, hosts));
  }

  uv_thread_create(&latency_thread_id, latency_thread, new ThreadArg(0, hosts));

  for (int i = 0; i < NUM_THREADS; ++i) {
    uv_thread_join(&thread_ids[i]);
  }

  uv_thread_join(&latency_thread_id);

  return 0;
}