promise allocator example

promise_allocator.cpp

#include <chrono>
#include <future>
#include <iostream>
#include <memory>
#include <thread>
#include <vector>


static size_t global_allocs = 0;
static size_t global_runtime_allocs = 0;

static size_t global_runtime_deallocs = 0;
static size_t global_deallocs = 0;

static size_t num_allocs = 0;
static size_t num_deallocs = 0;

static bool test_init = false;


// Override global new
void * operator new(std::size_t size)
{
  if (test_init) {
      std::cerr << "Allocation size: " << size << std::endl;
      global_runtime_allocs++;
  } else {
    global_allocs++;
  }
  return std::malloc(size);
}

void operator delete(void * ptr) noexcept
{
  if (ptr != nullptr) {
    if (test_init) {
      global_runtime_deallocs++;
    } else {
      global_deallocs++;
    }
    std::free(ptr);
    ptr = nullptr;
  }
}

void operator delete(void * ptr, size_t) noexcept
{
  if (ptr != nullptr) {
    if (test_init) {
      global_runtime_deallocs++;
    } else {
      global_deallocs++;
    }
    std::free(ptr);
    ptr = nullptr;
  }
}

// Necessary for using custom allocator with std::basic_string in GCC 4.8
// See: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=56437
template<typename T>
struct allocator_pointer_traits
{
  using size_type = std::size_t;
  using pointer = T *;
  using const_pointer = const T *;
  using difference_type = typename std::pointer_traits<pointer>::difference_type;
  using reference = T &;
  using const_reference = const T &;
};

template<>
struct allocator_pointer_traits<void>
{
};

// Simple instrumented allocator
template<typename T>
struct InstrumentedAllocator : public allocator_pointer_traits<T>
{
public:
  using value_type = T;

  InstrumentedAllocator() noexcept
  {
  }

  ~InstrumentedAllocator() noexcept
  {
  }

  template<typename U>
  InstrumentedAllocator(const InstrumentedAllocator<U> &) noexcept
  {
  }

  // During execution, publish/subscribe should only malloc from here
  T * allocate(size_t size, const void * = 0)
  {
    if (size == 0) {
      return nullptr;
    }
    num_allocs++;
    return static_cast<T *>(std::malloc(size * sizeof(T)));
  }

  void deallocate(T * ptr, size_t size)
  {
    (void)size;
    if (!ptr) {
      return;
    }
    num_deallocs++;
    std::free(ptr);
    ptr = nullptr;
  }

  // construct and destroy should not be required to implement allocator_traits,
  // this is a bug in gcc 4.8 for usage in std::pair, std::map
  template<typename U, typename ... Args,
  typename std::enable_if<!std::is_const<U>::value>::type * = nullptr>
  void
  construct(U * ptr, Args && ... args)
  {
    ::new(ptr)U(std::forward<Args>(args) ...);
  }

  template<typename U>
  void
  destroy(U * ptr)
  {
    ptr->~U();
  }

  // rebind should not be required to implement allocator_traits, this is a bug in gcc 4.8
  template<typename U>
  struct rebind
  {
    typedef InstrumentedAllocator<U> other;
  };
};

template<typename T, typename U>
constexpr bool operator==(const InstrumentedAllocator<T> &,
  const InstrumentedAllocator<U> &) noexcept
{
  return true;
}

template<typename T, typename U>
constexpr bool operator!=(const InstrumentedAllocator<T> &,
  const InstrumentedAllocator<U> &) noexcept
{
  return false;
}


struct Foo {
  std::vector<int, InstrumentedAllocator<int>> bar;
};

int main(int argc, char ** argv) {
  (void) argc;
  (void) argv;
  InstrumentedAllocator<void> alloc;

  std::promise<Foo> promise_(std::allocator_arg, alloc);
  std::shared_future<Foo> future_ = promise_.get_future().share();

  // Start a thread that blocks for a few ms and sets the future value
  std::thread result_thread(
    [&promise_]() {
      Foo result;
      result.bar.push_back(1);
      result.bar.push_back(2);
      result.bar.push_back(3);
      test_init = true;
      std::this_thread::sleep_for(std::chrono::milliseconds(5));
      promise_.set_value(result);
      test_init = false;
    }
  );
  future_.wait();
  result_thread.join();
  std::cout << "Runtime global allocations: " << global_runtime_allocs << " (expected: 0)" << std::endl;
  std::cout << "Runtime global deallocations: " << global_runtime_deallocs << " (expected: 0)" << std::endl;
}