YugabyteDB (2.13.1.0-b60, 21121d69985fbf76aa6958d8f04a9bfa936293b5)

//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.

//  This source code is licensed under the BSD-style license found in the

//  LICENSE file in the root directory of this source tree. An additional grant

//  of patent rights can be found in the PATENTS file in the same directory.

//

// The following only applies to changes made to this file as part of YugaByte development.

//

// Portions Copyright (c) YugaByte, 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.

//

// 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.

// Arena is an implementation of Allocator class. For a request of small size,

// it allocates a block with pre-defined block size. For a request of big

// size, it uses malloc to directly get the requested size.

#ifndef YB_ROCKSDB_UTIL_ARENA_H

#define YB_ROCKSDB_UTIL_ARENA_H

#pragma once

#ifndef OS_WIN

#include <sys/mman.h>

#endif

#include <assert.h>

#include <stdint.h>

#include <cstddef>

#include <cerrno>

#include <vector>

#include "yb/rocksdb/util/allocator.h"

#include "yb/rocksdb/util/mutexlock.h"

namespace yb {

class MemTracker;

}

namespace rocksdb {

class Arena : public Allocator {

 public:

  // No copying allowed

  Arena(const Arena&) = delete;

  void operator=(const Arena&) = delete;

  static const size_t kInlineSize = 2048;

  static const size_t kMinBlockSize;

  static const size_t kMaxBlockSize;

  // huge_page_size: if 0, don't use huge page TLB. If > 0 (should set to the

  // supported hugepage size of the system), block allocation will try huge

  // page TLB first. If allocation fails, will fall back to normal case.

  explicit Arena(size_t block_size = kMinBlockSize, size_t huge_page_size = 0);

  ~Arena();

  char* Allocate(size_t bytes) override;

  void SetMemTracker(std::shared_ptr<yb::MemTracker> mem_tracker);

  // huge_page_size: if >0, will try to allocate from huage page TLB.

  // The argument will be the size of the page size for huge page TLB. Bytes

  // will be rounded up to multiple of the page size to allocate through mmap

  // anonymous option with huge page on. The extra  space allocated will be

  // wasted. If allocation fails, will fall back to normal case. To enable it,

  // need to reserve huge pages for it to be allocated, like:

  //     sysctl -w vm.nr_hugepages=20

  // See linux doc Documentation/vm/hugetlbpage.txt for details.

  // huge page allocation can fail. In this case it will fail back to

  // normal cases. The messages will be logged to logger. So when calling with

  // huge_page_tlb_size > 0, we highly recommend a logger is passed in.

  // Otherwise, the error message will be printed out to stderr directly.

  char* AllocateAligned(size_t bytes, size_t huge_page_size = 0,

                        Logger* logger = nullptr) override;

  // Returns an estimate of the total memory usage of data allocated

  // by the arena (exclude the space allocated but not yet used for future

  // allocations).

  size_t ApproximateMemoryUsage() const {

    return blocks_memory_ + blocks_.capacity() * sizeof(char*) -

           alloc_bytes_remaining_;

  }

  size_t MemoryAllocatedBytes() const { return blocks_memory_; }

  size_t AllocatedAndUnused() const { return alloc_bytes_remaining_; }

  // If an allocation is too big, we'll allocate an irregular block with the

  // same size of that allocation.

  size_t IrregularBlockNum() const { return irregular_block_num; }

  size_t BlockSize() const override { return kBlockSize; }

 private:

  void Consumed(size_t size);

  char inline_block_[kInlineSize] __attribute__((__aligned__(sizeof(void*))));

  // Number of bytes allocated in one block

  const size_t kBlockSize;

  // Array of new[] allocated memory blocks

  typedef std::vector<char*> Blocks;

  Blocks blocks_;

  struct MmapInfo {

    void* addr_;

    size_t length_;

    MmapInfo(void* addr, size_t length) : addr_(addr), length_(length) {}

  };

  std::vector<MmapInfo> huge_blocks_;

  size_t irregular_block_num = 0;

  // Stats for current active block.

  // For each block, we allocate aligned memory chucks from one end and

  // allocate unaligned memory chucks from the other end. Otherwise the

  // memory waste for alignment will be higher if we allocate both types of

  // memory from one direction.

  char* unaligned_alloc_ptr_ = nullptr;

  char* aligned_alloc_ptr_ = nullptr;

  // How many bytes left in currently active block?

  size_t alloc_bytes_remaining_ = 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);

  // Bytes of memory in blocks allocated so far

  size_t blocks_memory_ = 0;

  std::shared_ptr<yb::MemTracker> mem_tracker_;

};

inline char* Arena::Allocate(size_t bytes) {

  // The semantics of what to return are a bit messy if we allow

  // 0-byte allocations, so we disallow them here (we don't need

  // them for our internal use).

  assert(bytes > 0);

  if (bytes <= alloc_bytes_remaining_) {

    unaligned_alloc_ptr_ -= bytes;

    alloc_bytes_remaining_ -= bytes;

    return unaligned_alloc_ptr_;

  }

  return AllocateFallback(bytes, false /* unaligned */);

}

// check and adjust the block_size so that the return value is

//  1. in the range of [kMinBlockSize, kMaxBlockSize].

//  2. the multiple of align unit.

extern size_t OptimizeBlockSize(size_t block_size);

}  // namespace rocksdb

#endif // YB_ROCKSDB_UTIL_ARENA_H