Add basic code for support of DSM cache.

danolivo · danolivo · commit ae557bd514ff · 2022-05-16T14:26:34.000+05:00
diff --git a/Makefile b/Makefile
@@ -6,7 +6,7 @@ PGFILEDESC = "AQO - Adaptive Query Optimization"
 MODULE_big = aqo
 OBJS = aqo.o auto_tuning.o cardinality_estimation.o cardinality_hooks.o \
 hash.o machine_learning.o path_utils.o postprocessing.o preprocessing.o \
-selectivity_cache.o storage.o utils.o learn_cache.o $(WIN32RES)
+selectivity_cache.o storage.o utils.o learn_cache.o aqo_shared.o $(WIN32RES)
 
 TAP_TESTS = 1
 
diff --git a/aqo.c b/aqo.c
@@ -18,6 +18,7 @@
 #include "utils/selfuncs.h"
 
 #include "aqo.h"
+#include "aqo_shared.h"
 #include "cardinality_hooks.h"
 #include "path_utils.h"
 #include "preprocessing.h"
@@ -125,7 +126,7 @@ _PG_init(void)
 {
 	/*
 	 * In order to create our shared memory area, we have to be loaded via
-	 * shared_preload_libraries.  If not, report an ERROR.
+	 * shared_preload_libraries. If not, report an ERROR.
 	 */
 	if (!process_shared_preload_libraries_in_progress)
 		ereport(ERROR,
@@ -198,6 +199,8 @@ _PG_init(void)
 							 NULL
 	);
 
+	prev_shmem_startup_hook						= shmem_startup_hook;
+	shmem_startup_hook							= aqo_init_shmem;
 	prev_planner_hook							= planner_hook;
 	planner_hook								= aqo_planner;
 	prev_ExecutorStart_hook						= ExecutorStart_hook;
@@ -239,6 +242,10 @@ _PG_init(void)
 											 ALLOCSET_DEFAULT_SIZES);
 	RegisterResourceReleaseCallback(aqo_free_callback, NULL);
 	RegisterAQOPlanNodeMethods();
+
+	MarkGUCPrefixReserved("aqo");
+	RequestAddinShmemSpace(MAXALIGN(sizeof(AQOSharedState)));
+
 	lc_init();
 }
 
diff --git a/aqo_shared.c b/aqo_shared.c
@@ -0,0 +1,61 @@
+/*
+ *
+ */
+
+#include "postgres.h"
+
+#include "storage/shmem.h"
+
+#include "aqo_shared.h"
+
+shmem_startup_hook_type prev_shmem_startup_hook = NULL;
+static AQOSharedState *aqo_state = NULL;
+unsigned long temp_storage_size = 1024 * 1024; /* Storage size, in bytes */
+void *temp_storage = NULL;
+
+static void
+attach_dsm_segment(void)
+{
+	dsm_segment *seg;
+
+	LWLockAcquire(&aqo_state->lock, LW_EXCLUSIVE);
+
+	if (aqo_state->dsm_handler != DSM_HANDLE_INVALID)
+	{
+		seg = dsm_attach(aqo_state->dsm_handler);
+	}
+	else
+	{
+		seg = dsm_create(temp_storage_size, 0);
+		aqo_state->dsm_handler = dsm_segment_handle(seg);
+	}
+
+	temp_storage = dsm_segment_address(seg);
+	LWLockRelease(&aqo_state->lock);
+}
+
+static void
+aqo_detach_shmem(int code, Datum arg)
+{
+	dsm_handle handler = *(dsm_handle *) arg;
+	dsm_detach(dsm_find_mapping(handler));
+}
+
+void
+aqo_init_shmem(void)
+{
+	bool		found;
+
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	aqo_state = ShmemInitStruct("aqo", sizeof(AQOSharedState), &found);
+	if (!found)
+	{
+		/* First time through ... */
+		LWLockInitialize(&aqo_state->lock, LWLockNewTrancheId());
+		aqo_state->dsm_handler = DSM_HANDLE_INVALID;
+	}
+	LWLockRelease(AddinShmemInitLock);
+
+	LWLockRegisterTranche(aqo_state->lock.tranche, "aqo");
+	on_shmem_exit(aqo_detach_shmem, (Datum) &aqo_state->dsm_handler);
+}
diff --git a/aqo_shared.h b/aqo_shared.h
@@ -0,0 +1,22 @@
+#ifndef AQO_SHARED_H
+#define AQO_SHARED_H
+
+
+#include "storage/dsm.h"
+#include "storage/ipc.h"
+#include "storage/lwlock.h"
+
+
+typedef struct AQOSharedState
+{
+	LWLock		lock;			/* mutual exclusion */
+	dsm_handle	dsm_handler;
+} AQOSharedState;
+
+
+extern shmem_startup_hook_type prev_shmem_startup_hook;
+
+
+extern void aqo_init_shmem(void);
+
+#endif /* AQO_SHARED_H */
