ThreadCachedArenaTest.cpp

/*
 * Copyright 2013 Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "folly/ThreadCachedArena.h"
#include "folly/Memory.h"

#include <mutex>
#include <thread>
#include <iterator>
#include <algorithm>
#include <random>
#include <unordered_map>

#include <glog/logging.h>
#include <gtest/gtest.h>

#include "folly/Range.h"
#include "folly/Benchmark.h"

using namespace folly;

namespace {

class ArenaTester {
 public:
  explicit ArenaTester(ThreadCachedArena& arena) : arena_(&arena) { }

  void allocate(size_t count, size_t maxSize);
  void verify();
  void merge(ArenaTester&& other);

 private:
  std::mutex mergeMutex_;
  std::vector<std::pair<uint8_t, Range<uint8_t*>>> areas_;
  ThreadCachedArena* arena_;
};

void ArenaTester::allocate(size_t count, size_t maxSize) {
  // Allocate chunks of memory of random sizes
  std::mt19937 rnd;
  std::uniform_int_distribution<uint32_t> sizeDist(1, maxSize - 1);
  areas_.clear();
  areas_.reserve(count);
  for (size_t i = 0; i < count; i++) {
    size_t size = sizeDist(rnd);
    uint8_t* p = static_cast<uint8_t*>(arena_->allocate(size));
    areas_.emplace_back(rnd() & 0xff, Range<uint8_t*>(p, size));
  }

  // Fill each area with a different value, to prove that they don't overlap
  // Fill in random order.
  std::random_shuffle(
      areas_.begin(), areas_.end(),
      [&rnd] (int n) -> int {
        return std::uniform_int_distribution<uint32_t>(0, n-1)(rnd);
      });

  for (auto& p : areas_) {
    std::fill(p.second.begin(), p.second.end(), p.first);
  }
}

void ArenaTester::verify() {
  for (auto& p : areas_) {
    for (auto v : p.second) {
      EXPECT_EQ(p.first, v);
    }
  }
}

void ArenaTester::merge(ArenaTester&& other) {
  {
    std::lock_guard<std::mutex> lock(mergeMutex_);
    std::move(other.areas_.begin(), other.areas_.end(),
              std::back_inserter(areas_));
  }
  other.areas_.clear();
}

}  // namespace

TEST(ThreadCachedArena, BlockSize) {
  struct Align { char c; } __attribute__((aligned));
  static const size_t alignment = alignof(Align);
  static const size_t requestedBlockSize = 64;

  ThreadCachedArena arena(requestedBlockSize);
  size_t blockSize = alignment;
  uint8_t* prev = static_cast<uint8_t*>(arena.allocate(1));

  // Keep allocating until we're no longer one single alignment away from the
  // previous allocation -- that's when we've gotten to the next block.
  uint8_t* p;
  while ((p = static_cast<uint8_t*>(arena.allocate(1))) ==
         prev + alignment) {
    prev = p;
    blockSize += alignment;
  }

  VLOG(1) << "Requested block size: " << requestedBlockSize << ", actual: "
          << blockSize;
  EXPECT_LE(requestedBlockSize, blockSize);
}

TEST(ThreadCachedArena, SingleThreaded) {
  static const size_t requestedBlockSize = 64;
  ThreadCachedArena arena(requestedBlockSize);
  ArenaTester tester(arena);
  tester.allocate(100, 100 << 10);
  tester.verify();
}

TEST(ThreadCachedArena, MultiThreaded) {
  static const size_t requestedBlockSize = 64;
  ThreadCachedArena arena(requestedBlockSize);
  ArenaTester mainTester(arena);

  // Do this twice, to catch the possibility that memory from the first
  // round gets freed
  static const size_t numThreads = 20;
  for (size_t i = 0; i < 2; i++) {
    std::vector<std::thread> threads;
    threads.reserve(numThreads);
    for (size_t j = 0; j < numThreads; j++) {
      threads.emplace_back(
          [&arena, &mainTester] () {
            ArenaTester tester(arena);
            tester.allocate(500, 1 << 10);
            tester.verify();
            mainTester.merge(std::move(tester));
          });
    }
    for (auto& t : threads) {
      t.join();
    }
  }

  mainTester.verify();
}

TEST(ThreadCachedArena, StlAllocator) {
  typedef std::unordered_map<
    int, int, std::hash<int>, std::equal_to<int>,
    StlAllocator<ThreadCachedArena, std::pair<const int, int>>> Map;

  static const size_t requestedBlockSize = 64;
  ThreadCachedArena arena(requestedBlockSize);

  Map map {0, std::hash<int>(), std::equal_to<int>(),
           StlAllocator<ThreadCachedArena, std::pair<const int, int>>(&arena)};

  for (int i = 0; i < 1000; i++) {
    map[i] = i;
  }

  for (int i = 0; i < 1000; i++) {
    EXPECT_EQ(i, map[i]);
  }
}

namespace {

static const int kNumValues = 10000;

BENCHMARK(bmUMStandard, iters) {
  typedef std::unordered_map<int, int> Map;

  while (iters--) {
    Map map {0};
    for (int i = 0; i < kNumValues; i++) {
      map[i] = i;
    }
  }
}

BENCHMARK(bmUMArena, iters) {
  typedef std::unordered_map<
    int, int, std::hash<int>, std::equal_to<int>,
    StlAllocator<ThreadCachedArena, std::pair<const int, int>>> Map;

  while (iters--) {
    ThreadCachedArena arena;

    Map map {0, std::hash<int>(), std::equal_to<int>(),
             StlAllocator<ThreadCachedArena, std::pair<const int, int>>(
                 &arena)};

    for (int i = 0; i < kNumValues; i++) {
      map[i] = i;
    }
  }
}

BENCHMARK_DRAW_LINE()

BENCHMARK(bmMStandard, iters) {
  typedef std::map<int, int> Map;

  while (iters--) {
    Map map;
    for (int i = 0; i < kNumValues; i++) {
      map[i] = i;
    }
  }
}

BENCHMARK_DRAW_LINE()

BENCHMARK(bmMArena, iters) {
  typedef std::map<
    int, int, std::less<int>,
    StlAllocator<ThreadCachedArena, std::pair<const int, int>>> Map;

  while (iters--) {
    ThreadCachedArena arena;

    Map map {std::less<int>(),
             StlAllocator<ThreadCachedArena, std::pair<const int, int>>(
                 &arena)};

    for (int i = 0; i < kNumValues; i++) {
      map[i] = i;
    }
  }
}

BENCHMARK_DRAW_LINE()

}  // namespace


// Benchmark                               Iters   Total t    t/iter iter/sec
// ----------------------------------------------------------------------------
// Comparing benchmarks: bmUMStandard,bmUMArena
//  + 143% bmUMStandard                     1570  2.005 s   1.277 ms  782.9
// *       bmUMArena                        3817  2.003 s   524.7 us  1.861 k
// ----------------------------------------------------------------------------
// Comparing benchmarks: bmMStandard,bmMArena
//  +79.0% bmMStandard                      1197  2.009 s   1.678 ms  595.8
// *       bmMArena                         2135  2.002 s   937.6 us  1.042 k
// ----------------------------------------------------------------------------

int main(int argc, char *argv[]) {
  testing::InitGoogleTest(&argc, argv);
  google::ParseCommandLineFlags(&argc, &argv, true);
  auto ret = RUN_ALL_TESTS();
  if (!ret && FLAGS_benchmark) {
    folly::runBenchmarks();
  }
  return ret;
}