[PGPRO-7038] bgworker for aqo_cleanup

Author: Andrey Kazarinov

aqo.c

@@ -26,6 +26,8 @@
 #include "preprocessing.h"
 #include "learn_cache.h"
 #include "storage.h"
+#include "postmaster/bgworker.h"
+#include "catalog/objectaccess.h"
 
 
 PG_MODULE_MAGIC;
@@ -77,6 +79,8 @@ int			auto_tuning_infinite_loop = 8;
 int			aqo_k = 3;
 double		log_selectivity_lower_bound = -30;
 
+bool		cleanup_bgworker = false;
+
 /*
  * Currently we use it only to store query_text string which is initialized
  * after a query parsing and is used during the query planning.
@@ -112,6 +116,7 @@ get_parameterized_joinrel_size_hook_type	prev_get_parameterized_joinrel_size_hoo
 ExplainOnePlan_hook_type					prev_ExplainOnePlan_hook;
 ExplainOneNode_hook_type					prev_ExplainOneNode_hook;
 static shmem_request_hook_type				prev_shmem_request_hook = NULL;
+object_access_hook_type						prev_object_access_hook;
 
 /*****************************************************************************
  *
@@ -147,6 +152,93 @@ aqo_shmem_request(void)
 	RequestAddinShmemSpace(aqo_memsize());
 }
 
+/*
+ * Entry point for CleanupWorker's process.
+ */
+void
+aqo_bgworker_cleanup(void)
+{
+	int	fs_num;
+	int	fss_num;
+
+	cleanup_aqo_database(true, &fs_num, &fss_num);
+}
+
+/* 
+ * Object access hook
+ */
+void
+aqo_drop_access_hook(ObjectAccessType access,
+					 Oid classId,
+					 Oid objectId,
+					 int subId,
+					 void *arg)
+{
+	if (prev_object_access_hook)
+		(*prev_object_access_hook) (access, classId, objectId, subId, arg);
+
+	if (access == OAT_DROP && cleanup_bgworker)
+	{
+		MemoryContext old_ctx;
+		int status = BGWH_STOPPED;
+		pid_t pid;
+
+		old_ctx = MemoryContextSwitchTo(AQOTopMemCtx);
+		LWLockAcquire(&aqo_state->lock, LW_EXCLUSIVE);
+		if (aqo_state->bgw_handle != NULL)
+		{
+			status = GetBackgroundWorkerPid(aqo_state->bgw_handle, &pid);
+		}
+		LWLockRelease(&aqo_state->lock);
+		if (status != BGWH_STARTED)
+		{
+			aqo_bgworker_startup();
+		}
+		MemoryContextSwitchTo(old_ctx);
+	}
+}
+
+void
+aqo_bgworker_startup(void)
+{
+	BackgroundWorker		worker;
+	BackgroundWorkerHandle	*handle;
+	BgwHandleStatus			status;
+	pid_t					pid;
+
+	MemSet(&worker, 0, sizeof(worker));
+
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_ConsistentState;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	worker.bgw_main_arg = Int32GetDatum(0);
+	worker.bgw_extra[0] = 0;
+	memcpy(worker.bgw_function_name, "aqo_bgworker_cleanup", 21);
+	memcpy(worker.bgw_library_name, "aqo", 4);
+	memcpy(worker.bgw_name, "aqo cleanup", 12);
+
+	/* must set notify PID to wait for startup */
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		ereport(NOTICE,
+				(errcode(ERRCODE_CONFIGURATION_LIMIT_EXCEEDED),
+				 errmsg("could not register background process"),
+				 errhint("You might need to increase max_worker_processes.")));
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+	if (status != BGWH_STARTED)
+		ereport(NOTICE,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+
+	LWLockAcquire(&aqo_state->lock, LW_EXCLUSIVE);
+	aqo_state->bgw_handle = handle;
+	LWLockRelease(&aqo_state->lock);
+}
+
 void
 _PG_init(void)
 {
@@ -309,6 +401,18 @@ _PG_init(void)
 							NULL
 	);
 
+	DefineCustomBoolVariable("aqo.cleanup_bgworker",
+							 "Enable bgworker which responsible for doing cleanup after drop",
+							 NULL,
+							 &cleanup_bgworker,
+							 false,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL
+	);
+
 	prev_shmem_startup_hook						= shmem_startup_hook;
 	shmem_startup_hook							= aqo_init_shmem;
 	prev_planner_hook							= planner_hook;
@@ -349,6 +453,9 @@ _PG_init(void)
 	prev_shmem_request_hook = shmem_request_hook;
 	shmem_request_hook = aqo_shmem_request;
 
+	prev_object_access_hook						= object_access_hook;
+	object_access_hook							= aqo_drop_access_hook;
+
 	init_deactivated_queries_storage();
 
 	/*

aqo.h

@@ -142,6 +142,7 @@
 #include "utils/rel.h"
 #include "utils/fmgroids.h"
 #include "utils/snapmgr.h"
+#include "catalog/objectaccess.h"
 
 #include "machine_learning.h"
 //#include "storage.h"
@@ -299,5 +300,10 @@ extern void selectivity_cache_clear(void);
 
 extern bool IsQueryDisabled(void);
 
+extern void aqo_drop_access_hook(ObjectAccessType access, Oid classId, 
+								 Oid objectId, int subId, void *arg);
+extern void aqo_bgworker_startup(void);
+extern PGDLLEXPORT void aqo_bgworker_cleanup(void);
+
 extern List *cur_classes;
 #endif

aqo_shared.c

@@ -199,6 +199,7 @@ aqo_init_shmem(void)
 		aqo_state->stat_changed = false;
 		aqo_state->data_changed = false;
 		aqo_state->queries_changed = false;
+		aqo_state->bgw_handle = NULL;
 
 		LWLockInitialize(&aqo_state->lock, LWLockNewTrancheId());
 		LWLockInitialize(&aqo_state->stat_lock, LWLockNewTrancheId());

aqo_shared.h

@@ -7,6 +7,7 @@
 #include "storage/lwlock.h"
 #include "utils/dsa.h"
 #include "lib/dshash.h"
+#include "postmaster/bgworker.h"
 
 #define AQO_SHARED_MAGIC	0x053163
 
@@ -43,6 +44,8 @@ typedef struct AQOSharedState
 
 	LWLock		queries_lock;  /* lock for access to queries storage */
 	bool		queries_changed;
+
+	BackgroundWorkerHandle	*bgw_handle;
 } AQOSharedState;
 
 

expected/cleanup_bgworker.out

@@ -0,0 +1,131 @@
+CREATE TABLE aqo_test0(a int, b int, c int, d int);
+WITH RECURSIVE t(a, b, c, d)
+AS (
+   VALUES (0, 0, 0, 0)
+   UNION ALL
+   SELECT t.a + 1, t.b + 1, t.c + 1, t.d + 1 FROM t WHERE t.a < 2000
+) INSERT INTO aqo_test0 (SELECT * FROM t);
+CREATE INDEX aqo_test0_idx_a ON aqo_test0 (a);
+ANALYZE aqo_test0;
+CREATE EXTENSION aqo;
+SET aqo.join_threshold = 0;
+SET aqo.mode = 'learn';
+SET aqo.cleanup_bgworker = 'true';
+EXPLAIN SELECT t1.a, t2.b FROM aqo_test0 AS t1, aqo_test0 AS t2
+WHERE t1.a < 1 AND t1.b < 1 AND t2.c < 1 AND t2.d < 1;
+                                        QUERY PLAN                                        
+------------------------------------------------------------------------------------------
+ Nested Loop  (cost=0.28..49.32 rows=1 width=8)
+   ->  Index Scan using aqo_test0_idx_a on aqo_test0 t1  (cost=0.28..8.30 rows=1 width=4)
+         Index Cond: (a < 1)
+         Filter: (b < 1)
+   ->  Seq Scan on aqo_test0 t2  (cost=0.00..41.02 rows=1 width=4)
+         Filter: ((c < 1) AND (d < 1))
+(6 rows)
+
+-- Test 1 : delete 1 out of 1 entry
+SELECT count(*) FROM aqo_queries;
+ count 
+-------
+     2
+(1 row)
+
+DROP TABLE aqo_test0;
+CREATE TABLE IF NOT EXISTS aqo_test0(a int, b int, c int, d int);
+WITH RECURSIVE t(a, b, c, d)
+AS (
+   VALUES (0, 0, 0, 0)
+   UNION ALL
+   SELECT t.a + 1, t.b + 1, t.c + 1, t.d + 1 FROM t WHERE t.a < 2000
+) INSERT INTO aqo_test0 (SELECT * FROM t);
+CREATE INDEX aqo_test0_idx_a ON aqo_test0 (a);
+ANALYZE aqo_test0;
+CREATE TABLE aqo_test1(a int, b int);
+WITH RECURSIVE t(a, b)
+AS (
+   VALUES (1, 2)
+   UNION ALL
+   SELECT t.a + 1, t.b + 1 FROM t WHERE t.a < 20
+) INSERT INTO aqo_test1 (SELECT * FROM t);
+CREATE INDEX aqo_test1_idx_a ON aqo_test1 (a);
+ANALYZE aqo_test1;
+SELECT count(*) FROM aqo_queries;
+ count 
+-------
+     1
+(1 row)
+
+EXPLAIN SELECT * FROM aqo_test0
+WHERE a < 3 AND b < 3 AND c < 3 AND d < 3;
+                                    QUERY PLAN                                    
+----------------------------------------------------------------------------------
+ Index Scan using aqo_test0_idx_a on aqo_test0  (cost=0.28..8.35 rows=1 width=16)
+   Index Cond: (a < 3)
+   Filter: ((b < 3) AND (c < 3) AND (d < 3))
+(3 rows)
+
+EXPLAIN SELECT * FROM aqo_test1
+WHERE a < 4 AND b < 4;
+                       QUERY PLAN                        
+---------------------------------------------------------
+ Seq Scan on aqo_test1  (cost=0.00..1.30 rows=1 width=8)
+   Filter: ((a < 4) AND (b < 4))
+(2 rows)
+
+-- Test2: Must delete 1 out of 2 entries
+SELECT count(*) FROM aqo_queries;
+ count 
+-------
+     3
+(1 row)
+
+DROP TABLE aqo_test1;
+CREATE TABLE IF NOT EXISTS aqo_test1(a int, b int);
+WITH RECURSIVE t(a, b)
+AS (
+   VALUES (1, 2)
+   UNION ALL
+   SELECT t.a + 1, t.b + 1 FROM t WHERE t.a < 20
+) INSERT INTO aqo_test1 (SELECT * FROM t);
+CREATE INDEX aqo_test1_idx_a ON aqo_test1 (a);
+ANALYZE aqo_test1;
+SELECT count(*) FROM aqo_queries;
+ count 
+-------
+     1
+(1 row)
+
+EXPLAIN SELECT * FROM aqo_test0
+WHERE a < 3 AND b < 3 AND c < 3 AND d < 3;
+                                    QUERY PLAN                                    
+----------------------------------------------------------------------------------
+ Index Scan using aqo_test0_idx_a on aqo_test0  (cost=0.28..8.35 rows=1 width=16)
+   Index Cond: (a < 3)
+   Filter: ((b < 3) AND (c < 3) AND (d < 3))
+(3 rows)
+
+EXPLAIN SELECT * FROM aqo_test1
+WHERE a < 4 AND b < 4;
+                       QUERY PLAN                        
+---------------------------------------------------------
+ Seq Scan on aqo_test1  (cost=0.00..1.30 rows=1 width=8)
+   Filter: ((a < 4) AND (b < 4))
+(2 rows)
+
+-- Test3: delete 2 out of 2 entries
+SELECT count(*) FROM aqo_queries;
+ count 
+-------
+     3
+(1 row)
+
+DROP TABLE aqo_test0;
+DROP TABLE aqo_test1;
+SELECT count(*) FROM aqo_queries;
+ count 
+-------
+     1
+(1 row)
+
+SET aqo.cleanup_bgworker = 'false';
+DROP EXTENSION aqo;

expected/relocatable.out

@@ -114,9 +114,9 @@ ORDER BY (learn_aqo, use_aqo, auto_tuning);
  t         | t       | f
 (3 rows)
 
-RESET search_path;
 DROP TABLE test CASCADE;
 DROP SCHEMA IF EXISTS test CASCADE;
 NOTICE:  drop cascades to extension aqo
 DROP EXTENSION IF EXISTS aqo CASCADE;
 NOTICE:  extension "aqo" does not exist, skipping
+RESET search_path;

regress_schedule

@@ -20,3 +20,4 @@ test: top_queries
 test: relocatable
 test: look_a_like
 test: feature_subspace
+test: cleanup_bgworker