Hide the AQO Statement Timeout feature under a GUC.

Use aqo.learn_statement_timeout to enable this feature. On more function here
is to do cleanup on this cache and memory context.

Author: a.pervushina

aqo.c

@@ -201,6 +201,19 @@ _PG_init(void)
 							 NULL
 	);
 
+	DefineCustomBoolVariable(
+							 "aqo.learn_statement_timeout",
+							 "Learn on a plan interrupted by statement timeout.",
+							 "ML data stored in a backend cache, so it works only locally.",
+							 &aqo_learn_statement_timeout,
+							 false,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 lc_assign_hook,
+							 NULL
+	);
+
 	prev_planner_hook							= planner_hook;
 	planner_hook								= aqo_planner;
 	prev_ExecutorStart_hook						= ExecutorStart_hook;

learn_cache.c

@@ -13,6 +13,7 @@
  */
 
 #include "postgres.h"
+#include "miscadmin.h"
 
 #include "aqo.h"
 #include "learn_cache.h"
@@ -39,6 +40,8 @@ typedef struct
 static HTAB *fss_htab = NULL;
 MemoryContext LearnCacheMemoryContext = NULL;
 
+bool aqo_learn_statement_timeout = false;
+
 void
 lc_init(void)
 {
@@ -66,7 +69,7 @@ lc_update_fss(uint64 fs, int fss, int nrows, int ncols,
 	int				i;
 	MemoryContext	memctx = MemoryContextSwitchTo(LearnCacheMemoryContext);
 
-	Assert(fss_htab);
+	Assert(fss_htab && aqo_learn_statement_timeout);
 
 	entry = (htab_entry *) hash_search(fss_htab, &key, HASH_ENTER, &found);
 	if (found)
@@ -99,6 +102,9 @@ lc_has_fss(uint64 fs, int fss)
 	htab_key	key = {fs, fss};
 	bool		found;
 
+	if (!aqo_learn_statement_timeout)
+		return false;
+
 	Assert(fss_htab);
 
 	(void) hash_search(fss_htab, &key, HASH_FIND, &found);
@@ -118,7 +124,7 @@ lc_load_fss(uint64 fs, int fss, int ncols, double **matrix,
 	bool		found;
 	int			i;
 
-	Assert(fss_htab);
+	Assert(fss_htab && aqo_learn_statement_timeout);
 
 	entry = (htab_entry *) hash_search(fss_htab, &key, HASH_FIND, &found);
 	if (!found)
@@ -150,6 +156,9 @@ lc_remove_fss(uint64 fs, int fss)
 	bool		found;
 	int			i;
 
+	if (!aqo_learn_statement_timeout)
+		return;
+
 	Assert(fss_htab);
 
 	entry = (htab_entry *) hash_search(fss_htab, &key, HASH_FIND, &found);
@@ -161,3 +170,31 @@ lc_remove_fss(uint64 fs, int fss)
 	pfree(entry->targets);
 	hash_search(fss_htab, &key, HASH_REMOVE, NULL);
 }
+
+/*
+ * Main purpose of this hook is to cleanup a backend cache in some way to prevent
+ * memory leaks - in large queries we could have many unused fss nodes.
+ */
+void
+lc_assign_hook(bool newval, void *extra)
+{
+	HASH_SEQ_STATUS		status;
+	htab_entry		   *entry;
+
+	if (!fss_htab || !IsUnderPostmaster)
+		return;
+
+	/* Remove all entries, reset memory context. */
+
+	elog(DEBUG5, "[AQO] Cleanup local cache of ML data.");
+
+	/* Remove all frozen plans from a plancache. */
+	hash_seq_init(&status, fss_htab);
+	while ((entry = (htab_entry *) hash_seq_search(&status)) != NULL)
+	{
+		if (!hash_search(fss_htab, (void *) &entry->key, HASH_REMOVE, NULL))
+			elog(ERROR, "[AQO] The local ML cache is corrupted.");
+	}
+
+	MemoryContextReset(LearnCacheMemoryContext);
+}

learn_cache.h

@@ -3,6 +3,8 @@
 
 #include "nodes/pg_list.h"
 
+extern bool aqo_learn_statement_timeout;
+
 extern void lc_init(void);
 extern bool lc_update_fss(uint64 fhash, int fsshash, int nrows, int ncols,
 							  double **matrix, double *targets, List *relids);
@@ -11,5 +13,6 @@ extern bool lc_load_fss(uint64 fhash, int fsshash, int ncols,
 							double **matrix, double *targets, int *nrows,
 							List **relids);
 extern void lc_remove_fss(uint64 fhash, int fss_hash);
+extern void lc_assign_hook(bool newval, void *extra);
 
 #endif /* LEARN_CACHE_H */

postprocessing.c

@@ -684,7 +684,7 @@ set_timeout_if_need(QueryDesc *queryDesc)
 {
 	TimestampTz	fin_time;
 
-	if (!get_timeout_active(STATEMENT_TIMEOUT))
+	if (!get_timeout_active(STATEMENT_TIMEOUT) || !aqo_learn_statement_timeout)
 		return false;
 
 	if (!ExtractFromQueryEnv(queryDesc))
@@ -696,7 +696,7 @@ set_timeout_if_need(QueryDesc *queryDesc)
 
 	/*
 	 * Statement timeout exists. AQO should create user timeout right before the
-	 * statement timeout.
+	 * timeout.
 	 */
 
 	if (timeoutCtl.id < USER_TIMEOUT)

storage.c

@@ -369,10 +369,12 @@ load_fss_ext(uint64 fs, int fss,
 			 int ncols, double **matrix, double *targets, int *rows,
 			 List **relids, bool isSafe)
 {
-	if (isSafe && !lc_has_fss(fs, fss))
+	if (isSafe && (!aqo_learn_statement_timeout || !lc_has_fss(fs, fss)))
 		return load_fss(fs, fss, ncols, matrix, targets, rows, relids);
 	else
 	{
+		Assert(aqo_learn_statement_timeout);
+
 		if (matrix == NULL && targets == NULL && rows == NULL)
 			return true;