oceanbase/deps/easy/test/memory/easy_mem_pool_test.c

/**
 * Copyright (c) 2021 OceanBase
 * OceanBase CE is licensed under Mulan PubL v2.
 * You can use this software according to the terms and conditions of the Mulan PubL v2.
 * You may obtain a copy of Mulan PubL v2 at:
 *          http://license.coscl.org.cn/MulanPubL-2.0
 * 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 PubL v2 for more details.
 */

#include <pthread.h>
#include "memory/easy_mem_pool.h"
#include <easy_test.h>

#if __WORDSIZE == 64 && __GNUC__ >= 4
#define ATOMIC_ADD(val, addv) (void)__sync_fetch_and_add((val), addv)

const int64_t thread_num = 4;
const int64_t limit = 10 * 1024 * 1024;
volatile int64_t ref_cnt_free = 0;
volatile int64_t ref_cnt_assert = 0;
volatile int64_t ref_cnt_total = 0;

typedef struct buf_t {
  struct buf_t* next;
} buf_t;

typedef void (*set_memlimit_pt)(int64_t);
typedef void* (*realloc_pt)(void*, size_t);
typedef int64_t (*get_memtotal_pt)(void);

void test_realloc(realloc_pt mp_realloc, get_memtotal_pt mp_get_memtotal, int64_t size, int64_t add_size, char checker)
{
  char* old_buf = (char*)easy_mempool_global_realloc(NULL, size);
  EXPECT_TRUE(NULL != old_buf);
  memset(old_buf, checker, size);

  char* new_buf = (char*)easy_mempool_global_realloc(old_buf, size + add_size);
  EXPECT_TRUE(NULL != new_buf);
  EXPECT_NE(old_buf, new_buf);
  int64_t i = 0;

  for (i = 0; i < size; i++) {
    EXPECT_EQ(checker, new_buf[i]);
  }

  easy_mempool_global_realloc(new_buf, 0);
}

void test_dup_realloc(realloc_pt mp_realloc, get_memtotal_pt mp_get_memtotal)
{
  const int64_t max_size = EASY_MEMPOOL_PAGE_SIZE;
  int64_t i = 1;
  int8_t c = 0;

  for (i = max_size - 32; i < max_size + 32; i++) {
    test_realloc(mp_realloc, mp_get_memtotal, i, 1, c++);
  }
}

void test_limit(
    set_memlimit_pt mp_set_memlimit, realloc_pt mp_realloc, get_memtotal_pt mp_get_memtotal, int64_t volatile* total)
{
  int64_t size = 2 * sizeof(buf_t);
  buf_t* head = NULL;
  mp_set_memlimit(limit);

  while (1) {
    buf_t* buf = (buf_t*)mp_realloc(NULL, size);

    if (NULL == buf) {
      break;
    }

    buf->next = head;
    head = buf;
    *(buf_t*)((char*)buf + size - sizeof(buf_t)) = *buf;
    ATOMIC_ADD(total, size);
    size += 1;
  }

  ATOMIC_ADD(&ref_cnt_assert, 1);
  ATOMIC_ADD(&ref_cnt_total, -1);

  while (0 != ref_cnt_total)
    ;

  EXPECT_TRUE(*total + size > limit);
  EXPECT_EQ(*total, mp_get_memtotal());
  ATOMIC_ADD(&ref_cnt_free, 1);
  ATOMIC_ADD(&ref_cnt_assert, -1);

  while (0 != ref_cnt_assert)
    ;

  buf_t* iter = head;

  while (NULL != iter) {
    size -= 1;
    buf_t* next = iter->next;
    buf_t* check = (buf_t*)((char*)iter + size - sizeof(buf_t));
    EXPECT_TRUE(0 == memcmp(check, iter, sizeof(buf_t)));
    mp_realloc(iter, 0);
    ATOMIC_ADD(total, 0 - size);
    iter = next;
  }

  EXPECT_EQ(2 * sizeof(buf_t), size);
  ATOMIC_ADD(&ref_cnt_free, -1);

  while (0 != ref_cnt_free)
    ;

  EXPECT_EQ(0, *total);
  EXPECT_EQ(2 * sizeof(buf_t), size);
}

void* test_set_thread_memlimit_thread_func(void* data)
{
  int64_t volatile* total_ptr = (int64_t volatile*)data;
  test_limit(
      easy_mempool_set_thread_memlimit, easy_mempool_thread_realloc, easy_mempool_get_thread_memtotal, total_ptr);
  return NULL;
}

void* test_thread_realloc_thread_func(void* data)
{
  test_dup_realloc(easy_mempool_thread_realloc, easy_mempool_get_thread_memtotal);
  return NULL;
}

void* test_set_global_memlimit_thread_func(void* data)
{
  int64_t volatile* total_ptr = (int64_t volatile*)data;
  test_limit(
      easy_mempool_set_global_memlimit, easy_mempool_global_realloc, easy_mempool_get_global_memtotal, total_ptr);
  return NULL;
}

void* test_global_realloc_thread_func(void* data)
{
  test_dup_realloc(easy_mempool_global_realloc, easy_mempool_get_global_memtotal);
  return NULL;
}

typedef struct buf_list_t {
  volatile buf_t* head;
  buf_t* tail;
  int64_t total;
  pthread_spinlock_t lock;
  easy_mempool_t* pool;
} buf_list_t;

void* producer(void* data)
{
  void* tmp = easy_mempool_thread_realloc(NULL, 1);
  buf_list_t* list = (buf_list_t*)data;
  int loop = 1;
  uint32_t seed = 0;

  while (loop) {
    int32_t size = rand_r(&seed) % (32 * 1024) + sizeof(buf_t);
    buf_t* buf = (buf_t*)easy_mempool_alloc(list->pool, size);

    // buf_t *buf = easy_mempool_global_realloc(NULL, size);
    // buf_t *buf = easy_mempool_thread_realloc(NULL, size);
    if (NULL != buf) {
      // memset(buf, 0, sizeof(buf_t));
      pthread_spin_lock(&(list->lock));
      buf->next = NULL;

      if (NULL != list->tail) {
        list->tail->next = buf;
      }

      list->tail = buf;

      if (NULL == list->head) {
        list->head = list->tail;
      }

      list->total--;

      if (0 == list->total) {
        loop = 0;
      }

      pthread_spin_unlock(&(list->lock));
    } else {
      // usleep(1);
    }
  }

  easy_mempool_thread_realloc(tmp, 0);
  return NULL;
}

void* consummer(void* data)
{
  buf_list_t* list = (buf_list_t*)data;
  int loop = 1;

  while (loop) {
    buf_t* buf = NULL;

    if (NULL == list->head && 0 != list->total) {
      continue;
    }

    pthread_spin_lock(&(list->lock));
    buf = (buf_t*)(list->head);

    if (NULL != buf) {
      list->head = buf->next;

      if (buf == list->tail) {
        list->tail = NULL;
      }
    } else {
      if (0 == list->total) {
        loop = 0;
      }
    }

    pthread_spin_unlock(&(list->lock));

    if (NULL != buf) {
      easy_mempool_free(list->pool, buf);
      // easy_mempool_global_realloc(buf, 0);
      // easy_mempool_thread_realloc(buf, 0);
    }
  }

  return NULL;
}

void test_thread_press(int64_t total, int64_t thread)
{
  buf_list_t* list = (buf_list_t*)easy_malloc(sizeof(buf_list_t) * thread);
  pthread_t* p_pd = (pthread_t*)easy_malloc(sizeof(pthread_t) * thread);
  pthread_t* c_pd = (pthread_t*)easy_malloc(sizeof(pthread_t) * thread);
  easy_mempool_t* pool = easy_mempool_create(0);
  int64_t i = 0;

  for (i = 0; i < thread; i++) {
    list[i].head = NULL;
    list[i].tail = NULL;
    list[i].total = total;
    pthread_spin_init(&(list[i].lock), PTHREAD_PROCESS_PRIVATE);
    list[i].pool = pool;
    pthread_create(&(p_pd[i]), NULL, producer, &list[i]);
  }

  for (i = 0; i < thread; i++) {
    pthread_create(&(c_pd[i]), NULL, consummer, &list[i]);
  }

  for (i = 0; i < thread; i++) {
    pthread_join(p_pd[i], NULL);
  }

  for (i = 0; i < thread; i++) {
    pthread_join(c_pd[i], NULL);
  }

  EXPECT_EQ(0, easy_mempool_get_memtotal(pool));
  easy_mempool_destroy(pool);
  easy_free(c_pd);
  easy_free(p_pd);
}

TEST(easy_mem_pool, set_global_memlimit)
{
  volatile int64_t total = 0;
  int64_t i = 0;
  ref_cnt_total = thread_num;
  pthread_t pd[thread_num];

  for (i = 0; i < thread_num; i++) {
    pthread_create(&(pd[i]), NULL, test_set_global_memlimit_thread_func, (void*)&total);
  }

  for (i = 0; i < thread_num; i++) {
    pthread_join(pd[i], NULL);
  }

  EXPECT_EQ(0, easy_mempool_get_global_memtotal());
}

TEST(easy_mem_pool, set_thread_memlimit)
{
  volatile int64_t total = 0;
  int64_t i = 0;
  ref_cnt_total = thread_num;
  pthread_t pd[thread_num];

  for (i = 0; i < thread_num; i++) {
    pthread_create(&(pd[i]), NULL, test_set_thread_memlimit_thread_func, (void*)&total);
  }

  for (i = 0; i < thread_num; i++) {
    pthread_join(pd[i], NULL);
  }

  EXPECT_EQ(0, easy_mempool_get_thread_memtotal());
}

TEST(easy_mem_pool, global_realloc)
{
  int64_t i = 0;
  const int64_t thread_num = 4;
  pthread_t pd[thread_num];

  for (i = 0; i < thread_num; i++) {
    pthread_create(&(pd[i]), NULL, test_global_realloc_thread_func, NULL);
  }

  for (i = 0; i < thread_num; i++) {
    pthread_join(pd[i], NULL);
  }

  EXPECT_EQ(0, easy_mempool_get_global_memtotal());
}

TEST(easy_mem_pool, thread_realloc)
{
  int64_t i = 0;
  const int64_t thread_num = 4;
  pthread_t pd[thread_num];

  for (i = 0; i < thread_num; i++) {
    pthread_create(&(pd[i]), NULL, test_thread_realloc_thread_func, NULL);
  }

  for (i = 0; i < thread_num; i++) {
    pthread_join(pd[i], NULL);
  }

  EXPECT_EQ(0, easy_mempool_get_thread_memtotal());
}

TEST(easy_mem_pool, alloc_limit)
{
  easy_mempool_t* pool = easy_mempool_create(0);

  int64_t i = 0;
  int64_t total = 0;
  int64_t limit = 1024 * 1024;
  buf_t* head = NULL;
  easy_mempool_set_memlimit(pool, limit);

  for (i = 0; total < limit; i++) {
    int32_t alloc_size = sizeof(buf_t) + i;
    buf_t* buf = (buf_t*)easy_mempool_alloc(pool, alloc_size);

    if (total + alloc_size > limit) {
      assert(NULL == buf);
    } else {
      assert(NULL != buf);
    }

    if (NULL == buf) {
      break;
    }

    total += alloc_size;
    buf->next = head;
    head = buf;
  }

  buf_t* iter = head;

  while (NULL != iter) {
    buf_t* next = iter->next;
    easy_mempool_free(pool, iter);
    iter = next;
  }

  easy_mempool_destroy(pool);
}

TEST(easy_mem_pool, thread_press)
{
  test_thread_press(1000000, 2);
}
#endif