深入理解 PostgreSQL 中的内存上下文（MmeoryContext）

typedef struct MemoryContextMethods
{
    void    *(*alloc) (MemoryContext context, Size size);                    // 内存分配
    void     (*free_p) (MemoryContext context, void *pointer);               // 内存释放
    void    *(*realloc) (MemoryContext context, void *pointer, Size size);   // 内存重分配
    void     (*reset) (MemoryContext context);                               // 内存重置
    void     (*delete_context) (MemoryContext context);                      // 删除某个内存上下文
    Size     (*get_chunk_space) (MemoryContext context, void *pointer);      // 获取内存片大小
    bool     (*is_empty) (MemoryContext context);                            // 判断内存上下文是否为空
    void     (*stats) (MemoryContext context,                                
                          MemoryStatsPrintFunc printfunc, void *passthru,
                          MemoryContextCounters *totals, bool print_to_stderr);
} MemoryContextMethods;

typedef struct AllocSetContext
{
    MemoryContextData header;       // header 信息，保存了内存上下文之间的关联关系
    
    AllocBlock	blocks;             // 当前内存上下文中所有内存块所组成的双向链表
    AllocChunk	freelist[ALLOCSET_NUM_FREELISTS];	// 当前内存上下文中空闲内存片的数组
    
    Size        initBlockSize;	    // 初始内存块的大小
    Size        maxBlockSize;	    // 允许申请的最大内存块大小
    Size        nextBlockSize;	    // 下一个要分配的内存块大小
    Size        allocChunkLimit;    // 分配内存片的尺寸阈值	
    
    AllocBlock	keeper;             // 保留在 keeper 中的内存块在上下文重置时会保留，只做重置操作，而不进行删除
                                    
    int         freeListIndex;      // 在 context_freelists 中的顺序。
                                    // 0 表示默认 freeList，1 表示小内存 freeList，-1 表示不需要进入 freeList
} AllocSetContext;

set->nextBlockSize <<= 1;                       // 设置为上一次的两倍
if (set->nextBlockSize > set->maxBlockSize)
    set->nextBlockSize = set->maxBlockSize;     // 若超过阈值，则按阈值计算

typedef struct AllocBlockData *AllocBlock;
typedef struct AllocChunkData *AllocChunk;

typedef struct AllocBlockData
{
    AllocSet	aset;           // 该内存块所处的 AllocSetContext
    AllocBlock	prev;		// 前驱指针
    AllocBlock	next;		// 后继指针
    char       *freeptr;	// 指向该内存块空闲区域的首地址
    char       *endptr;		// 指向该内存块空闲区域的结束地址
} AllocBlockData;

typedef struct AllocChunkData
{
    Size    size;           // 内存片的实际大小，以 2 的幂为大小进行向上取整
    Size    requested_size; // debug 使用
    
    void   *aset;           // 该指针有两个作用，使用时指向 AllocSet，空闲时作为 next 指针链接其空闲链表
} AllocChunkData;

void * palloc(Size size)
{
	void	   *ret;
	
	// 在当前内存上下文中进行内存分配
	MemoryContext context = CurrentMemoryContext;

	// 将 isReset 标志位设置为 false，那么在释放内存上下文时就需要清理其内存
	context->isReset = false;
	
	// 此处为多态实现，目前只有 AllocSetAlloc() 这一个实现
	ret = context->methods->alloc(context, size);

	if (unlikely(ret == NULL))
	{
	    // 此处将打印 OOM 错误信息
	}

	return ret;
}

// 若 size 大于 allocChunkLimit 阈值，则需申请新的内存块
if (size > set->allocChunkLimit)
{
	// 字节对齐
	chunk_size = MAXALIGN(size);
	
	// 除开申请的内存片以外，还需要为 AllocBlockData 和 AllocChunkData 预留空间
	blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
	
	// 向 OS 申请内存，这里使用的是 malloc
	block = (AllocBlock) malloc(blksize);
	if (block == NULL)
		return NULL;
	
	// 更新计数器
	context->mem_allocated += blksize;
	
	// 设置内存块信息
	block->aset = set;
	// 此时将不会有空闲空间用于切割更多的 Chunk 了，因此 freeptr 和 endptr 将指向末尾
	block->freeptr = block->endptr = ((char *) block) + blksize;
	
	// 设置内存片信息
	chunk = (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ);
	chunk->aset = set;
	chunk->size = chunk_size;
	
	// 需要将该内存块添加至上下文的 blocks 双向链表中，以便后续的释放
	if (set->blocks != NULL)
	{
		// 链表头不为空，则将该 Block 插入至双向链表的第二个位置。
		// 这是因为链表的第一个 Block 通常用于切割 Chunk，而该 Block 显然没有能力再切割
		block->prev = set->blocks;
		block->next = set->blocks->next;
		if (block->next)
			block->next->prev = block;
		set->blocks->next = block;
	}
	else
	{
		// 链表头为空，则将其插入至链表头
		block->prev = NULL;
		block->next = NULL;
		set->blocks = block;
	}

	return AllocChunkGetPointer(chunk);
}

// size <= allocChunkLimit

// 计算 size 所对应的 freelist index
fidx = AllocSetFreeIndex(size);
// 获取 index 所对应的内存片链表
chunk = set->freelist[fidx];
if (chunk != NULL)
{
	Assert(chunk->size >= size);
	
	// 使 header 指向下一个链表元素
	set->freelist[fidx] = (AllocChunk) chunk->aset;
	
	// 设置内存上下文信息
	chunk->aset = (void *) set;

	return AllocChunkGetPointer(chunk);
}

// 此时需要创建一个新的内存块
if (block == NULL)
{
	Size		required_size;

	blksize = set->nextBlockSize;
	
	// 设置下一次要分配的 Block 大小为当前的 2 倍
	set->nextBlockSize <<= 1;
	// 若超过了阈值，则按阈值取值
	if (set->nextBlockSize > set->maxBlockSize)
		set->nextBlockSize = set->maxBlockSize;
	
	// corner case: 若 blksize 小于预设的 Block 大小，则需要继续扩大 Block 大小
	required_size = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
	while (blksize < required_size)
		blksize <<= 1;

	/* Try to allocate it */
	block = (AllocBlock) malloc(blksize);
	
	// 申请失败时使用更小的 blksize 进行重试
	while (block == NULL && blksize > 1024 * 1024)
	{
		blksize >>= 1;
		if (blksize < required_size)
			break;
		block = (AllocBlock) malloc(blksize);
	}

	if (block == NULL)
		return NULL;
	
	// 更新计数器
	context->mem_allocated += blksize;
	
	// 设置内存块信息
	block->aset = set;
	block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
	block->endptr = ((char *) block) + blksize;

	/* Mark unallocated space NOACCESS. */
	VALGRIND_MAKE_MEM_NOACCESS(block->freeptr,
							   blksize - ALLOC_BLOCKHDRSZ);
	
	// 将该 block 插入至 blocks 链表的头部，而非第二个位置，因为下次 Chunk 切割需要从此 Block 开始
	block->prev = NULL;
	block->next = set->blocks;
	if (block->next)
		block->next->prev = block;
	set->blocks = block;
}

// 从 freeptr 开始切割出一个 Chunk
chunk = (AllocChunk) (block->freeptr);

// 更新 freeptr 指针指向
block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
Assert(block->freeptr <= block->endptr);

// 设置内存片信息
chunk->aset = (void *) set;
chunk->size = chunk_size;

return AllocChunkGetPointer(chunk);

static inline MemoryContext
MemoryContextSwitchTo(MemoryContext context)
{
	MemoryContext old = CurrentMemoryContext;
	CurrentMemoryContext = context;
	return old;
}

// 根据现有分布策略复制一个出来
GpPolicy *root_dist = GpPolicyCopy(rel_dist);
// 修改分布策略的 segments 数量
root_dist->numsegments = new_numsegments;

// 将 relation 的分布策略替换成 root_dist
GpPolicyReplace(relid, root_dist);
// 更新 relation 的 rd_cdbpolicy 指针
rel->rd_cdbpolicy = root_dist;

// 通过 GetMemoryChunkContext 方法获取到 rel 所在内存上下文
oldcontext = MemoryContextSwitchTo(GetMemoryChunkContext(rel));
// 在 rel 所在内存上下文中复制出一个 GpPolicy 出来
new_policy = GpPolicyCopy(rel->rd_cdbpolicy); 
new_policy->numsegments = new_numsegments;
// 切换至原有内存上下文，此处为 PortalHeapMemory
MemoryContextSwitchTo(oldcontext);

GpPolicyReplace(relid, new_policy);
rel->rd_cdbpolicy = new_policy;