folly/test/IndexedMemPoolTest.cpp

   1 /*
   2  * Copyright 2016 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/IndexedMemPool.h>
  18 #include <folly/test/DeterministicSchedule.h>
  19 #include <folly/portability/Unistd.h>
  20 #include <string>
  21 #include <thread>
  22 #include <semaphore.h>
  23 #include <gtest/gtest.h>
  24
  25 using namespace folly;
  26 using namespace folly::test;
  27
  28 TEST(IndexedMemPool, unique_ptr) {
  29   typedef IndexedMemPool<size_t> Pool;
  30   Pool pool(100);
  31
  32   for (size_t i = 0; i < 100000; ++i) {
  33     auto ptr = pool.allocElem();
  34     EXPECT_TRUE(!!ptr);
  35     *ptr = i;
  36   }
  37
  38   std::vector<Pool::UniquePtr> leak;
  39   while (true) {
  40     auto ptr = pool.allocElem();
  41     if (!ptr) {
  42       // good, we finally ran out
  43       break;
  44     }
  45     leak.emplace_back(std::move(ptr));
  46     EXPECT_LT(leak.size(), 10000);
  47   }
  48 }
  49
  50 TEST(IndexedMemPool, no_starvation) {
  51   const int count = 1000;
  52   const int poolSize = 100;
  53
  54   typedef DeterministicSchedule Sched;
  55   Sched sched(Sched::uniform(0));
  56
  57   typedef IndexedMemPool<int,8,8,DeterministicAtomic> Pool;
  58   Pool pool(poolSize);
  59
  60   for (auto pass = 0; pass < 10; ++pass) {
  61     int fd[2];
  62     EXPECT_EQ(pipe(fd), 0);
  63
  64     // makes sure we wait for available nodes, rather than fail allocIndex
  65     sem_t allocSem;
  66     sem_init(&allocSem, 0, poolSize);
  67
  68     // this semaphore is only needed for deterministic replay, so that we
  69     // always block in an Sched:: operation rather than in a read() syscall
  70     sem_t readSem;
  71     sem_init(&readSem, 0, 0);
  72
  73     std::thread produce = Sched::thread([&]() {
  74       for (auto i = 0; i < count; ++i) {
  75         Sched::wait(&allocSem);
  76         uint32_t idx = pool.allocIndex();
  77         EXPECT_NE(idx, 0);
  78         EXPECT_LE(idx,
  79             poolSize + (pool.NumLocalLists - 1) * pool.LocalListLimit);
  80         pool[idx] = i;
  81         EXPECT_EQ(write(fd[1], &idx, sizeof(idx)), sizeof(idx));
  82         Sched::post(&readSem);
  83       }
  84     });
  85
  86     std::thread consume = Sched::thread([&]() {
  87       for (auto i = 0; i < count; ++i) {
  88         uint32_t idx;
  89         Sched::wait(&readSem);
  90         EXPECT_EQ(read(fd[0], &idx, sizeof(idx)), sizeof(idx));
  91         EXPECT_NE(idx, 0);
  92         EXPECT_GE(idx, 1);
  93         EXPECT_LE(idx,
  94             poolSize + (Pool::NumLocalLists - 1) * Pool::LocalListLimit);
  95         EXPECT_EQ(pool[idx], i);
  96         pool.recycleIndex(idx);
  97         Sched::post(&allocSem);
  98       }
  99     });
 100
 101     Sched::join(produce);
 102     Sched::join(consume);
 103     close(fd[0]);
 104     close(fd[1]);
 105   }
 106 }
 107
 108 TEST(IndexedMemPool, st_capacity) {
 109   // only one local list => capacity is exact
 110   typedef IndexedMemPool<int,1,32> Pool;
 111   Pool pool(10);
 112
 113   EXPECT_EQ(pool.capacity(), 10);
 114   EXPECT_EQ(Pool::maxIndexForCapacity(10), 10);
 115   for (auto i = 0; i < 10; ++i) {
 116     EXPECT_NE(pool.allocIndex(), 0);
 117   }
 118   EXPECT_EQ(pool.allocIndex(), 0);
 119 }
 120
 121 TEST(IndexedMemPool, mt_capacity) {
 122   typedef IndexedMemPool<int,16,32> Pool;
 123   Pool pool(1000);
 124
 125   std::thread threads[10];
 126   for (auto i = 0; i < 10; ++i) {
 127     threads[i] = std::thread([&]() {
 128       for (auto j = 0; j < 100; ++j) {
 129         uint32_t idx = pool.allocIndex();
 130         EXPECT_NE(idx, 0);
 131       }
 132     });
 133   }
 134
 135   for (auto i = 0; i < 10; ++i) {
 136     threads[i].join();
 137   }
 138
 139   for (auto i = 0; i < 16 * 32; ++i) {
 140     pool.allocIndex();
 141   }
 142   EXPECT_EQ(pool.allocIndex(), 0);
 143 }
 144
 145 TEST(IndexedMemPool, locate_elem) {
 146   IndexedMemPool<int> pool(1000);
 147
 148   for (auto i = 0; i < 1000; ++i) {
 149     auto idx = pool.allocIndex();
 150     EXPECT_FALSE(idx == 0);
 151     int* elem = &pool[idx];
 152     EXPECT_TRUE(idx == pool.locateElem(elem));
 153   }
 154
 155   EXPECT_EQ(pool.locateElem(nullptr), 0);
 156 }
 157
 158 struct NonTrivialStruct {
 159   static FOLLY_TLS int count;
 160
 161   int elem_;
 162
 163   NonTrivialStruct() {
 164     elem_ = 0;
 165     ++count;
 166   }
 167
 168   NonTrivialStruct(std::unique_ptr<std::string>&& arg1, int arg2) {
 169     elem_ = arg1->length() + arg2;
 170     ++count;
 171   }
 172
 173   ~NonTrivialStruct() {
 174     --count;
 175   }
 176 };
 177
 178 FOLLY_TLS int NonTrivialStruct::count;
 179
 180 TEST(IndexedMemPool, eager_recycle) {
 181   typedef IndexedMemPool<NonTrivialStruct> Pool;
 182   Pool pool(100);
 183
 184   EXPECT_EQ(NonTrivialStruct::count, 0);
 185
 186   for (size_t i = 0; i < 10; ++i) {
 187     {
 188       std::unique_ptr<std::string> arg{ new std::string{ "abc" } };
 189       auto ptr = pool.allocElem(std::move(arg), 100);
 190       EXPECT_EQ(NonTrivialStruct::count, 1);
 191       EXPECT_EQ(ptr->elem_, 103);
 192       EXPECT_TRUE(!!ptr);
 193     }
 194     EXPECT_EQ(NonTrivialStruct::count, 0);
 195   }
 196 }
 197
 198 TEST(IndexedMemPool, late_recycle) {
 199   {
 200     typedef IndexedMemPool<NonTrivialStruct, 8, 8, std::atomic, false, false>
 201         Pool;
 202     Pool pool(100);
 203
 204     EXPECT_EQ(NonTrivialStruct::count, 0);
 205
 206     for (size_t i = 0; i < 10; ++i) {
 207       {
 208         auto ptr = pool.allocElem();
 209         EXPECT_TRUE(NonTrivialStruct::count > 0);
 210         EXPECT_TRUE(!!ptr);
 211         ptr->elem_ = i;
 212       }
 213       EXPECT_TRUE(NonTrivialStruct::count > 0);
 214     }
 215   }
 216   EXPECT_EQ(NonTrivialStruct::count, 0);
 217 }