Add palloc_aligned() to allow aligned memory allocations

This introduces palloc_aligned() and MemoryContextAllocAligned() which
allow callers to obtain memory which is allocated to the given size and
also aligned to the specified alignment boundary.  The alignment
boundaries may be any power-of-2 value.  Currently, the alignment is
capped at 2^26, however, we don't expect values anything like that large.
The primary expected use case is to align allocations to perhaps CPU
cache line size or to maybe I/O page size.  Certain use cases can benefit
from having aligned memory by either having better performance or more
predictable performance.

The alignment is achieved by requesting 'alignto' additional bytes from
the underlying allocator function and then aligning the address that is
returned to the requested alignment.  This obviously does waste some
memory, so alignments should be kept as small as what is required.

It's also important to note that these alignment bytes eat into the
maximum allocation size.  So something like:

palloc_aligned(MaxAllocSize, 64, 0);

will not work as we cannot request MaxAllocSize + 64 bytes.

Additionally, because we're just requesting the requested size plus the
alignment requirements from the given MemoryContext, if that context is
the Slab allocator, then since slab can only provide chunks of the size
that's specified when the slab context is created, then this is not going
to work.  Slab will generate an error to indicate that the requested size
is not supported.

The alignment that is requested in palloc_aligned() is stored along with
the allocated memory.  This allows the alignment to remain intact through
repalloc() calls.

Author: Andres Freund, David Rowley
Reviewed-by: Maxim Orlov, Andres Freund, John Naylor
Discussion: https://postgr.es/m/CAApHDvpxLPUMV1mhxs6g7GNwCP6Cs6hfnYQL5ffJQTuFAuxt8A%40mail.gmail.com

5 files changed, 284 insertions, 10 deletions

diff --git a/src/backend/utils/cache/catcache.c b/src/backend/utils/cache/catcache.c
index 30ef0ba39c3..9e635177c80 100644
--- a/src/backend/utils/cache/catcache.c
+++ b/src/backend/utils/cache/catcache.c
@@ -763,7 +763,6 @@ InitCatCache(int id,
 {
 	CatCache   *cp;
 	MemoryContext oldcxt;
-	size_t		sz;
 	int			i;
 
 	/*
@@ -807,8 +806,8 @@ InitCatCache(int id,
 	 *
 	 * Note: we rely on zeroing to initialize all the dlist headers correctly
 	 */
-	sz = sizeof(CatCache) + PG_CACHE_LINE_SIZE;
-	cp = (CatCache *) CACHELINEALIGN(palloc0(sz));
+	cp = (CatCache *) palloc_aligned(sizeof(CatCache), PG_CACHE_LINE_SIZE,
+									 MCXT_ALLOC_ZERO);
 	cp->cc_bucket = palloc0(nbuckets * sizeof(dlist_head));
 
 	/*
diff --git a/src/backend/utils/mmgr/Makefile b/src/backend/utils/mmgr/Makefile
index 3b4cfdbd520..dae3432c98a 100644
--- a/src/backend/utils/mmgr/Makefile
+++ b/src/backend/utils/mmgr/Makefile
@@ -13,6 +13,7 @@ top_builddir = ../../../..
 include $(top_builddir)/src/Makefile.global
 
 OBJS = \
+	alignedalloc.o \
 	aset.o \
 	dsa.o \
 	freepage.o \
diff --git a/src/backend/utils/mmgr/alignedalloc.c b/src/backend/utils/mmgr/alignedalloc.c
new file mode 100644
index 00000000000..62f9dbfe89a
--- /dev/null
+++ b/src/backend/utils/mmgr/alignedalloc.c
@@ -0,0 +1,132 @@
+/*-------------------------------------------------------------------------
+ *
+ * alignedalloc.c
+ *	  Allocator functions to implement palloc_aligned
+ *
+ * This is not a fully-fledged MemoryContext type as there is no means to
+ * create a MemoryContext of this type.  The code here only serves to allow
+ * operations such as pfree() and repalloc() to work correctly on a memory
+ * chunk that was allocated by palloc_aligned().
+ *
+ * Portions Copyright (c) 2022, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *	  src/backend/utils/mmgr/alignedalloc.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "utils/memdebug.h"
+#include "utils/memutils_memorychunk.h"
+
+/*
+ * AlignedAllocFree
+*		Frees allocated memory; memory is removed from its owning context.
+*/
+void
+AlignedAllocFree(void *pointer)
+{
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	void	   *unaligned;
+
+	Assert(!MemoryChunkIsExternal(chunk));
+
+	/* obtain the original (unaligned) allocated pointer */
+	unaligned = MemoryChunkGetBlock(chunk);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* Test for someone scribbling on unused space in chunk */
+	if (!sentinel_ok(pointer, chunk->requested_size))
+		elog(WARNING, "detected write past chunk end in %s %p",
+			 GetMemoryChunkContext(unaligned)->name, chunk);
+#endif
+
+	pfree(unaligned);
+}
+
+/*
+ * AlignedAllocRealloc
+ *		Change the allocated size of a chunk and return possibly a different
+ *		pointer to a memory address aligned to the same boundary as the
+ *		originally requested alignment.  The contents of 'pointer' will be
+ *		copied into the returned pointer up until 'size'.  Any additional
+ *		memory will be uninitialized.
+ */
+void *
+AlignedAllocRealloc(void *pointer, Size size)
+{
+	MemoryChunk *redirchunk = PointerGetMemoryChunk(pointer);
+	Size		alignto = MemoryChunkGetValue(redirchunk);
+	void	   *unaligned = MemoryChunkGetBlock(redirchunk);
+	MemoryContext ctx;
+	Size		old_size;
+	void	   *newptr;
+
+	/* sanity check this is a power of 2 value */
+	Assert((alignto & (alignto - 1)) == 0);
+
+	/*
+	 * Determine the size of the original allocation.  We can't determine this
+	 * exactly as GetMemoryChunkSpace() returns the total space used for the
+	 * allocation, which for contexts like aset includes rounding up to the
+	 * next power of 2.  However, this value is just used to memcpy() the old
+	 * data into the new allocation, so we only need to concern ourselves with
+	 * not reading beyond the end of the original allocation's memory.  The
+	 * drawback here is that we may copy more bytes than we need to, which
+	 * only amounts to wasted effort.  We can safely subtract the extra bytes
+	 * that we requested to allow us to align the pointer.  We must also
+	 * subtract the space for the unaligned pointer's MemoryChunk since
+	 * GetMemoryChunkSpace should have included that.  This does assume that
+	 * all context types use MemoryChunk as a chunk header.
+	 */
+	old_size = GetMemoryChunkSpace(unaligned) -
+		PallocAlignedExtraBytes(alignto) - sizeof(MemoryChunk);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* check that GetMemoryChunkSpace returned something realistic */
+	Assert(old_size >= redirchunk->requested_size);
+#endif
+
+	ctx = GetMemoryChunkContext(unaligned);
+	newptr = MemoryContextAllocAligned(ctx, size, alignto, 0);
+
+	/*
+	 * We may memcpy beyond the end of the original allocation request size,
+	 * so we must mark the entire allocation as defined.
+	 */
+	VALGRIND_MAKE_MEM_DEFINED(pointer, old_size);
+	memcpy(newptr, pointer, Min(size, old_size));
+	pfree(unaligned);
+
+	return newptr;
+}
+
+/*
+ * AlignedAllocGetChunkContext
+ *		Return the MemoryContext that 'pointer' belongs to.
+ */
+MemoryContext
+AlignedAllocGetChunkContext(void *pointer)
+{
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+
+	Assert(!MemoryChunkIsExternal(chunk));
+
+	return GetMemoryChunkContext(MemoryChunkGetBlock(chunk));
+}
+
+/*
+ * AlignedAllocGetChunkSpace
+ *		Given a currently-allocated chunk, determine the total space
+ *		it occupies (including all memory-allocation overhead).
+ */
+Size
+AlignedAllocGetChunkSpace(void *pointer)
+{
+	MemoryChunk *redirchunk = PointerGetMemoryChunk(pointer);
+	void	   *unaligned = MemoryChunkGetBlock(redirchunk);
+
+	return GetMemoryChunkSpace(unaligned);
+}
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 57bd6690ca0..814ada36cd6 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -30,6 +30,7 @@
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
 #include "utils/memutils_internal.h"
+#include "utils/memutils_memorychunk.h"
 
 
 static void BogusFree(void *pointer);
@@ -84,6 +85,21 @@ static const MemoryContextMethods mcxt_methods[] = {
 	[MCTX_SLAB_ID].check = SlabCheck,
 #endif
 
+	/* alignedalloc.c */
+	[MCTX_ALIGNED_REDIRECT_ID].alloc = NULL,	/* not required */
+	[MCTX_ALIGNED_REDIRECT_ID].free_p = AlignedAllocFree,
+	[MCTX_ALIGNED_REDIRECT_ID].realloc = AlignedAllocRealloc,
+	[MCTX_ALIGNED_REDIRECT_ID].reset = NULL,	/* not required */
+	[MCTX_ALIGNED_REDIRECT_ID].delete_context = NULL,	/* not required */
+	[MCTX_ALIGNED_REDIRECT_ID].get_chunk_context = AlignedAllocGetChunkContext,
+	[MCTX_ALIGNED_REDIRECT_ID].get_chunk_space = AlignedAllocGetChunkSpace,
+	[MCTX_ALIGNED_REDIRECT_ID].is_empty = NULL, /* not required */
+	[MCTX_ALIGNED_REDIRECT_ID].stats = NULL,	/* not required */
+#ifdef MEMORY_CONTEXT_CHECKING
+	[MCTX_ALIGNED_REDIRECT_ID].check = NULL,	/* not required */
+#endif
+
+
 	/*
 	 * Unused (as yet) IDs should have dummy entries here.  This allows us to
 	 * fail cleanly if a bogus pointer is passed to pfree or the like.  It
@@ -110,11 +126,6 @@ static const MemoryContextMethods mcxt_methods[] = {
 	[MCTX_UNUSED4_ID].realloc = BogusRealloc,
 	[MCTX_UNUSED4_ID].get_chunk_context = BogusGetChunkContext,
 	[MCTX_UNUSED4_ID].get_chunk_space = BogusGetChunkSpace,
-
-	[MCTX_UNUSED5_ID].free_p = BogusFree,
-	[MCTX_UNUSED5_ID].realloc = BogusRealloc,
-	[MCTX_UNUSED5_ID].get_chunk_context = BogusGetChunkContext,
-	[MCTX_UNUSED5_ID].get_chunk_space = BogusGetChunkSpace,
 };
 
 /*
@@ -1299,6 +1310,125 @@ palloc_extended(Size size, int flags)
 }
 
 /*
+ * MemoryContextAllocAligned
+ *		Allocate 'size' bytes of memory in 'context' aligned to 'alignto'
+ *		bytes.
+ *
+ * Currently, we align addresses by requesting additional bytes from the
+ * MemoryContext's standard allocator function and then aligning the returned
+ * address by the required alignment.  This means that the given MemoryContext
+ * must support providing us with a chunk of memory that's larger than 'size'.
+ * For allocators such as Slab, that's not going to work, as slab only allows
+ * chunks of the size that's specified when the context is created.
+ *
+ * 'alignto' must be a power of 2.
+ * 'flags' may be 0 or set the same as MemoryContextAllocExtended().
+ */
+void *
+MemoryContextAllocAligned(MemoryContext context,
+						  Size size, Size alignto, int flags)
+{
+	MemoryChunk *alignedchunk;
+	Size		alloc_size;
+	void	   *unaligned;
+	void	   *aligned;
+
+	/* wouldn't make much sense to waste that much space */
+	Assert(alignto < (128 * 1024 * 1024));
+
+	/* ensure alignto is a power of 2 */
+	Assert((alignto & (alignto - 1)) == 0);
+
+	/*
+	 * If the alignment requirements are less than what we already guarantee
+	 * then just use the standard allocation function.
+	 */
+	if (unlikely(alignto <= MAXIMUM_ALIGNOF))
+		return MemoryContextAllocExtended(context, size, flags);
+
+	/*
+	 * We implement aligned pointers by simply allocating enough memory for
+	 * the requested size plus the alignment and an additional "redirection"
+	 * MemoryChunk.  This additional MemoryChunk is required for operations
+	 * such as pfree when used on the pointer returned by this function.  We
+	 * use this redirection MemoryChunk in order to find the pointer to the
+	 * memory that was returned by the MemoryContextAllocExtended call below.
+	 * We do that by "borrowing" the block offset field and instead of using
+	 * that to find the offset into the owning block, we use it to find the
+	 * original allocated address.
+	 *
+	 * Here we must allocate enough extra memory so that we can still align
+	 * the pointer returned by MemoryContextAllocExtended and also have enough
+	 * space for the redirection MemoryChunk.  Since allocations will already
+	 * be at least aligned by MAXIMUM_ALIGNOF, we can subtract that amount
+	 * from the allocation size to save a little memory.
+	 */
+	alloc_size = size + PallocAlignedExtraBytes(alignto);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* ensure there's space for a sentinal byte */
+	alloc_size += 1;
+#endif
+
+	/* perform the actual allocation */
+	unaligned = MemoryContextAllocExtended(context, alloc_size, flags);
+
+	/* set the aligned pointer */
+	aligned = (void *) TYPEALIGN(alignto, (char *) unaligned +
+								 sizeof(MemoryChunk));
+
+	alignedchunk = PointerGetMemoryChunk(aligned);
+
+	/*
+	 * We set the redirect MemoryChunk so that the block offset calculation is
+	 * used to point back to the 'unaligned' allocated chunk.  This allows us
+	 * to use MemoryChunkGetBlock() to find the unaligned chunk when we need
+	 * to perform operations such as pfree() and repalloc().
+	 *
+	 * We store 'alignto' in the MemoryChunk's 'value' so that we know what
+	 * the alignment was set to should we ever be asked to realloc this
+	 * pointer.
+	 */
+	MemoryChunkSetHdrMask(alignedchunk, unaligned, alignto,
+						  MCTX_ALIGNED_REDIRECT_ID);
+
+	/* double check we produced a correctly aligned pointer */
+	Assert((void *) TYPEALIGN(alignto, aligned) == aligned);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	alignedchunk->requested_size = size;
+	/* set mark to catch clobber of "unused" space */
+	set_sentinel(aligned, size);
+#endif
+
+	/* Mark the bytes before the redirection header as noaccess */
+	VALGRIND_MAKE_MEM_NOACCESS(unaligned,
+							   (char *) alignedchunk - (char *) unaligned);
+	return aligned;
+}
+
+/*
+ * palloc_aligned
+ *		Allocate 'size' bytes returning a pointer that's aligned to the
+ *		'alignto' boundary.
+ *
+ * Currently, we align addresses by requesting additional bytes from the
+ * MemoryContext's standard allocator function and then aligning the returned
+ * address by the required alignment.  This means that the given MemoryContext
+ * must support providing us with a chunk of memory that's larger than 'size'.
+ * For allocators such as Slab, that's not going to work, as slab only allows
+ * chunks of the size that's specified when the context is created.
+ *
+ * 'alignto' must be a power of 2.
+ * 'flags' may be 0 or set the same as MemoryContextAllocExtended().
+ */
+void *
+palloc_aligned(Size size, Size alignto, int flags)
+{
+	return MemoryContextAllocAligned(CurrentMemoryContext, size, alignto, flags);
+}
+
+/*
  * pfree
  *		Release an allocated chunk.
  */
@@ -1306,11 +1436,16 @@ void
 pfree(void *pointer)
 {
 #ifdef USE_VALGRIND
+	MemoryContextMethodID method = GetMemoryChunkMethodID(pointer);
 	MemoryContext context = GetMemoryChunkContext(pointer);
 #endif
 
 	MCXT_METHOD(pointer, free_p) (pointer);
-	VALGRIND_MEMPOOL_FREE(context, pointer);
+
+#ifdef USE_VALGRIND
+	if (method != MCTX_ALIGNED_REDIRECT_ID)
+		VALGRIND_MEMPOOL_FREE(context, pointer);
+#endif
 }
 
 /*
@@ -1320,6 +1455,9 @@ pfree(void *pointer)
 void *
 repalloc(void *pointer, Size size)
 {
+#ifdef USE_VALGRIND
+	MemoryContextMethodID method = GetMemoryChunkMethodID(pointer);
+#endif
 #if defined(USE_ASSERT_CHECKING) || defined(USE_VALGRIND)
 	MemoryContext context = GetMemoryChunkContext(pointer);
 #endif
@@ -1346,7 +1484,10 @@ repalloc(void *pointer, Size size)
 						   size, cxt->name)));
 	}
 
-	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+#ifdef USE_VALGRIND
+	if (method != MCTX_ALIGNED_REDIRECT_ID)
+		VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+#endif
 
 	return ret;
 }
diff --git a/src/backend/utils/mmgr/meson.build b/src/backend/utils/mmgr/meson.build
index f0dfc91a7ba..cd52654cb1e 100644
--- a/src/backend/utils/mmgr/meson.build
+++ b/src/backend/utils/mmgr/meson.build
@@ -1,6 +1,7 @@
 # Copyright (c) 2022, PostgreSQL Global Development Group
 
 backend_sources += files(
+  'alignedalloc.c',
   'aset.c',
   'dsa.c',
   'freepage.c',