深入淺出帶你走進 RocksDB

語言: CN / TW / HK

時間 2022-09-30 17:01:31 開務資料庫

主題: relaxed glibc tracker rocksdb

RocksDB 是基於 Google LevelDB 研發的高效能 Key-Value 持久化儲存引擎，以庫元件形式嵌入程式中，為大規模分散式應用在 SSD 上執行提供優化。RocksDB 重點是提供工具支援，具體實現將交給上層應用。

正是這種高度可定製化能力，使得 RocksDB 可涵蓋包括工作負載等眾多場景的應用，今天我們將逐一為大家介紹：

"為什麼需要記憶體管理器？為什麼不使用現有的記憶體管理器？ RocksDB 究竟是如何實現的？"

01 為什麼需要記憶體管理器？

RocksDB 有很多核心場景需要分配記憶體的，包括但不限於 Memtable、 Cache、Iterator 等，高效優質的記憶體分配器就顯得尤為重要。一個優秀的通用記憶體分配器需要具備以下特性:

儘量避免記憶體碎片；併發分配效能好；
額外的空間損耗盡量少；
兼具通用性、相容性、可移植性且易除錯。

記憶體管理可以分為 3 個層次，自下而上分別是：

作業系統核心的記憶體管理；
glibc 層使用系統呼叫維護的記憶體管理；
應用程式從 glibc 層動態分配記憶體後，根據應用程式本身的程式特性進行優化，比如使用引用計數、std::shared_ptr、apache 等記憶體池方式的記憶體管理。

目前大部分服務端程式使用 glibc 提供的 malloc/free 系列函式，而 glibc 使用的 ptmalloc2 在效能上遠遠落後於 Google 的 Tcmalloc 和 Facebook 的 Jemalloc 。後兩者只需使用 LD_PRELOAD 環境變數啟動程式即可，甚至並不需要重新編譯。

02 為什麼不使用現有記憶體管理器？

RocksDB 使用記憶體的場景集中在 Memtable、Cache、Iterator 等核心鏈路上，在保證高效能的同時，還需要實現對記憶體分配器內部控制（包括監控等）。

當前現有的記憶體分配器不具有主動上報記憶體監控細節，所以從長遠看，依舊需要自己實現 RocksDB 專有的記憶體管理器。記憶體管理器設計思路如下：

小記憶體的申請通過 Thread-cache 或者 Per-cpu cache，保證了 CPU 不會頻繁的 cache-miss；
若支援 MAP_HUGETLB，則會直接 mmap；
若大頁記憶體則直接從底層的分配器去申請。

03 RocksDB 是如何實現的？

如圖所示，RocksDB 的記憶體管理器是支援併發的，接下來讓我們一起從原始碼入手，看看具體如何實現的。

Allocator


// Abstract interface for allocating memory in blocks. This memory is freed
// When the allocator object is destroyed. See the Arena class for more info.
class Allocator {
 public:
  virtual ~Allocator() {}
  // 分配記憶體
  virtual char* Allocate(size_t bytes) = 0;                     
  // 對齊分配記憶體
  virtual char* AllocateAligned(size_t bytes, size_t huge_page_size = 0,
                                Logger* logger = nullptr) = 0;  
  // 塊大小
  virtual size_t BlockSize() const = 0;                         
};

Arena

Arena 負責實現 RocksDB 的記憶體分配，我們從中可以看到其針對不同大小的記憶體分配請求，採取不同的分配策略，從而減少記憶體碎片。‍

class Arena : public Allocator {
    static const size_t kInlineSize = 2048;
    static const size_t kMinBlockSize;  // 4K
    static const size_t kMaxBlockSize;  // 2G
    
private:
    char inline_block_[kInlineSize] __attribute__((__aligned__(alignof(max_align_t))));
    const size_t kBlockSize;       // 每個 Block 的大小
    using Blocks = std::vector<char*>;
    Blocks blocks_;                // 分配的新 Block 地址（不夠了就從新分配）
    
    struct MmapInfo {
        void* addr_;
        size_t length_;
        MmapInfo(void* addr, size_t length) : addr_(addr), length_(length) {}
    };
    std::vector<MmapInfo> huge_blocks_;     // 大塊使用 mmap
    size_t irregular_block_num = 0;         // 不整齊的塊分配次數（塊大於 kBlockSize/4）
    
    char* unaligned_alloc_ptr_ = nullptr;   // 未對齊的一端指標（高地址開始）
    char* aligned_alloc_ptr_ = nullptr;     // 對齊一端指標（低地址開始）
    size_t alloc_bytes_remaining_ = 0;      // 記憶體剩餘
    
    size_t blocks_memory_ = 0;              // 已經分配的記憶體大小
    
    #ifdef MAP_HUGETLB
    size_t hugetlb_size_ = 0;
    #endif  // MAP_HUGETLB
    char* AllocateFromHugePage(size_t bytes);
    char* AllocateFallback(size_t bytes, bool aligned);
    char* AllocateNewBlock(size_t block_bytes);
    
    AllocTracker* tracker_;
}

針對 Allocate 和 AllocateAligned，我們採用對同一塊 Block 的兩端進行分配。AllocateAligned 從記憶體塊的低地址開始分配，Allocate 從高地址開始分配。分配流程圖如下：

ConcurrentArena

記憶體分配器不僅需要減少記憶體碎片，同樣需要保證併發分配的效能，那麼 RocksDB 是如何實現 ConcurrentArena 的呢？

從記憶體管理架構圖可以看出，RocksDB 維護了 CoreLocal 記憶體陣列，每個執行緒從所在 CPU 對應的本地 Shard 上分配記憶體，若不足再去主記憶體 Arena 進行分配。我們從幾個核心類開始逐一介紹：

1、ConcurrentArena


class ConcurrentArena : public Allocator {
 public:
  // block_size and huge_page_size are the same as for Arena (and are
  // in fact just passed to the constructor of arena_.  The core-local
  // shards compute their shard_block_size as a fraction of block_size
  // that varies according to the hardware concurrency level.
  explicit ConcurrentArena(size_t block_size = Arena::kMinBlockSize,
                           AllocTracker* tracker = nullptr,
                           size_t huge_page_size = 0);
                           
  char* Allocate(size_t bytes) override {
    return AllocateImpl(bytes, false /*force_arena*/,
                        [this, bytes]() { return arena_.Allocate(bytes); });
  }

 private:
  ...
  CoreLocalArray<Shard> shards_;    // 維護了一個 CoreLocal 記憶體陣列
  Arena arena_;                     // 主記憶體
  ...
};

2、CoreLocalArray


// An array of core-local values. Ideally the value type, T, is cache aligned to
// prevent false sharing.
template <typename T>
class CoreLocalArray {
 public:
  ...
  // returns pointer to element for the specified core index. This can be used,
  // e.g., for aggregation, or if the client caches core index.
  T* AccessAtCore(size_t core_idx) const;

 private:
  std::unique_ptr<T[]> data_;
  int size_shift_;
};

3、併發分配流程


template <typename Func>
  char* AllocateImpl(size_t bytes, bool force_arena, const Func& func) {
    size_t cpu;

    // 1：大塊則直接從 Arena 上分配，需要加鎖
    std::unique_lock<SpinMutex> arena_lock(arena_mutex_, std::defer_lock);
    if (bytes > shard_block_size_ / 4 || force_arena ||
        ((cpu = tls_cpuid) == 0 &&
         !shards_.AccessAtCore(0)->allocated_and_unused_.load(
             std::memory_order_relaxed) &&
         arena_lock.try_lock())) {
      if (!arena_lock.owns_lock()) {
        arena_lock.lock();
      }
      auto rv = func();
      Fixup();
      return rv;
    }

    // 2：挑選 CPU 對應的 Shard
    Shard* s = shards_.AccessAtCore(cpu & (shards_.Size() - 1));
    if (!s->mutex.try_lock()) {
      s = Repick();
      s->mutex.lock();
    }
    std::unique_lock<SpinMutex> lock(s->mutex, std::adopt_lock);

    size_t avail = s->allocated_and_unused_.load(std::memory_order_relaxed);
    
    // 2.1：若當前 Shard 可用記憶體不夠時，去 Arena 分配存入 Shard
    if (avail < bytes) {
      // reload
      std::lock_guard<SpinMutex> reload_lock(arena_mutex_);

      // If the arena's current block is within a factor of 2 of the right
      // size, we adjust our request to avoid arena waste.
      auto exact = arena_allocated_and_unused_.load(std::memory_order_relaxed);
      assert(exact == arena_.AllocatedAndUnused());

      if (exact >= bytes && arena_.IsInInlineBlock()) {
        // If we haven't exhausted arena's inline block yet, allocate from arena
        // directly. This ensures that we'll do the first few small allocations
        // without allocating any blocks.
        // In particular this prevents empty memtables from using
        // disproportionately large amount of memory: a memtable allocates on
        // the order of 1 KB of memory when created; we wouldn't want to
        // allocate a full arena block (typically a few megabytes) for that,
        // especially if there are thousands of empty memtables.
        auto rv = func();
        Fixup();
        return rv;
      }

      avail = exact >= shard_block_size_ / 2 && exact < shard_block_size_ * 2
                  ? exact
                  : shard_block_size_;
      s->free_begin_ = arena_.AllocateAligned(avail);
      Fixup();
    }
    s->allocated_and_unused_.store(avail - bytes, std::memory_order_relaxed);

    // 3：根據是否對齊，判斷是從高地址/低地址分配
    char* rv;
    if ((bytes % sizeof(void*)) == 0) {
      // aligned allocation from the beginning
      rv = s->free_begin_;
      s->free_begin_ += bytes;
    } else {
      // unaligned from the end
      rv = s->free_begin_ + avail - bytes;
    }
    return rv;
  }

總結

入參 Func 就是上面傳入的 Lambda 表示式：this, bytes { returnarena_.Allocate(bytes) ；
當請求記憶體塊較大時，直接從 Arena 分配且需要加鎖；否則直接從當前 CPU 對應的 Shard 分配；
若當前 Shard 可用記憶體不夠，需要從 Arena 再次請求；4、根據是否對齊，判斷從高/低地址分配。

「其他文章」