mem_pool.h

/* Copyright (c) 2021 OceanBase and/or its affiliates. All rights reserved.
miniob is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
         http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details. */
 
//
// Created by Longda on 2010
//
 
#pragma once
 
#include <sstream>
 
#include "common/lang/mutex.h"
#include "common/lang/string.h"
#include "common/lang/set.h"
#include "common/lang/list.h"
#include "common/lang/memory.h"
#include "common/lang/sstream.h"
#include "common/log/log.h"
#include "common/os/os.h"
 
#include <sanitizer/asan_interface.h>
 
namespace common {
 
#define DEFAULT_ITEM_NUM_PER_POOL 128
#define DEFAULT_POOL_NUM 1
 
typedef bool (*match)(void *item, void *input_arg);
 
template <class T>
class MemPool
{
public:
  MemPool(const char *tag) : name(tag)
  {
    this->size = 0;
 
    pthread_mutexattr_t mutexatr;
    pthread_mutexattr_init(&mutexatr);
    pthread_mutexattr_settype(&mutexatr, PTHREAD_MUTEX_RECURSIVE);
 
    MUTEX_INIT(&mutex, &mutexatr);
  }
 
  virtual ~MemPool() { MUTEX_DESTROY(&mutex); }
 
  virtual int init(
      bool dynamic = true, int pool_num = DEFAULT_POOL_NUM, int item_num_per_pool = DEFAULT_ITEM_NUM_PER_POOL) = 0;
 
  virtual void cleanup() = 0;
 
  virtual int extend() = 0;
 
  virtual T *alloc() = 0;
 
  virtual void free(T *item) = 0;
 
  virtual string to_string() = 0;
 
  const string get_name() const { return name; }
  bool         is_dynamic() const { return dynamic; }
  int          get_size() const { return size; }
 
protected:
  pthread_mutex_t mutex;
  int             size;
  bool            dynamic;
  string          name;
};
 
template <class T>
class MemPoolSimple : public MemPool<T>
{
public:
  MemPoolSimple(const char *tag) : MemPool<T>(tag) {}
 
  virtual ~MemPoolSimple() { cleanup(); }
 
  int init(bool dynamic = true, int pool_num = DEFAULT_POOL_NUM, int item_num_per_pool = DEFAULT_ITEM_NUM_PER_POOL);
 
  void cleanup();
 
  int extend();
 
  T *alloc();
 
  void free(T *item);
 
  string to_string();
 
  int get_item_num_per_pool() const { return item_num_per_pool; }
 
  int get_used_num()
  {
    MUTEX_LOCK(&this->mutex);
    auto num = used.size();
    MUTEX_UNLOCK(&this->mutex);
    return num;
  }
 
protected:
  list<T *> pools;
  set<T *>  used;
  list<T *> frees;
  int       item_num_per_pool;
 
private:
  inline void asan_poison(void *addr, size_t size) { ASAN_POISON_MEMORY_REGION(addr, size); }
  inline void asan_unpoison(void *addr, size_t size) { ASAN_UNPOISON_MEMORY_REGION(addr, size); }
};
 
template <class T>
int MemPoolSimple<T>::init(bool dynamic, int pool_num, int item_num_per_pool)
{
  if (pools.empty() == false) {
    LOG_WARN("Memory pool has been initialized, but still begin to be initialized, this->name:%s.", this->name.c_str());
    return 0;
  }
 
  if (pool_num <= 0 || item_num_per_pool <= 0) {
    LOG_ERROR("Invalid arguments,  pool_num:%d, item_num_per_pool:%d, this->name:%s.",
              pool_num, item_num_per_pool, this->name.c_str());
    return -1;
  }
 
  this->item_num_per_pool = item_num_per_pool;
  // in order to init memory pool, enable dynamic here
  this->dynamic = true;
  for (int i = 0; i < pool_num; i++) {
    if (extend() < 0) {
      cleanup();
      return -1;
    }
  }
  this->dynamic = dynamic;
 
  LOG_INFO("Extend one pool, this->size:%d, item_num_per_pool:%d, this->name:%s.",
           this->size, item_num_per_pool, this->name.c_str());
  return 0;
}
 
template <class T>
void MemPoolSimple<T>::cleanup()
{
  if (pools.empty() == true) {
    LOG_WARN("Begin to do cleanup, but there is no memory pool, this->name:%s!", this->name.c_str());
    return;
  }
  MUTEX_LOCK(&this->mutex);
  for (auto &&i : frees) {
    asan_unpoison(i, sizeof(T));
  }
  used.clear();
  frees.clear();
  this->size = 0;
 
  for (typename list<T *>::iterator iter = pools.begin(); iter != pools.end(); iter++) {
    T *pool = *iter;
    delete[] pool;
  }
  pools.clear();
  MUTEX_UNLOCK(&this->mutex);
  LOG_INFO("Successfully do cleanup, this->name:%s.", this->name.c_str());
}
 
template <class T>
int MemPoolSimple<T>::extend()
{
  if (this->dynamic == false) {
    LOG_ERROR("Disable dynamic extend memory pool, but begin to extend, this->name:%s", this->name.c_str());
    return -1;
  }
 
  MUTEX_LOCK(&this->mutex);
  T *pool = new T[item_num_per_pool];
  if (pool == nullptr) {
    MUTEX_UNLOCK(&this->mutex);
    LOG_ERROR("Failed to extend memory pool, this->size:%d, item_num_per_pool:%d, this->name:%s.",
              this->size, item_num_per_pool, this->name.c_str());
    return -1;
  }
 
  pools.push_back(pool);
  this->size += item_num_per_pool;
  for (int i = 0; i < item_num_per_pool; i++) {
    frees.push_back(pool + i);
    asan_poison(pool + i, sizeof(T));
  }
  MUTEX_UNLOCK(&this->mutex);
 
  LOG_INFO("Extend one pool, this->size:%d, item_num_per_pool:%d, this->name:%s.",
           this->size, item_num_per_pool, this->name.c_str());
  return 0;
}
 
template <class T>
T *MemPoolSimple<T>::alloc()
{
  MUTEX_LOCK(&this->mutex);
  if (frees.empty() == true) {
    if (this->dynamic == false) {
      MUTEX_UNLOCK(&this->mutex);
      return nullptr;
    }
 
    if (extend() < 0) {
      MUTEX_UNLOCK(&this->mutex);
      return nullptr;
    }
  }
  T *buffer = frees.front();
  asan_unpoison(buffer, sizeof(T));
  frees.pop_front();
 
  used.insert(buffer);
 
  MUTEX_UNLOCK(&this->mutex);
  buffer->reinit();
  return buffer;
}
 
template <class T>
void MemPoolSimple<T>::free(T *buf)
{
  buf->reset();
 
  MUTEX_LOCK(&this->mutex);
 
  size_t num = used.erase(buf);
  if (num == 0) {
    MUTEX_UNLOCK(&this->mutex);
    LOG_WARN("No entry of %p in %s.", buf, this->name.c_str());
    print_stacktrace();
    return;
  }
  asan_poison(buf, sizeof(T));
  frees.push_back(buf);
 
  MUTEX_UNLOCK(&this->mutex);
  return;  // TODO for test
}
 
template <class T>
string MemPoolSimple<T>::to_string()
{
  stringstream ss;
 
  ss << "name:" << this->name << ","
     << "dyanmic:" << this->dynamic << ","
     << "size:" << this->size << ","
     << "pool_size:" << this->pools.size() << ","
     << "used_size:" << this->used.size() << ","
     << "free_size:" << this->frees.size();
  return ss.str();
}
 
class MemPoolItem
{
public:
  using item_unique_ptr = unique_ptr<void, function<void(void *const)>>;
 
public:
  MemPoolItem(const char *tag) : name(tag)
  {
    this->size = 0;
 
    pthread_mutexattr_t mutexatr;
    pthread_mutexattr_init(&mutexatr);
    pthread_mutexattr_settype(&mutexatr, PTHREAD_MUTEX_RECURSIVE);
 
    MUTEX_INIT(&mutex, &mutexatr);
  }
 
  virtual ~MemPoolItem()
  {
    cleanup();
    MUTEX_DESTROY(&mutex);
  }
 
  int init(int item_size, bool dynamic = true, int pool_num = DEFAULT_POOL_NUM,
      int item_num_per_pool = DEFAULT_ITEM_NUM_PER_POOL);
 
  void cleanup();
 
  int extend();
 
  void           *alloc();
  item_unique_ptr alloc_unique_ptr();
 
  void free(void *item);
 
  bool is_used(void *item)
  {
    MUTEX_LOCK(&mutex);
    auto it = used.find(item);
    MUTEX_UNLOCK(&mutex);
    return it != used.end();
  }
 
  string to_string()
  {
 
    stringstream ss;
 
    ss << "name:" << this->name << ","
       << "dyanmic:" << this->dynamic << ","
       << "size:" << this->size << ","
       << "pool_size:" << this->pools.size() << ","
       << "used_size:" << this->used.size() << ","
       << "free_size:" << this->frees.size();
    return ss.str();
  }
 
  const string get_name() const { return name; }
  bool         is_dynamic() const { return dynamic; }
  int          get_size() const { return size; }
  int          get_item_size() const { return item_size; }
  int          get_item_num_per_pool() const { return item_num_per_pool; }
 
  int get_used_num()
  {
    MUTEX_LOCK(&mutex);
    auto num = used.size();
    MUTEX_UNLOCK(&mutex);
    return num;
  }
 
protected:
  pthread_mutex_t mutex;
  string          name;
  bool            dynamic;
  int             size;
  int             item_size;
  int             item_num_per_pool;
 
  list<void *> pools;
  set<void *>  used;
  list<void *> frees;
};
 
}  // namespace common