folly/test/AHMIntStressTest.cpp

   1 /*
   2  * Copyright 2015 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 <gtest/gtest.h>
  18
  19 #include <thread>
  20 #include <memory>
  21 #include <mutex>
  22
  23 #include <folly/AtomicHashMap.h>
  24 #include <folly/ScopeGuard.h>
  25 #include <folly/Memory.h>
  26
  27 namespace {
  28
  29 struct MyObject {
  30   explicit MyObject(int i) : i(i) {}
  31   int i;
  32 };
  33
  34 typedef folly::AtomicHashMap<int,std::shared_ptr<MyObject>> MyMap;
  35 typedef std::lock_guard<std::mutex> Guard;
  36
  37 std::unique_ptr<MyMap> newMap() { return folly::make_unique<MyMap>(100); }
  38
  39 struct MyObjectDirectory {
  40   MyObjectDirectory()
  41     : cur_(newMap())
  42     , prev_(newMap())
  43   {}
  44
  45   std::shared_ptr<MyObject> get(int key) {
  46     auto val = tryGet(key);
  47     if (val) {
  48       return val;
  49     }
  50
  51     std::shared_ptr<MyMap> cur;
  52     {
  53       Guard g(lock_);
  54       cur = cur_;
  55     }
  56
  57     auto ret = cur->insert(key, std::make_shared<MyObject>(key));
  58     return ret.first->second;
  59   }
  60
  61   std::shared_ptr<MyObject> tryGet(int key) {
  62     std::shared_ptr<MyMap> cur;
  63     std::shared_ptr<MyMap> prev;
  64     {
  65       Guard g(lock_);
  66       cur = cur_;
  67       prev = prev_;
  68     }
  69
  70     auto it = cur->find(key);
  71     if (it != cur->end()) {
  72       return it->second;
  73     }
  74
  75     it = prev->find(key);
  76     if (it != prev->end()) {
  77       auto ret = cur->insert(key, it->second);
  78       return ret.first->second;
  79     }
  80
  81     return nullptr;
  82   }
  83
  84   void archive() {
  85     std::shared_ptr<MyMap> cur(newMap());
  86
  87     Guard g(lock_);
  88     prev_ = cur_;
  89     cur_ = cur;
  90   }
  91
  92   std::mutex lock_;
  93   std::shared_ptr<MyMap> cur_;
  94   std::shared_ptr<MyMap> prev_;
  95 };
  96
  97 }
  98
  99 //////////////////////////////////////////////////////////////////////
 100
 101 /*
 102  * This test case stresses ThreadLocal allocation/deallocation heavily
 103  * via ThreadCachedInt and AtomicHashMap, and a bunch of other
 104  * mallocing.
 105  */
 106 TEST(AHMIntStressTest, Test) {
 107   auto const objs = new MyObjectDirectory();
 108   SCOPE_EXIT { delete objs; };
 109
 110   std::vector<std::thread> threads;
 111   for (int threadId = 0; threadId < 64; ++threadId) {
 112     threads.emplace_back(
 113       [objs,threadId] {
 114         for (int recycles = 0; recycles < 500; ++recycles) {
 115           for (int i = 0; i < 10; i++) {
 116             auto val = objs->get(i);
 117           }
 118
 119           objs->archive();
 120         }
 121       }
 122     );
 123   }
 124
 125   for (auto& t : threads) t.join();
 126 }