1 files changed, 47 insertions, 0 deletions
diff --git a/src/rocksdb/memory/concurrent_arena.cc b/src/rocksdb/memory/concurrent_arena.cc
new file mode 100644
index 000000000..3333f94db
--- /dev/null
+++ b/src/rocksdb/memory/concurrent_arena.cc
@@ -0,0 +1,47 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+
+#include "memory/concurrent_arena.h"
+#include <thread>
+#include "port/port.h"
+#include "util/random.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+#ifdef ROCKSDB_SUPPORT_THREAD_LOCAL
+__thread size_t ConcurrentArena::tls_cpuid = 0;
+#endif
+
+namespace {
+// If the shard block size is too large, in the worst case, every core
+// allocates a block without populate it. If the shared block size is
+// 1MB, 64 cores will quickly allocate 64MB, and may quickly trigger a
+// flush. Cap the size instead.
+const size_t kMaxShardBlockSize = size_t{128 * 1024};
+}  // namespace
+
+ConcurrentArena::ConcurrentArena(size_t block_size, AllocTracker* tracker,
+                                 size_t huge_page_size)
+    : shard_block_size_(std::min(kMaxShardBlockSize, block_size / 8)),
+      shards_(),
+      arena_(block_size, tracker, huge_page_size) {
+  Fixup();
+}
+
+ConcurrentArena::Shard* ConcurrentArena::Repick() {
+  auto shard_and_index = shards_.AccessElementAndIndex();
+#ifdef ROCKSDB_SUPPORT_THREAD_LOCAL
+  // even if we are cpu 0, use a non-zero tls_cpuid so we can tell we
+  // have repicked
+  tls_cpuid = shard_and_index.second | shards_.Size();
+#endif
+  return shard_and_index.first;
+}
+
+}  // namespace ROCKSDB_NAMESPACE