Be more careful with locks of relations and syscaches in get_list_of_relids() routine

danolivo · Alena Rybakina · commit 24da0292c483 · 2022-09-26T14:42:37.000+03:00
Switch on feature 'search on neighbour feature spaces' by a GUC (disabled by default).

Some mistakes fixed.
diff --git a/aqo.c b/aqo.c
@@ -228,6 +228,19 @@ _PG_init(void)
 							 NULL
 	);
 
+	DefineCustomBoolVariable(
+							 "aqo.wide_search",
+							 "Search ML data in neighbour feature spaces.",
+							 NULL,
+							 &use_wide_search,
+							 false,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL
+	);
+
 	DefineCustomIntVariable("aqo.join_threshold",
 							"Sets the threshold of number of JOINs in query beyond which AQO is used.",
 							NULL,
diff --git a/aqo.h b/aqo.h
@@ -173,6 +173,7 @@ extern bool	force_collect_stat;
 extern bool aqo_show_hash;
 extern bool aqo_show_details;
 extern int aqo_join_threshold;
+extern bool use_wide_search;
 
 /* Parameters for current query */
 typedef struct QueryContextData
diff --git a/aqo_shared.c b/aqo_shared.c
@@ -191,16 +191,18 @@ aqo_init_shmem(void)
 	{
 		/* First time through ... */
 
-		LWLockInitialize(&aqo_state->lock, LWLockNewTrancheId());
 		aqo_state->dsm_handler = DSM_HANDLE_INVALID;
-
 		aqo_state->qtexts_dsa_handler = DSM_HANDLE_INVALID;
+		aqo_state->data_dsa_handler = DSM_HANDLE_INVALID;
+
 		aqo_state->qtext_trancheid = LWLockNewTrancheId();
+
 		aqo_state->qtexts_changed = false;
-		aqo_state->data_dsa_handler = DSM_HANDLE_INVALID;
+		aqo_state->stat_changed = false;
 		aqo_state->data_changed = false;
 		aqo_state->queries_changed = false;
 
+		LWLockInitialize(&aqo_state->lock, LWLockNewTrancheId());
 		LWLockInitialize(&aqo_state->stat_lock, LWLockNewTrancheId());
 		LWLockInitialize(&aqo_state->qtexts_lock, LWLockNewTrancheId());
 		LWLockInitialize(&aqo_state->data_lock, LWLockNewTrancheId());
@@ -245,7 +247,7 @@ aqo_init_shmem(void)
 	LWLockRegisterTranche(aqo_state->data_lock.tranche, "AQO Data Lock Tranche");
 	LWLockRegisterTranche(aqo_state->queries_lock.tranche, "AQO Queries Lock Tranche");
 
-	if (!IsUnderPostmaster)
+	if (!IsUnderPostmaster && !found)
 	{
 		before_shmem_exit(on_shmem_shutdown, (Datum) 0);
 
@@ -261,8 +263,16 @@ aqo_init_shmem(void)
 static void
 on_shmem_shutdown(int code, Datum arg)
 {
+	Assert(!IsUnderPostmaster);
+
+	/*
+	 * Save ML data to a permanent storage. Do it on postmaster shutdown only
+	 * to save time. We can't do so for query_texts and aqo_data because of DSM
+	 * limits.
+	 */
 	aqo_stat_flush();
 	aqo_queries_flush();
+	return;
 }
 
 Size
diff --git a/cardinality_estimation.c b/cardinality_estimation.c
@@ -24,6 +24,9 @@
 #include "machine_learning.h"
 #include "storage.h"
 
+
+bool use_wide_search = false;
+
 #ifdef AQO_DEBUG_PRINT
 static void
 predict_debug_output(List *clauses, List *selectivities,
@@ -90,7 +93,7 @@ predict_for_relation(List *clauses, List *selectivities, List *relsigns,
 		 */
 
 		/* Try to search in surrounding feature spaces for the same node */
-		if (!load_aqo_data(query_context.fspace_hash, *fss, data, NULL, true))
+		if (!load_aqo_data(query_context.fspace_hash, *fss, data, NULL, use_wide_search))
 			result = -1;
 		else
 		{
diff --git a/conf.add b/conf.add
@@ -1,3 +1,4 @@
 autovacuum = off
 shared_preload_libraries = 'postgres_fdw, aqo'
 max_parallel_workers_per_gather = 1 # switch off parallel workers because of unsteadiness
+aqo.wide_search = 'on'
diff --git a/path_utils.c b/path_utils.c
@@ -151,6 +151,8 @@ hashTempTupleDesc(TupleDesc desc)
 	return s;
 }
 
+#include "storage/lmgr.h"
+
 /*
  * Get list of relation indexes and prepare list of permanent table reloids,
  * list of temporary table reloids (can be changed between query launches) and
@@ -173,6 +175,8 @@ get_list_of_relids(PlannerInfo *root, Relids relids, RelSortOut *rels)
 		HeapTuple		htup;
 		Form_pg_class	classForm;
 		char		   *relname = NULL;
+		Oid				relrewrite;
+		char			relpersistence;
 
 		entry = planner_rt_fetch(index, root);
 
@@ -187,34 +191,39 @@ get_list_of_relids(PlannerInfo *root, Relids relids, RelSortOut *rels)
 		if (!HeapTupleIsValid(htup))
 			elog(PANIC, "cache lookup failed for reloid %u", entry->relid);
 
+		/* Copy the fields from syscache and release the slot as quickly as possible. */
 		classForm = (Form_pg_class) GETSTRUCT(htup);
+		relpersistence = classForm->relpersistence;
+		relrewrite = classForm->relrewrite;
+		relname = pstrdup(NameStr(classForm->relname));
+		ReleaseSysCache(htup);
 
-		if (classForm->relpersistence == RELPERSISTENCE_TEMP)
+		if (relpersistence == RELPERSISTENCE_TEMP)
 		{
 			/* The case of temporary table */
 
-			Relation trel = relation_open(entry->relid, NoLock);
-			TupleDesc tdesc = RelationGetDescr(trel);
+			Relation	trel;
+			TupleDesc	tdesc;
 
+			trel = relation_open(entry->relid, NoLock);
+			tdesc = RelationGetDescr(trel);
+			Assert(CheckRelationLockedByMe(trel, AccessShareLock, true));
 			hashes = lappend_uint64(hashes, hashTempTupleDesc(tdesc));
 			relation_close(trel, NoLock);
 		}
 		else
 		{
 			/* The case of regular table */
 			relname = quote_qualified_identifier(
-				get_namespace_name(get_rel_namespace(entry->relid)),
-				classForm->relrewrite ?
-										get_rel_name(classForm->relrewrite) :
-										NameStr(classForm->relname));
+						get_namespace_name(get_rel_namespace(entry->relid)),
+							relrewrite ? get_rel_name(relrewrite) : relname);
+
 			hashes = lappend_uint64(hashes, DatumGetInt64(hash_any_extended(
 											(unsigned char *) relname,
 											strlen(relname), 0)));
 
 			hrels = lappend_oid(hrels, entry->relid);
 		}
-
-		ReleaseSysCache(htup);
 	}
 
 	rels->hrels = list_concat(rels->hrels, hrels);
diff --git a/postprocessing.c b/postprocessing.c
@@ -173,7 +173,6 @@ restore_selectivities(List *clauselist, List *relidslist, JoinType join_type,
 {
 	List		*lst = NIL;
 	ListCell	*l;
-	int			i = 0;
 	bool		parametrized_sel;
 	int			nargs;
 	int			*args_hash;
@@ -222,7 +221,6 @@ restore_selectivities(List *clauselist, List *relidslist, JoinType join_type,
 		Assert(cur_sel > 0);
 
 		lst = lappend(lst, cur_sel);
-		i++;
 	}
 
 	if (parametrized_sel)
diff --git a/storage.c b/storage.c
@@ -303,7 +303,9 @@ aqo_stat_store(uint64 queryid, bool use_aqo,
 		entry->exec_time[pos] = exec_time;
 		entry->est_error[pos] = est_error;
 	}
+
 	entry = memcpy(palloc(sizeof(StatEntry)), entry, sizeof(StatEntry));
+	aqo_state->stat_changed = true;
 	LWLockRelease(&aqo_state->stat_lock);
 	return entry;
 }
@@ -425,14 +427,24 @@ aqo_stat_flush(void)
 	int				ret;
 	long			entries;
 
-	LWLockAcquire(&aqo_state->stat_lock, LW_SHARED);
+	/* Use exclusive lock to prevent concurrent flushing in different backends. */
+	LWLockAcquire(&aqo_state->stat_lock, LW_EXCLUSIVE);
+
+	if (!aqo_state->stat_changed)
+		/* Hash table wasn't changed, meaningless to store it in permanent storage */
+		goto end;
+
 	entries = hash_get_num_entries(stat_htab);
 	hash_seq_init(&hash_seq, stat_htab);
 	ret = data_store(PGAQO_STAT_FILE, _form_stat_record_cb, entries,
 					 (void *) &hash_seq);
 	if (ret != 0)
 		hash_seq_term(&hash_seq);
+	else
+		/* Hash table and disk storage are now consistent */
+		aqo_state->stat_changed = false;
 
+end:
 	LWLockRelease(&aqo_state->stat_lock);
 }
 
@@ -469,7 +481,7 @@ aqo_qtexts_flush(void)
 	long			entries;
 
 	dsa_init();
-	LWLockAcquire(&aqo_state->qtexts_lock, LW_SHARED);
+	LWLockAcquire(&aqo_state->qtexts_lock, LW_EXCLUSIVE);
 
 	if (!aqo_state->qtexts_changed)
 		/* XXX: mull over forced mode. */
@@ -481,7 +493,9 @@ aqo_qtexts_flush(void)
 					 (void *) &hash_seq);
 	if (ret != 0)
 		hash_seq_term(&hash_seq);
-	aqo_state->qtexts_changed = false;
+	else
+		/* Hash table and disk storage are now consistent */
+		aqo_state->qtexts_changed = false;
 
 end:
 	LWLockRelease(&aqo_state->qtexts_lock);
@@ -531,7 +545,7 @@ aqo_data_flush(void)
 	long			entries;
 
 	dsa_init();
-	LWLockAcquire(&aqo_state->data_lock, LW_SHARED);
+	LWLockAcquire(&aqo_state->data_lock, LW_EXCLUSIVE);
 
 	if (!aqo_state->data_changed)
 		/* XXX: mull over forced mode. */
@@ -548,6 +562,7 @@ aqo_data_flush(void)
 		 */
 		hash_seq_term(&hash_seq);
 	else
+		/* Hash table and disk storage are now consistent */
 		aqo_state->data_changed = false;
 end:
 	LWLockRelease(&aqo_state->data_lock);
@@ -574,14 +589,22 @@ aqo_queries_flush(void)
 	int				ret;
 	long			entries;
 
-	LWLockAcquire(&aqo_state->queries_lock, LW_SHARED);
+	LWLockAcquire(&aqo_state->queries_lock, LW_EXCLUSIVE);
+
+	if (!aqo_state->queries_changed)
+		goto end;
+
 	entries = hash_get_num_entries(queries_htab);
 	hash_seq_init(&hash_seq, queries_htab);
 	ret = data_store(PGAQO_QUERIES_FILE, _form_queries_record_cb, entries,
 					 (void *) &hash_seq);
 	if (ret != 0)
 		hash_seq_term(&hash_seq);
+	else
+		/* Hash table and disk storage are now consistent */
+		aqo_state->queries_changed = false;
 
+end:
 	LWLockRelease(&aqo_state->queries_lock);
 }
 
@@ -621,7 +644,8 @@ data_store(const char *filename, form_record_t callback,
 		goto error;
 	}
 
-	(void) durable_rename(tmpfile, filename, LOG);
+	/* Parallel (re)writing into a file haven't happen. */
+	(void) durable_rename(tmpfile, filename, PANIC);
 	elog(LOG, "[AQO] %d records stored in file %s.", counter, filename);
 	return 0;
 
@@ -839,7 +863,7 @@ aqo_queries_load(void)
 
 	LWLockAcquire(&aqo_state->queries_lock, LW_EXCLUSIVE);
 
-	/* Load on postmaster sturtup. So no any concurrent actions possible here. */
+	/* Load on postmaster startup. So no any concurrent actions possible here. */
 	Assert(hash_get_num_entries(queries_htab) == 0);
 
 	data_load(PGAQO_QUERIES_FILE, _deform_queries_record_cb, NULL);
@@ -926,6 +950,9 @@ data_load(const char *filename, deform_record_t callback, void *ctx)
 static void
 on_shmem_shutdown(int code, Datum arg)
 {
+	/*
+	 * XXX: It can be expensive to rewrite a file on each shutdown of a backend.
+	 */
 	aqo_qtexts_flush();
 	aqo_data_flush();
 }
@@ -1201,6 +1228,7 @@ _aqo_data_remove(data_key *key)
 
 		if (hash_search(data_htab, key, HASH_REMOVE, NULL) == NULL)
 			elog(PANIC, "[AQO] Inconsistent data hash table");
+
 		aqo_state->data_changed = true;
 	}
 
@@ -1270,8 +1298,9 @@ aqo_data_store(uint64 fs, int fss, OkNNrdata *data, List *reloids)
 	char	   *ptr;
 	ListCell   *lc;
 	size_t		size;
-	bool			tblOverflow;
-	HASHACTION		action;
+	bool		tblOverflow;
+	HASHACTION	action;
+	bool		result;
 
 	Assert(!LWLockHeldByMe(&aqo_state->data_lock));
 
@@ -1387,8 +1416,9 @@ aqo_data_store(uint64 fs, int fss, OkNNrdata *data, List *reloids)
 
 	aqo_state->data_changed = true;
 end:
+	result = aqo_state->data_changed;
 	LWLockRelease(&aqo_state->data_lock);
-	return aqo_state->data_changed;
+	return result;
 }
 
 static void
@@ -1496,7 +1526,7 @@ load_aqo_data(uint64 fs, int fss, OkNNrdata *data, List **reloids,
 
 	dsa_init();
 
-	LWLockAcquire(&aqo_state->data_lock, LW_EXCLUSIVE);
+	LWLockAcquire(&aqo_state->data_lock, LW_SHARED);
 
 	if (!wideSearch)
 	{
@@ -1631,7 +1661,8 @@ aqo_data(PG_FUNCTION_ARGS)
 		ptr += sizeof(data_key);
 
 		if (entry->cols > 0)
-		values[AD_FEATURES] = PointerGetDatum(form_matrix((double *)ptr, entry->rows, entry->cols));
+			values[AD_FEATURES] = PointerGetDatum(form_matrix((double *) ptr,
+													entry->rows, entry->cols));
 		else
 			nulls[AD_FEATURES] = true;
 
@@ -1719,7 +1750,9 @@ aqo_data_reset(void)
 			elog(ERROR, "[AQO] hash table corrupted");
 		num_remove++;
 	}
-	aqo_state->data_changed = true;
+
+	if (num_remove > 0)
+		aqo_state->data_changed = true;
 	LWLockRelease(&aqo_state->data_lock);
 	if (num_remove != num_entries)
 		elog(ERROR, "[AQO] Query ML memory storage is corrupted or parallel access without a lock has detected.");
@@ -1831,6 +1864,7 @@ aqo_queries_store(uint64 queryid,
 	entry->use_aqo = use_aqo;
 	entry->auto_tuning = auto_tuning;
 
+	aqo_state->queries_changed = true;
 	LWLockRelease(&aqo_state->queries_lock);
 	return true;
 }
@@ -1856,7 +1890,10 @@ aqo_queries_reset(void)
 			elog(ERROR, "[AQO] hash table corrupted");
 		num_remove++;
 	}
-	aqo_state->queries_changed = true;
+
+	if (num_remove > 0)
+		aqo_state->queries_changed = true;
+
 	LWLockRelease(&aqo_state->queries_lock);
 
 	if (num_remove != num_entries - 1)

Original file line number	Diff line number	Diff line change
`@@ -173,7 +173,6 @@ restore_selectivities(List clauselist, List relidslist, JoinType join_type,`
`173`	`173`	`{`
`174`	`174`	`List *lst = NIL;`
`175`	`175`	`ListCell *l;`
`176`		`- int i = 0;`
`177`	`176`	`bool parametrized_sel;`
`178`	`177`	`int nargs;`
`179`	`178`	`int *args_hash;`
`@@ -222,7 +221,6 @@ restore_selectivities(List clauselist, List relidslist, JoinType join_type,`
`222`	`221`	`Assert(cur_sel > 0);`
`223`	`222`
`224`	`223`	`lst = lappend(lst, cur_sel);`
`225`		`- i++;`
`226`	`224`	`}`
`227`	`225`
`228`	`226`	`if (parametrized_sel)`