include/MemPool.hpp

   1 #ifndef SCX_MEM_POOL
   2 #define SCX_MEM_POOL
   3
   4 #ifdef SCX_LOG
   5 #include <iostream>
   6 #endif
   7 #include <cstdlib>
   8 #include "Alias.hpp"
   9
  10 namespace scx {
  11
  12 /**
  13  * The BLOCK_STEP at least should be sizeof(void*).
  14  * It means on amd64 the min step is 8, while on i386 the min step is 4.
  15  */
  16 template<int BLOCK_STEP = 8, int BLOCK_COUNT = 16>
  17 class MemPool {
  18 public:
  19     explicit MemPool(int secCount = 1024):
  20         mSecCount(secCount),
  21         mMemList(NULL),
  22         mPoolSize(0)
  23     {
  24         Init();
  25     }
  26
  27     ~MemPool()
  28     {
  29         Free();
  30     }
  31
  32 private:
  33     void Init()
  34     {
  35         for (int i = 0; i < BLOCK_COUNT; ++i) {
  36             mFreeBlocks[i] = NULL;
  37         }
  38     }
  39
  40 public:
  41     /**
  42      * If the size is large than the max block size,<br/>
  43      * we'll directly use malloc(). <br/>
  44      * So the potential problem is: <br/>
  45      * you should always remember to call Reuse() to free such memory, <br/>
  46      * as such memory won't be record in mMemList(so, it won't be freed by Free()/dtor).
  47      */
  48     void* Alloc(size_t size)
  49     {
  50         size = SCX_ALIGN(size, BLOCK_STEP);
  51         int index = (size-1)/BLOCK_STEP;
  52         if (index < BLOCK_COUNT) {
  53             if (mFreeBlocks[index] == NULL) {
  54 #ifdef SCX_LOG
  55                 std::cout << "." << std::flush;
  56 #endif
  57                 size_t total = sizeof(Mem) + mSecCount*size;
  58                 Mem* mem = (Mem*)malloc(total);
  59                 mem->buf = (char*)mem + sizeof(Mem);
  60                 mem->next = mMemList;
  61                 mMemList = mem;
  62                 mPoolSize += total;
  63                 Sec* sec = (Sec*)mMemList->buf;
  64                 mFreeBlocks[index] = sec;
  65                 for (int i = 0; i < mSecCount-1; ++i) {
  66                     sec->next = (Sec*)((char*)sec+size);
  67                     sec = sec->next;
  68                 }
  69                 sec->next = NULL;
  70             }
  71             Sec* sec = mFreeBlocks[index];
  72             mFreeBlocks[index] = sec->next;
  73             return sec;
  74         } else {
  75 #ifdef SCX_LOG
  76             std::cout << "!" << size << std::flush;
  77 #endif
  78             mPoolSize += size;
  79             return malloc(size);
  80         }
  81     }
  82
  83     void Reuse(void* p, int size)
  84     {
  85         size = SCX_ALIGN(size, BLOCK_STEP);
  86         int index = (size-1)/BLOCK_STEP;
  87         if (index < BLOCK_COUNT) {
  88             Sec* sec = (Sec*)p;
  89             sec->next = mFreeBlocks[index];
  90             mFreeBlocks[index] = sec;
  91         } else {
  92             free(p);
  93         }
  94     }
  95
  96     /**
  97      * Free all the memory we've allocated.
  98      */
  99     void Free()
 100     {
 101         for (Mem* mem = mMemList; mMemList != NULL; mem = mMemList) {
 102             mMemList = mem->next;
 103             free(mem);
 104         }
 105         mPoolSize = 0;
 106         Init();
 107     }
 108
 109     /**
 110      * The pool size means how many bytes we've allocated,<br/>
 111      * not the available size for Alloc() before next allocation.
 112      */
 113     size_t Size() const
 114     {
 115         return mPoolSize;
 116     }
 117
 118     void SetSecCount(int secCount)
 119     {
 120         mSecCount = secCount;
 121     }
 122
 123     int GetSecCount() const
 124     {
 125         return mSecCount;
 126     }
 127
 128     void PreAlloc(size_t size)
 129     {
 130
 131     }
 132
 133 private:
 134     struct Sec
 135     {
 136         Sec* next;
 137     };
 138
 139     struct Mem
 140     {
 141         void* buf;
 142         Mem* next;
 143     };
 144
 145     int mSecCount;
 146     Mem* mMemList;
 147     size_t mPoolSize;
 148     Sec* mFreeBlocks[BLOCK_COUNT];
 149 };
 150
 151 }
 152 #endif