TDengine代碼學習(2):內存池分配

代碼介紹

學習TDengine tmempool.c 中的代碼，主要學習下內存池的數據結構設計。
這個有點類似於linux 內核中的kmem_cache, 預先創建內存塊大小爲 blockSize，數量爲numOfBlock的內存池，然後每次從內存池中獲取大小爲blockSize的內存塊。

typedef struct {
  int             numOfFree;  /* number of free slots */
  int             first;      /* the first free slot  */
  int             numOfBlock; /* the number of blocks */
  int             blockSize;  /* block size in bytes  */
  int *           freeList;   /* the index list       */
  char *          pool;       /* the actual mem block */
  pthread_mutex_t mutex;
} pool_t;

結構體pool_t中的成員

• pool：是申請的連續內存空間，大小爲 (blockSize * numOfBlock)
• freeList: 一個int 數組，每個項保存指向 pool 中的內存塊的偏移index
• first：第一個可用內存塊的freeList數組index
• numOfFree: 內存池中可用內存塊的數目

freeList 可以看成是一個用來模擬鏈表的數組。index 在 [first, first + numOfFree) 範圍內的freeList 數組的值表示pool中可分配內存塊的偏移index。這裏的 first + numOfFree 值在作爲freeList 數組index時實際是取除以numOfBlock 的餘數，即 (first + numOfFree) % numOfBlock 的值。

內存池分配釋放示例

假設創建了一個numOfBlock 爲6 的內存池，通過一個例子來看下內存池的結構變化。

1. 初始化時的內存池如圖1， first 值爲0，numOfFree值爲6，圖中的 first + numOfFree 值實際是 (first + numOfFree) % numOfBlock，所以值爲0。

圖1

2. 用戶從內存池中分配了3個內存塊給a1，a2，a3以後的內存池如圖2，first 目前值爲3，numOfFree值爲3，freeList 數組中紅色字體的值表示已經分配出去的內存塊偏移值。可以看到，分配內存時， (first + numOfFree) % numOfBlock 的值是保持不變的。

圖2

3. 用戶按順序釋放了a3，a2 以後的內存池如圖3，因爲先釋放了a3，所以freeList[0]的值設爲2。釋放內存時，first的值保持不變。

圖3

測試效果

按照上面例子寫的測試程序運行結果如下，和圖示是一致的。

完整測試代碼

在linux 下面運行測試。
編譯需要加上 -lpthread 選項。例子如下：
gcc -o mempool mempool.c -lpthread

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>

#define mpool_h void *

typedef struct {
   int				numOfFree;
   int 			first;
   int 			numOfBlock;
   int 			blockSize;
   int *			freeList;
   char * 			pool;
   pthread_mutex_t	mutex;
} pool_t;

mpool_h memPoolInit(int numOfBlock, int blockSize)
{
   int i = 0;
   pool_t * pool_p = NULL;

   if(numOfBlock <= 1 || blockSize <= 1)
   {
   	printf("invalid parameter in memPoolInit\n");
   	return NULL;
   }

   pool_p = (pool_t *)malloc(sizeof(pool_t));
   if(pool_p == NULL)
   {
   	printf("mempool malloc failed\n");
   	return NULL;
   }

   memset(pool_p, 0, sizeof(pool_t));

   pool_p->blockSize = blockSize;
   pool_p->numOfBlock = numOfBlock;
   pool_p->pool = (char *)malloc((size_t)(blockSize * numOfBlock));
   pool_p->freeList = (int *)malloc(sizeof(int) * (size_t)numOfBlock);

   if(pool_p->pool == NULL || pool_p->freeList == NULL)
   {
   	printf("failed to allocate memory\n");
   	free(pool_p->freeList);
   	free(pool_p->pool);
   	free(pool_p);
   }

   pthread_mutex_init(&(pool_p->mutex), NULL);

   for(i = 0; i < pool_p->numOfBlock; i++)
   	pool_p->freeList[i] = i;

   pool_p->first = 0;
   pool_p->numOfFree= pool_p->numOfBlock;

   return (mpool_h)pool_p;
}

char * memPoolMalloc(mpool_h handle)
{
   char * pos = NULL;
   pool_t * pool_p = (pool_t *)handle;

   pthread_mutex_lock(&pool_p->mutex);

   if(pool_p->numOfFree <= 0)
   {
   	printf("mempool: out of memory");
   }
   else
   {
   	pos = pool_p->pool + pool_p->blockSize * (pool_p->freeList[pool_p->first]);
   	pool_p->first = (pool_p->first + 1) % pool_p->numOfBlock;
   	pool_p->numOfFree--;
   }

   pthread_mutex_unlock(&pool_p->mutex);
   if(pos != NULL) memset(pos, 0, (size_t)pool_p->blockSize);
   return pos;
}

void memPoolFree(mpool_h handle, char * pMem)
{
   int index = 0;
   pool_t * pool_p = (pool_t *)handle;

   if(pool_p == NULL || pMem == NULL) return;

   pthread_mutex_lock(&pool_p->mutex);

   index = (int)(pMem - pool_p->pool) % pool_p->blockSize;
   if(index != 0)
   {
   	printf("invalid free address:%p\n", pMem);
   }
   else
   {
   	index = (int)((pMem - pool_p->pool) / pool_p->blockSize);
   	if(index < 0 || index >= pool_p->numOfBlock)
   	{
   		printf("mempool: error, invalid address:%p\n", pMem);
   	}
   	else
   	{
   		pool_p->freeList[(pool_p->first + pool_p->numOfFree) % pool_p->numOfBlock] = index;
   		pool_p->numOfFree++;
   		memset(pMem, 0, (size_t)pool_p->blockSize);
   	}
   }
   
   pthread_mutex_unlock(&pool_p->mutex);
}

void memPoolCleanup(mpool_h handle)
{
   pool_t *pool_p = (pool_t *)handle;

   pthread_mutex_destroy(&pool_p->mutex);
   if(pool_p->pool) free(pool_p->pool);
   if(pool_p->freeList) free(pool_p->freeList);
}

void displayMemPool(mpool_h handle)
{
   int i = 0, index = 0;
   pool_t * pool_p = (pool_t *)handle;

   if(pool_p == NULL) return;

   printf("Mempool Info\nfirst:%d, numOfFree:%d\n", pool_p->first, pool_p->numOfFree);

   for(i = pool_p->first; i < (pool_p->first + pool_p->numOfBlock); i++)
   {
   	index = i % pool_p->numOfBlock;
   	if(i >= pool_p->first + pool_p->numOfFree)
   		printf("\tfreeList[%d] => used\n", index);
   	else
   		printf("\tfreeList[%d] => pool[%d]\n", index, pool_p->freeList[index]);
   }
}

#define BLOCK_NUM 6
void testMemPool(mpool_h handle)
{
   int i = 0;

   displayMemPool(handle);

   void * a[BLOCK_NUM] = {0};
   for(i = 0; i < 3; i++)
   	a[i] = memPoolMalloc(handle);

   displayMemPool(handle);
   
   memPoolFree(handle, (char *)a[2]);
   memPoolFree(handle, (char *)a[1]);

   displayMemPool(handle);
}

int main()
{	
   void * pool = memPoolInit(BLOCK_NUM, 100);
   testMemPool(pool);
   memPoolCleanup(pool);
}