oceanbase

1
#include <assert.h>
2
#include <stdio.h>
3
#include <malloc.h>
4
#include <stdlib.h>
5
#include <pthread.h>
6
#include <stdint.h>
7
#include "memory/easy_mem_pool.h"
8

9
#if __GNUC__ >= 4
10
/**
11
 * mem buffer
12
 */
13

14
// 私有数据结构
15

16
enum {
17
    EASY_MEMPOOL_ALLOC = 0,     // 内存由pool申请
18
    EASY_MEMPOOL_DIRECT_ALLOC = 1,   // 内存直接从allocator申请
19
};
20

21
typedef union easy_mempool_atomic_t {
22
    volatile uint64_t       atomic;
23
    struct {
24
        volatile int32_t        ref_cnt;
25
        volatile int32_t        seq_num;
26
    };
27
} easy_mempool_atomic_t;
28

29
// page结构 后面直接跟page内存
30
typedef struct easy_mempool_page_t {
31
    volatile int32_t        base;
32
    int32_t                 reserve;
33
    struct easy_mempool_page_t *next;
34
} easy_mempool_page_t;
35

36
// 管理page的元数据结构 16Byte
37
typedef struct easy_mempool_page_meta_t {
38
    union {
39
        volatile uint64_t       atomic;
40
        struct {
41
            volatile int32_t        ref_cnt;
42
            volatile int32_t        seq_num;
43
        };
44
    };
45
    easy_mempool_page_t *volatile page;
46
} easy_mempool_page_meta_t;
47

48
// 分配内存的头部结构 16Byte
49
typedef struct easy_mempool_buf_t {
50
    uint16_t                magic_num;
51
    uint8_t                 alloc_type;
52
    uint8_t                 reserve1;
53
    uint16_t                page_pos;
54
    uint32_t                size;
55
    uint32_t                reserve2;
56
} easy_mempool_buf_t;
57

58
// 一个pool结构
59
struct easy_mempool_t {
60
    volatile int32_t        cur_page_pos;
61
    volatile int32_t        direct_alloc_cnt;
62
    volatile int64_t        mem_total;
63
    int64_t                 mem_limit;
64

65
    easy_mempool_allocator_t *allocator;
66

67
    int32_t                 page_size;
68
    int32_t                 page_num;
69
    easy_mempool_page_meta_t *page_metas;
70

71
    int32_t                 free_num;
72
    int32_t                 reserve2;
73
    easy_mempool_page_t     *free_list;
74
    pthread_spinlock_t      free_list_lock;
75
};
76

77
typedef struct easy_mempool_thread_info_t {
78
    volatile int64_t        ref_cnt;
79
    easy_mempool_t          *pool;
80
} easy_mempool_thread_info_t;
81

82
easy_mempool_allocator_t easy_mempool_g_allocator = {memalign, free};
83

84
static int64_t          EASY_MEMPOOL_BUF_MAGIC_NUM = 0xabcd;
85
static int64_t          EASY_MEMPOOL_BUF_FREE_FLAG = 0Xef12;
86
static easy_mempool_t   *easy_mempool_g_pool = NULL;
87
static pthread_key_t    easy_mempool_g_thread_key = INT32_MAX;
88
static int64_t          easy_mempool_g_thread_memlimit = INT64_MAX;
89
static volatile int64_t easy_mempool_g_thread_memtotal = 0;
90

91
////////////////////////////////////////////////////////////////////////////////////////////////////
92

93
// 私有方法
94

95
#define ATOMIC_INC(val) (void)__sync_add_and_fetch((val), 1)
96
#define ATOMIC_INC_FETCH(val) __sync_add_and_fetch((val), 1)
97
#define ATOMIC_DEC(val) (void)__sync_sub_and_fetch((val), 1)
98
#define ATOMIC_DEC_FETCH(val) __sync_sub_and_fetch((val), 1)
99
#define ATOMIC_ADD(val, addv) (void)__sync_fetch_and_add((val), addv)
100
#define ATOMIC_SUB(val, addv) (void)__sync_fetch_and_sub((val), addv)
101
#define ATOMIC_CAS(val, cmpv, newv) __sync_val_compare_and_swap((val), (cmpv), (newv))
102

103
static void *easy_mempool_alloc_(easy_mempool_t *pool, uint32_t size, uint32_t align_size);
104

105
static void easy_mempool_destroy_free_list_(easy_mempool_t *pool);
106
static easy_mempool_page_t *easy_mempool_alloc_page_(easy_mempool_t *pool);
107
static void easy_mempool_free_page_(easy_mempool_t *pool, easy_mempool_page_t *page);
108

109
static easy_mempool_page_t *easy_mempool_get_cur_page_(easy_mempool_t *pool, int32_t ensure_size, int32_t *page_pos);
110
static void *easy_mempool_alloc_from_page_(easy_mempool_page_t *page, int32_t page_size, int32_t alloc_size);
111

112
static void easy_mempool_deref_page_(easy_mempool_t *pool, int32_t page_pos);
113
static int32_t easy_mempool_dec_ref_cnt_and_inc_seq_num_(easy_mempool_page_meta_t *page_meta);
114

115
static void easy_mempool_thread_destroy_callback_(void *data);
116

117
////////////////////////////////////////////////////////////////////////////////////////////////////
118

119
void easy_mempool_set_global_memlimit(int64_t limit)
120
{
121
    easy_mempool_set_memlimit(easy_mempool_g_pool, limit);
122
}
123

124
int64_t easy_mempool_get_global_memtotal()
125
{
126
    return easy_mempool_get_memtotal(easy_mempool_g_pool);
127
}
128

129
void *easy_mempool_global_realloc(void *ptr, size_t size)
130
{
131
    void                    *ret = NULL;
132

133
    if (0 != size) {
134
        ret = easy_mempool_alloc(easy_mempool_g_pool, size);
135
    }
136

137
    if (NULL != ptr && NULL != ret) {
138
        easy_mempool_buf_t      *buf = (easy_mempool_buf_t *)((char *)ptr - sizeof(easy_mempool_buf_t));
139

140
        if (NULL == buf) {
141
            easy_mempool_alloc((easy_mempool_t *)ret, 0);
142
        } else {
143
            memcpy(ret, ptr, buf->size > size ? size : buf->size);
144
        }
145
    }
146

147
    easy_mempool_free(easy_mempool_g_pool, ptr);
148
    return ret;
149
}
150

151
void easy_mempool_set_thread_memlimit(int64_t limit)
152
{
153
    if (0 < limit) {
154
        easy_mempool_g_thread_memlimit = limit;
155
    }
156
}
157

158
int64_t easy_mempool_get_thread_memtotal()
159
{
160
    return easy_mempool_g_thread_memtotal;
161
}
162

163
void *easy_mempool_thread_realloc(void *ptr, size_t size)
164
{
165
    void                    *ret = NULL;
166
    void                    *alloc_ptr = NULL;
167

168
    easy_mempool_thread_info_t *thread_info = (easy_mempool_thread_info_t *)pthread_getspecific(easy_mempool_g_thread_key);
169

170
    if (NULL == thread_info) {
171
        thread_info = (easy_mempool_thread_info_t *)easy_mempool_g_allocator.memalign(EASY_MEMPOOL_ALIGNMENT,
172
                      sizeof(easy_mempool_thread_info_t));
173

174
        if (NULL != thread_info) {
175
            thread_info->ref_cnt = 1;
176
            thread_info->pool = easy_mempool_create(0);
177
        }
178

179
        pthread_setspecific(easy_mempool_g_thread_key, thread_info);
180
    }
181

182
    if (0 != size
183
            && NULL != thread_info
184
            && easy_mempool_g_thread_memlimit >= (int64_t)(easy_mempool_g_thread_memtotal + size)) {
185
        alloc_ptr = easy_mempool_alloc(thread_info->pool, size + sizeof(easy_mempool_thread_info_t *));
186

187
        if (NULL != alloc_ptr) {
188
            *(easy_mempool_thread_info_t **)(alloc_ptr) = thread_info;
189
            ret = (char *)alloc_ptr + sizeof(easy_mempool_thread_info_t *);
190
            ATOMIC_INC(&(thread_info->ref_cnt));
191
            ATOMIC_ADD(&easy_mempool_g_thread_memtotal, size);
192
        }
193
    }
194

195
    if (NULL != ptr && NULL != ret) {
196
        easy_mempool_buf_t      *buf = (easy_mempool_buf_t *)((char *)ptr -
197
                                       sizeof(easy_mempool_thread_info_t *) - sizeof(easy_mempool_buf_t));
198

199
        if (NULL == buf) {
200
            easy_mempool_free(thread_info->pool, alloc_ptr);
201
        } else {
202
            memcpy(ret, ptr, buf->size > size ? size : buf->size);
203
        }
204
    }
205

206
    if (NULL != ptr) {
207
        ptr = (char *)ptr - sizeof(easy_mempool_thread_info_t *);
208
        easy_mempool_buf_t      *buf = (easy_mempool_buf_t *)((char *)ptr - sizeof(easy_mempool_buf_t));
209
        easy_mempool_thread_info_t *host = *(easy_mempool_thread_info_t **)ptr;
210

211
        if (NULL != buf) {
212
            ATOMIC_SUB(&easy_mempool_g_thread_memtotal, buf->size - sizeof(easy_mempool_thread_info_t *));
213
        }
214

215
        if (NULL != host) {
216
            easy_mempool_free(host->pool, ptr);
217

218
            if (0 == ATOMIC_DEC_FETCH(&(host->ref_cnt))) {
219
                easy_mempool_destroy(host->pool);
220
                easy_mempool_g_allocator.free(host);
221
            }
222
        }
223
    }
224

225
    return ret;
226
}
227

228
////////////////////////////////////////////////////////////////////////////////////////////////////
229

230
easy_mempool_t *easy_mempool_create(uint32_t size)
231
{
232
    easy_mempool_t          *ret = NULL;
233
    int32_t                 page_num = EASY_MEMPOOL_PAGE_MAX_NUM;
234
    int32_t                 size2alloc = sizeof(easy_mempool_t) + sizeof(easy_mempool_page_t) * page_num;
235

236
    if (NULL != (ret = (easy_mempool_t *)easy_mempool_g_allocator.memalign(EASY_MEMPOOL_ALIGNMENT, size2alloc))) {
237
        ret->cur_page_pos = 0;
238
        ret->direct_alloc_cnt = 0;
239
        ret->mem_total = 0;
240
        ret->mem_limit = INT64_MAX;
241
        ret->allocator = &easy_mempool_g_allocator;
242
        ret->page_size = size ? size : EASY_MEMPOOL_PAGE_SIZE;
243
        ret->page_num = page_num;
244
        ret->page_metas = (easy_mempool_page_meta_t *)((char *)ret + sizeof(easy_mempool_t));
245
        memset(ret->page_metas, 0, sizeof(easy_mempool_page_t) * page_num);
246
        ret->page_metas[ret->cur_page_pos].ref_cnt = 1;
247
        ret->free_num = 0;
248
        ret->free_list = NULL;
249
        pthread_spin_init(&(ret->free_list_lock), PTHREAD_PROCESS_PRIVATE);
250
    }
251

252
    return ret;
253
}
254

255
void easy_mempool_destroy(easy_mempool_t *pool)
256
{
257
    if (NULL != pool) {
258
        easy_mempool_clear(pool);
259
        easy_mempool_g_allocator.free(pool);
260
    }
261
}
262

263
void easy_mempool_clear(easy_mempool_t *pool)
264
{
265
    if (NULL != pool) {
266
        int32_t                 unfreed_num = 0;
267
        int32_t                 i = 0;
268

269
        for (i = 0; i < EASY_MEMPOOL_PAGE_MAX_NUM; i++) {
270
            if (0 == pool->page_metas[i].ref_cnt
271
                    || (i == pool->cur_page_pos && 1 == pool->page_metas[i].ref_cnt)) {
272
                easy_mempool_free_page_(pool, pool->page_metas[i].page);
273
                pool->page_metas[i].page = NULL;
274
            } else {
275
                unfreed_num++;
276
            }
277
        }
278

279
        easy_mempool_destroy_free_list_(pool);
280

281
        if (0 != unfreed_num
282
                || 0 != pool->direct_alloc_cnt) {
283
            fprintf(stderr, "[WARN] there are still %d pool_buf or %d direct_buf used cannot free",
284
                    unfreed_num, pool->direct_alloc_cnt);
285
        }
286
    }
287
}
288

289
void *easy_mempool_alloc(easy_mempool_t *pool, uint32_t size)
290
{
291
    return easy_mempool_alloc_(pool, size, EASY_MEMPOOL_ALIGNMENT);
292
}
293

294
void easy_mempool_free(easy_mempool_t *pool, void *ptr)
295
{
296
    if (NULL != pool
297
            && NULL != ptr) {
298
        easy_mempool_buf_t      *buf = (easy_mempool_buf_t *)((char *)ptr - sizeof(easy_mempool_buf_t));
299

300
        if (EASY_MEMPOOL_BUF_MAGIC_NUM == buf->magic_num) {
301
            int64_t                 size = buf->size;
302
            buf->magic_num = EASY_MEMPOOL_BUF_FREE_FLAG;
303

304
            if (EASY_MEMPOOL_DIRECT_ALLOC == buf->alloc_type) {
305
                pool->allocator->free(buf);
306
                ATOMIC_DEC(&(pool->direct_alloc_cnt));
307
            } else {
308
                easy_mempool_deref_page_(pool, buf->page_pos);
309
            }
310

311
            ATOMIC_SUB(&(pool->mem_total), size);
312
        }
313
    }
314
}
315

316
void easy_mempool_set_memlimit(easy_mempool_t *pool, int64_t limit)
317
{
318
    if (NULL != pool && 0 < limit) {
319
        pool->mem_limit = limit;
320
    }
321
}
322

323
void easy_mempool_set_allocator(easy_mempool_t *pool, easy_mempool_allocator_t *allocator)
324
{
325
    if (NULL != pool
326
            && NULL != allocator->memalign
327
            && NULL != allocator->free) {
328
        pool->allocator = allocator;
329
    }
330
}
331

332
int64_t easy_mempool_get_memtotal(easy_mempool_t *pool)
333
{
334
    int64_t                 ret = 0;
335

336
    if (NULL != pool) {
337
        ret = pool->mem_total;
338
    }
339

340
    return ret;
341
}
342

343
////////////////////////////////////////////////////////////////////////////////////////////////////
344

345
static void *easy_mempool_alloc_(easy_mempool_t *pool, uint32_t size, uint32_t align_size)
346
{
347
    void                    *ret = NULL;
348
    int32_t                 alloc_size = size + sizeof(easy_mempool_buf_t);
349
    alloc_size = easy_mempool_align(alloc_size, align_size);
350

351
    if (NULL != pool) {
352
        if ((pool->mem_total + size) > pool->mem_limit) {
353
            // memory over limit
354
        } else if (pool->page_size < alloc_size) {
355
            easy_mempool_buf_t      *buf = (easy_mempool_buf_t *)pool->allocator->memalign(align_size, alloc_size);
356

357
            if (NULL != buf) {
358
                buf->magic_num = EASY_MEMPOOL_BUF_MAGIC_NUM;
359
                buf->alloc_type = EASY_MEMPOOL_DIRECT_ALLOC;
360
                buf->size = size;
361
                ret = (char *)buf + sizeof(easy_mempool_buf_t);
362
                ATOMIC_INC(&(pool->direct_alloc_cnt));
363
            }
364
        } else {
365
            easy_mempool_page_t     *page = NULL;
366
            easy_mempool_buf_t      *buf = NULL;
367
            int32_t                 page_pos = -1;
368

369
            while (1) {
370
                if (NULL == (page = easy_mempool_get_cur_page_(pool, alloc_size, &page_pos))) {
371
                    break;
372
                }
373

374
                buf = (easy_mempool_buf_t *)easy_mempool_align_ptr(easy_mempool_alloc_from_page_(page, pool->page_size, alloc_size), align_size);
375

376
                if (NULL != buf) {
377
                    ATOMIC_INC(&(pool->page_metas[page_pos].ref_cnt));
378
                }
379

380
                easy_mempool_deref_page_(pool, page_pos);
381

382
                if (NULL != buf) {
383
                    buf->magic_num = EASY_MEMPOOL_BUF_MAGIC_NUM;
384
                    buf->alloc_type = EASY_MEMPOOL_ALLOC;
385
                    buf->page_pos = page_pos;
386
                    buf->size = size;
387
                    ret = (char *)buf + sizeof(easy_mempool_buf_t);
388
                    break;
389
                }
390
            }
391
        }
392

393
        if (NULL != ret) {
394
            ATOMIC_ADD(&(pool->mem_total), size);
395
        }
396
    }
397

398
    return ret;
399
}
400

401
static void easy_mempool_destroy_free_list_(easy_mempool_t *pool)
402
{
403
    if (NULL != pool) {
404
        easy_mempool_page_t     *iter = pool->free_list;
405

406
        while (NULL != iter) {
407
            easy_mempool_page_t     *next = iter->next;
408
            pool->allocator->free(iter);
409
            iter = next;
410
        }
411
    }
412
}
413

414
static easy_mempool_page_t *easy_mempool_alloc_page_(easy_mempool_t *pool)
415
{
416
    easy_mempool_page_t     *ret = NULL;
417

418
    if (NULL != pool) {
419
        pthread_spin_lock(&(pool->free_list_lock));
420
        easy_mempool_page_t     *page = pool->free_list;
421

422
        if (NULL != page) {
423
            pool->free_list = page->next;
424
            pool->free_num--;
425
            ret = page;
426
        }
427

428
        pthread_spin_unlock(&(pool->free_list_lock));
429

430
        if (NULL == ret) {
431
            ret = (easy_mempool_page_t *)pool->allocator->memalign(EASY_MEMPOOL_ALIGNMENT, pool->page_size + sizeof(easy_mempool_page_t));
432
        }
433

434
        if (NULL != ret) {
435
            ret->base = 0;
436
        }
437
    }
438

439
    return ret;
440
}
441

442
static void easy_mempool_free_page_(easy_mempool_t *pool, easy_mempool_page_t *page)
443
{
444
    if (NULL != pool
445
            && NULL != page) {
446
        pthread_spin_lock(&(pool->free_list_lock));
447

448
        if (EASY_MEMPOOL_PAGE_FREE_NUM > pool->free_num) {
449
            page->next = pool->free_list;
450
            pool->free_list = page;
451
            pool->free_num++;
452
            page = NULL;
453
        }
454

455
        pthread_spin_unlock(&(pool->free_list_lock));
456

457
        if (NULL != page) {
458
            pool->allocator->free(page);
459
        }
460
    }
461
}
462

463
static void easy_mempool_deref_page_(easy_mempool_t *pool, int32_t page_pos)
464
{
465
    if (NULL != pool
466
            && pool->page_num > page_pos) {
467
        easy_mempool_page_t     *tmp_page = pool->page_metas[page_pos].page;
468

469
        if (0 == easy_mempool_dec_ref_cnt_and_inc_seq_num_(&(pool->page_metas[page_pos]))) {
470
            if (tmp_page == ATOMIC_CAS(&(pool->page_metas[page_pos].page), tmp_page, NULL)) {
471
                easy_mempool_free_page_(pool, tmp_page);
472
            }
473
        }
474
    }
475
}
476

477
static easy_mempool_page_t *easy_mempool_get_cur_page_(easy_mempool_t *pool, int32_t ensure_size, int32_t *page_pos)
478
{
479
    easy_mempool_page_t     *ret = NULL;
480

481
    if (NULL != pool) {
482
        volatile int32_t        oldv = -1;
483
        volatile int32_t        newv = -1;
484
        volatile int32_t        cmpv = -1;
485
        easy_mempool_page_t     *cur_page = NULL;
486

487
        while (oldv != pool->cur_page_pos) {
488
            oldv = pool->cur_page_pos;
489
            newv = oldv;
490
            cmpv = oldv;
491
            ATOMIC_INC(&(pool->page_metas[oldv].ref_cnt));
492

493
            if (NULL == pool->page_metas[oldv].page) {
494
                easy_mempool_page_t     *tmp_page = easy_mempool_alloc_page_(pool);
495

496
                if (NULL != tmp_page) {
497
                    if (NULL != ATOMIC_CAS(&(pool->page_metas[oldv].page), NULL, tmp_page)) {
498
                        easy_mempool_free_page_(pool, tmp_page);
499
                    }
500
                }
501
            }
502

503
            if (NULL == (cur_page = pool->page_metas[oldv].page)) {
504
                easy_mempool_deref_page_(pool, oldv);
505
                break;
506
            }
507

508
            if ((pool->page_size - cur_page->base) < ensure_size) {
509
                int32_t                 base = cur_page->base;
510
                easy_mempool_deref_page_(pool, oldv);
511

512
                if (0 == base) {
513
                    break;
514
                }
515

516
                int32_t                 counter = 0;
517

518
                while (++counter < pool->page_num) {
519
                    newv = (newv + 1) % pool->page_num;
520

521
                    if (0 == ATOMIC_CAS(&(pool->page_metas[newv].ref_cnt), 0, 1)) {
522
                        if (oldv == ATOMIC_CAS(&(pool->cur_page_pos), cmpv, newv)) {
523
                            easy_mempool_deref_page_(pool, oldv);
524
                        } else {
525
                            easy_mempool_deref_page_(pool, newv);
526
                        }
527

528
                        break;
529
                    }
530
                }
531
            } else {
532
                *page_pos = oldv;
533
                ret = cur_page;
534
                break;
535
            }
536
        }
537
    }
538

539
    return ret;
540
}
541

542
static void *easy_mempool_alloc_from_page_(easy_mempool_page_t *page, int32_t page_size, int32_t alloc_size)
543
{
544
    void                    *ret = NULL;
545

546
    if (NULL != page) {
547
        volatile int32_t        oldv = 0;
548
        volatile int32_t        newv = 0;
549
        volatile int32_t        cmpv = 0;
550

551
        while (1) {
552
            oldv = page->base;
553
            newv = oldv + alloc_size;
554
            cmpv = oldv;
555

556
            if (newv > page_size) {
557
                break;
558
            }
559

560
            if (oldv == ATOMIC_CAS(&(page->base), cmpv, newv)) {
561
                ret = (char *)page + sizeof(easy_mempool_page_t) + oldv;
562
                break;
563
            }
564
        }
565
    }
566

567
    return ret;
568
}
569

570
static int32_t easy_mempool_dec_ref_cnt_and_inc_seq_num_(easy_mempool_page_meta_t *page_meta)
571
{
572
    int32_t                 ret = -1;
573

574
    if (NULL != page_meta) {
575
        easy_mempool_atomic_t   oldv = {0};
576
        easy_mempool_atomic_t   newv = {0};
577
        easy_mempool_atomic_t   cmpv = {0};
578

579
        while (1) {
580
            oldv.atomic = page_meta->atomic;
581
            newv.atomic = oldv.atomic;
582
            cmpv.atomic = oldv.atomic;
583
            newv.ref_cnt -= 1;
584

585
            if (0 == newv.ref_cnt) {
586
                newv.seq_num += 1;
587
            }
588

589
            //assert(0 != oldv.ref_cnt);
590
            if (oldv.atomic == ATOMIC_CAS(&(page_meta->atomic), cmpv.atomic, newv.atomic)) {
591
                ret = newv.ref_cnt;
592
                break;
593
            }
594
        }
595
    }
596

597
    return ret;
598
}
599

600
static void easy_mempool_thread_destroy_callback_(void *data)
601
{
602
    if (NULL != data) {
603
        easy_mempool_thread_info_t *thread_info = (easy_mempool_thread_info_t *)data;
604

605
        if (0 == ATOMIC_DEC_FETCH(&(thread_info->ref_cnt))) {
606
            easy_mempool_destroy(thread_info->pool);
607
            easy_mempool_g_allocator.free(thread_info);
608
        }
609
    }
610
}
611

612
void __attribute__((constructor)) init_global_easy_mempool_()
613
{
614
    easy_mempool_g_pool = easy_mempool_create(0);
615
    pthread_key_create(&easy_mempool_g_thread_key, easy_mempool_thread_destroy_callback_);
616
}
617

618
void __attribute__((destructor)) destroy_global_easy_mempool_()
619
{
620
    easy_mempool_destroy(easy_mempool_g_pool);
621
    void                    *data = pthread_getspecific(easy_mempool_g_thread_key);
622
    easy_mempool_thread_destroy_callback_(data);
623
    pthread_key_delete(easy_mempool_g_thread_key);
624
}
625

626
#else
627

628
#endif
629

630