support implementing custom storage manager in extension

We add a new smgr and need to make the following changes:

Firstly we should to define our own f_smgr structure in extension.
struct f_smgr custom_smgr = {
	.smgr_name = "custom smgr",
	.smgr_name = = md_init,
	...
}

Then register our own structure in shared_preload_libraries phase.
smgr_register(&custom_smgr, CUSTOM_SMGR_ID);
CUSTOM_SMGR_ID should be unique and cannot be changed, and should
not conflict with other extensions.

Finally we need to correctly return the relation's smgr id in smgr_get_impl
the extension need to implement smgr_get_impl_hook correctly

Author: gongxun0928

src/backend/access/rmgrdesc/smgrdesc.c

@@ -27,7 +27,11 @@ smgr_desc(StringInfo buf, XLogReaderState *record)
 	{
 		xl_smgr_create *xlrec = (xl_smgr_create *) rec;
 		char	   *path = relpathperm(xlrec->rnode, xlrec->forkNum);
-		appendStringInfo(buf, "%s; smgr: %s", path, xlrec->impl == SMGR_MD ? "heap" : "ao");
+#ifndef FRONTEND
+		appendStringInfo(buf, "%s; smgr: %s", path, smgr_get_name(xlrec->impl));
+#else
+		appendStringInfo(buf, "%s; smgr: %s", path, xlrec->impl == SMGR_MD ? "heap" : (xlrec->impl == SMGR_AO ? "ao" : "unknown"));
+#endif
 		pfree(path);
 	}
 	else if (info == XLOG_SMGR_TRUNCATE)

src/backend/catalog/storage.c

@@ -213,8 +213,7 @@ RelationDropStorage(Relation rel)
 	pending->relnode.isTempRelation = rel->rd_backend == TempRelBackendId;
 	pending->atCommit = true;	/* delete if commit */
 	pending->nestLevel = GetCurrentTransactionNestLevel();
-	pending->relnode.smgr_which =
-		RelationIsAppendOptimized(rel) ? SMGR_AO : SMGR_MD;
+	pending->relnode.smgr_which = smgr_get_impl(rel);
 	pending->action = &storage_pending_rel_deletes_action;
 	RegisterPendingDelete(pending);
 

src/backend/commands/tablecmds.c

@@ -16559,7 +16559,7 @@ index_copy_data(Relation rel, RelFileNode newrnode)
 {
 	SMgrRelation dstrel;
 
-	SMgrImpl smgr_which = RelationIsAppendOptimized(rel) ? SMGR_AO : SMGR_MD;
+	SMgrImpl smgr_which = smgr_get_impl(rel);
 
 	dstrel = smgropen(newrnode, rel->rd_backend, smgr_which, rel);
 

src/backend/storage/smgr/smgr.c

@@ -50,14 +50,17 @@ file_extend_hook_type file_extend_hook = NULL;
 file_truncate_hook_type file_truncate_hook = NULL;
 file_unlink_hook_type file_unlink_hook = NULL;
 
+smgr_get_impl_hook_type smgr_get_impl_hook = NULL;
+
 /* Hook for plugins to get control in smgr */
 smgr_init_hook_type smgr_init_hook = NULL;
 smgr_hook_type smgr_hook = NULL;
 smgr_shutdown_hook_type smgr_shutdown_hook = NULL;
-
-static const f_smgr smgrsw[] = {
+#define SMGR_MAX_ID UINT8_MAX
+static f_smgr smgrsw[SMGR_MAX_ID + 1] = {
 	/* magnetic disk */
 	{
+		.smgr_name = "heap",
 		.smgr_init = mdinit,
 		.smgr_shutdown = NULL,
 		.smgr_open = mdopen,
@@ -82,6 +85,7 @@ static const f_smgr smgrsw[] = {
 	 * Append-optimized relation files currently fall in this category.
 	 */
 	{
+		.smgr_name = "ao",
 		.smgr_init = mdinit,
 		.smgr_shutdown = NULL,
 		.smgr_open = mdopen,
@@ -115,8 +119,6 @@ static const f_smgr_ao smgrswao[] = {
 };
 
 
-static const int NSmgr = lengthof(smgrsw);
-
 /*
  * Each backend has a hashtable that stores all extant SMgrRelation objects.
  * In addition, "unowned" SMgrRelation objects are chained together in a list.
@@ -128,6 +130,79 @@ static dlist_head unowned_relns;
 /* local function prototypes */
 static void smgrshutdown(int code, Datum arg);
 
+void smgr_register(const f_smgr *smgr, SMgrImpl smgr_impl)
+{
+	if (!process_shared_preload_libraries_in_progress)
+	{
+		ereport(ERROR, (errmsg("smgr_register not in shared_preload_libraries")));
+	}
+	
+	if (smgr_impl > SMGR_MAX_ID)
+	{
+		elog(ERROR, "smgr_impl is out of range");
+	}
+
+	if (smgr->smgr_name == NULL)
+	{
+		elog(ERROR, "smgr_name is not set");
+	}
+
+	// check if the smgr_impl is already registered, avoid conflict with existing smgr_impl
+	if (smgrsw[smgr_impl].smgr_name != NULL)
+	{
+		elog(ERROR, "smgr_impl is already registered");
+	}
+
+	smgrsw[smgr_impl] = *smgr;
+}
+
+const f_smgr *smgr_get(SMgrImpl smgr_impl)
+{
+	if (smgr_impl > SMGR_MAX_ID)
+	{
+		elog(ERROR, "smgr_impl is out of range");
+	}
+
+	if (smgrsw[smgr_impl].smgr_name == NULL)
+	{
+		elog(ERROR, "smgr_impl is not registered");
+	}
+
+	return &smgrsw[smgr_impl];
+}
+
+/*
+ * smgr_get_impl() is used to get the smgr id of the relation.
+ *
+ * FIXME: For PAX_AM_OID, Cloudberrydb reserves this value for ORCA, a
+ * predefined value is used here to reserve the smgr id for PAX_AM_OID.
+ * should we add a hook to get the smgr id of the relation?
+ */
+SMgrImpl smgr_get_impl(const Relation rel)
+{
+	SMgrImpl smgr_impl = SMGR_MD;
+
+	if (RelationIsAppendOptimized(rel))
+	{
+		smgr_impl = SMGR_AO;
+	}
+	else if (RelationIsPax(rel))
+	{
+		smgr_impl = SMGR_PAX;
+	}
+	else
+	{
+		if (smgr_get_impl_hook)
+		{
+			(*smgr_get_impl_hook)(rel, &smgr_impl);
+			Assert(smgr_impl >= SMGR_MD && smgr_impl <= SMGR_MAX_ID);
+			Assert(smgrsw[smgr_impl].smgr_name != NULL);
+		}
+	}
+
+
+	return smgr_impl;
+}
 
 /*
  *	smgrinit(), smgrshutdown() -- Initialize or shut down storage
@@ -142,7 +217,7 @@ smgrinit(void)
 {
 	int			i;
 
-	for (i = 0; i < NSmgr; i++)
+	for (i = 0; i <= SMGR_MAX_ID; i++)
 	{
 		if (smgrsw[i].smgr_init)
 			smgrsw[i].smgr_init();
@@ -163,7 +238,7 @@ smgrshutdown(int code, Datum arg)
 {
 	int			i;
 
-	for (i = 0; i < NSmgr; i++)
+	for (i = 0; i <= SMGR_MAX_ID; i++)
 	{
 		if (smgrsw[i].smgr_shutdown)
 			smgrsw[i].smgr_shutdown();
@@ -777,3 +852,10 @@ AtEOXact_SMgr(void)
 		smgrclose(rel);
 	}
 }
+
+const char *smgr_get_name(SMgrImpl impl)
+{
+	if (impl > SMGR_MAX_ID)
+		return "invalid";
+	return smgrsw[impl].smgr_name ? smgrsw[impl].smgr_name : "unknown";
+}

src/include/storage/smgr.h

@@ -24,11 +24,35 @@
 #include "storage/fd.h"
 #include "utils/relcache.h"
 
+/*
+ * Extension can register a new storage manager(smgr) by implementing
+ * its own struct f_smgr.
+ * Each storage manager has a unique implementation id, which is used
+ * to identify the smgr.
+ *
+ * The smgr_register function is used to register the smgr. it will
+ * check if the smgr is already registered, if so, return SMGR_INVALID.
+ * otherwise, register the smgr and return the implementation id.
+ *
+ * The smgr id of each relation is fixed and cannot be changed. because the
+ * type of smgr is logged in the commit log and wal log. If the value of
+ * smgr is changed, data corruption may occur.
+ *
+ * How to ensure that the smgr of each extension does not conflict? One way
+ * is to predefine the smgr id used by the extension in Cloudberry; the other
+ * way is to refer to the practice of custom rmgr and provide a wiki page
+ * [https://wiki.postgresql.org/wiki/CustomWALResourceManagers] to record
+ * the use of smgr to avoid conflicts.
+ *
+ * FIXME: For PAX_AM_OID, Cloudberrydb reserves this value for ORCA, a
+ * predefined value is used here to reserve the smgr id for PAX_AM_OID.
+ */
 typedef enum SMgrImplementation
 {
 	SMGR_INVALID = -1,
 	SMGR_MD = 0,
 	SMGR_AO = 1,
+	SMGR_PAX = 2,
 } SMgrImpl;
 
 struct f_smgr;
@@ -105,6 +129,7 @@ typedef SMgrRelationData *SMgrRelation;
  */
 typedef struct f_smgr
 {
+	const char 	*smgr_name;
 	void		(*smgr_init) (void);	/* may be NULL */
 	void		(*smgr_shutdown) (void);	/* may be NULL */
 	void		(*smgr_open) (SMgrRelation reln);
@@ -148,9 +173,16 @@ typedef void (*smgr_shutdown_hook_type) (void);
 extern PGDLLIMPORT smgr_init_hook_type smgr_init_hook;
 extern PGDLLIMPORT smgr_hook_type smgr_hook;
 extern PGDLLIMPORT smgr_shutdown_hook_type smgr_shutdown_hook;
-
 extern bool smgr_is_heap_relation(SMgrRelation reln);
 
+// must be registered in the shared_preload_libraries phase in extension
+// we should check whether smgr and smgr_impl is valid.
+extern void smgr_register(const f_smgr *smgr, SMgrImpl smgr_impl);
+
+extern const f_smgr *smgr_get(SMgrImpl smgr_impl);
+
+extern SMgrImpl smgr_get_impl(const Relation rel);
+
 extern void smgrinit(void);
 extern SMgrRelation smgropen(RelFileNode rnode, BackendId backend,
                              SMgrImpl smgr_which, Relation rel);
@@ -183,6 +215,8 @@ extern void AtEOXact_SMgr(void);
 
 extern const struct f_smgr_ao * smgrAOGetDefault(void);
 
+extern const char* smgr_get_name(SMgrImpl impl);
+
 
 /*
  * Hook for plugins to collect statistics from storage functions
@@ -201,4 +235,14 @@ extern PGDLLIMPORT file_truncate_hook_type file_truncate_hook;
 typedef void (*file_unlink_hook_type)(RelFileNodeBackend rnode);
 extern PGDLLIMPORT file_unlink_hook_type file_unlink_hook;
 
+/*
+ * This hook is used to get the smgr implementation id of the relation for extension.
+ * If the hook is not set, the default smgr implementation id is SMGR_MD.
+ * If the extension register a custom smgr to manage its own relation, it needs
+ * to implement this hook and then set the smgr_impl to the correct value for
+ * the relation which is managed by the extension, otherwise should ignore it.
+ */
+typedef void (*smgr_get_impl_hook_type)(const Relation rel, SMgrImpl* smgr_impl);
+extern PGDLLIMPORT smgr_get_impl_hook_type smgr_get_impl_hook;
+
 #endif							/* SMGR_H */

src/include/utils/rel.h

@@ -532,6 +532,10 @@ typedef struct ViewOptions
  * can't distinguish the PAX and renamed heap(heap_psql) in test `psql`.
  */
 #define PAX_AM_OID 7047
+
+#define RelationIsPax(relation) \
+	((relation)->rd_rel->relam == PAX_AM_OID)
+
 /*
  * CAUTION: this macro is a violation of the absraction that table AM and
  * index AM interfaces provide.  Use of this macro is discouraged.  If
@@ -548,7 +552,7 @@ typedef struct ViewOptions
  */
 #define RelationIsNonblockRelation(relation) \
 	(RelationIsAppendOptimized(relation) || \
-	 (relation)->rd_rel->relam == PAX_AM_OID)
+	 RelationIsPax(relation))
 
 /*
  * RelationIsBitmapIndex