From 0fa60d8241e25519e7b1d6f663d1b369b4c84412 Mon Sep 17 00:00:00 2001
From: Andreas Karlsson <andreas.karlsson@percona.com>
Date: Thu, 9 Oct 2025 19:45:54 +0200
Subject: [PATCH v6 6/6] Make AIO handle callbacks extensible

With the extensible SMGR patch, extensions can register their own
storage manager logic. However, this is not that helpful with the startreadv
as the AIO API provides no extensibility mechanisms:

* There's no way to extend the callback chain defined by md.c, as it is
  not visible to extensions
* And there's also no way to define another completely custom chain
  instead of it

This patch aims to improve this by making the md callback chain visible
to extensions, and also increasing the maximum number of callbacks so
that a few extensions can safely register their additions to it.
---
 src/backend/storage/aio/aio_callback.c | 20 +++++++++++++++++---
 src/backend/storage/smgr/md.c          |  5 ++---
 src/include/storage/aio.h              |  3 ++-
 src/include/storage/md.h               |  5 ++++-
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/backend/storage/aio/aio_callback.c b/src/backend/storage/aio/aio_callback.c
index 206b81f5028..49e3d512eed 100644
--- a/src/backend/storage/aio/aio_callback.c
+++ b/src/backend/storage/aio/aio_callback.c
@@ -27,8 +27,8 @@ static const PgAioHandleCallbacks aio_invalid_cb = {0};
 
 typedef struct PgAioHandleCallbacksEntry
 {
-	const PgAioHandleCallbacks *const cb;
-	const char *const name;
+	const PgAioHandleCallbacks *cb;
+	const char *name;
 } PgAioHandleCallbacksEntry;
 
 /*
@@ -36,7 +36,7 @@ typedef struct PgAioHandleCallbacksEntry
  * handle.  See PgAioHandleCallbackID's definition for an explanation for why
  * callbacks are not identified by a pointer.
  */
-static const PgAioHandleCallbacksEntry aio_handle_cbs[] = {
+static PgAioHandleCallbacksEntry aio_handle_cbs[PGAIO_HCB_MAX] = {
 #define CALLBACK_ENTRY(id, callback)  [id] = {.cb = &callback, .name = #callback}
 	CALLBACK_ENTRY(PGAIO_HCB_INVALID, aio_invalid_cb),
 
@@ -48,7 +48,21 @@ static const PgAioHandleCallbacksEntry aio_handle_cbs[] = {
 #undef CALLBACK_ENTRY
 };
 
+static PgAioHandleCallbackID aio_handle_max_cb_id = PGAIO_HCB_LOCAL_BUFFER_READV;
 
+PgAioHandleCallbackID
+pgaio_io_register_callback_entry(const PgAioHandleCallbacks *callback, const char *name)
+{
+	PgAioHandleCallbackID cb_id = ++aio_handle_max_cb_id;
+
+	if (cb_id > PGAIO_HCB_MAX)
+		elog(FATAL, "There can be at most %d AIO callback entries", PGAIO_HCB_MAX);
+
+	aio_handle_cbs[cb_id].cb = callback;
+	aio_handle_cbs[cb_id].name = name;
+
+	return cb_id;
+}
 
 /* --------------------------------------------------------------------------------
  * Public callback related functions operating on IO Handles
diff --git a/src/backend/storage/smgr/md.c b/src/backend/storage/smgr/md.c
index bea0b8229f5..b5960ba587a 100644
--- a/src/backend/storage/smgr/md.c
+++ b/src/backend/storage/smgr/md.c
@@ -241,10 +241,9 @@ static MdfdVec *_mdfd_getseg(MdSMgrRelation reln, ForkNumber forknum,
 static BlockNumber _mdnblocks(MdSMgrRelation reln, ForkNumber forknum,
 							  MdfdVec *seg);
 
-static PgAioResult md_readv_complete(PgAioHandle *ioh, PgAioResult prior_result, uint8 cb_data);
 static void md_readv_report(PgAioResult result, const PgAioTargetData *td, int elevel);
 
-const PgAioHandleCallbacks aio_md_readv_cb = {
+PgAioHandleCallbacks aio_md_readv_cb = {
 	.complete_shared = md_readv_complete,
 	.report = md_readv_report,
 };
@@ -2074,7 +2073,7 @@ mdfiletagmatches(const FileTag *ftag, const FileTag *candidate)
 /*
  * AIO completion callback for mdstartreadv().
  */
-static PgAioResult
+PgAioResult
 md_readv_complete(PgAioHandle *ioh, PgAioResult prior_result, uint8 cb_data)
 {
 	PgAioTargetData *td = pgaio_io_get_target_data(ioh);
diff --git a/src/include/storage/aio.h b/src/include/storage/aio.h
index ec543b78409..a0162b69be1 100644
--- a/src/include/storage/aio.h
+++ b/src/include/storage/aio.h
@@ -200,7 +200,7 @@ typedef enum PgAioHandleCallbackID
 	PGAIO_HCB_LOCAL_BUFFER_READV,
 } PgAioHandleCallbackID;
 
-#define PGAIO_HCB_MAX	PGAIO_HCB_LOCAL_BUFFER_READV
+#define PGAIO_HCB_MAX	15
 StaticAssertDecl(PGAIO_HCB_MAX < (1 << PGAIO_RESULT_ID_BITS),
 				 "PGAIO_HCB_MAX is too big for PGAIO_RESULT_ID_BITS");
 
@@ -308,6 +308,7 @@ extern PgAioTargetData *pgaio_io_get_target_data(PgAioHandle *ioh);
 extern char *pgaio_io_get_target_description(PgAioHandle *ioh);
 
 /* functions in aio_callback.c */
+extern PgAioHandleCallbackID pgaio_io_register_callback_entry(const PgAioHandleCallbacks *callback, const char *name);
 extern void pgaio_io_register_callbacks(PgAioHandle *ioh, PgAioHandleCallbackID cb_id,
 										uint8 cb_data);
 extern void pgaio_io_set_handle_data_64(PgAioHandle *ioh, uint64 *data, uint8 len);
diff --git a/src/include/storage/md.h b/src/include/storage/md.h
index 5a044ad9aee..506c5478309 100644
--- a/src/include/storage/md.h
+++ b/src/include/storage/md.h
@@ -20,7 +20,7 @@
 #include "storage/smgr.h"
 #include "storage/sync.h"
 
-extern PGDLLIMPORT const PgAioHandleCallbacks aio_md_readv_cb;
+extern PGDLLIMPORT PgAioHandleCallbacks aio_md_readv_cb;
 
 /* registration function for md storage manager */
 extern void mdsmgr_register(void);
@@ -34,4 +34,7 @@ extern int	mdsyncfiletag(const FileTag *ftag, char *path);
 extern int	mdunlinkfiletag(const FileTag *ftag, char *path);
 extern bool mdfiletagmatches(const FileTag *ftag, const FileTag *candidate);
 
+extern PgAioResult
+md_readv_complete(PgAioHandle *ioh, PgAioResult prior_result, uint8 cb_data);
+
 #endif							/* MD_H */
-- 
2.43.0

