Add handling of situation when AQO shmem storage is overflowed.

danolivo · danolivo · commit bd5b63a492fe · 2022-06-29T11:15:44.000+05:00
Our tactics here: log a problem, switch backend into CONTROLLED mode and go
further.

TODO:
1) change aqo.mode for all backends;
2) switch to FROZEN mode if data storage is full.
3) How to process overflow of DSM?
diff --git a/aqo.c b/aqo.c
@@ -208,7 +208,34 @@ _PG_init(void)
 							0,
 							NULL,
 							NULL,
-							NULL);
+							NULL
+	);
+
+	DefineCustomIntVariable("aqo.fs_max_items",
+							"Max number of feature spaces that AQO can operate with.",
+							NULL,
+							&fs_max_items,
+							1000,
+							1, INT_MAX,
+							PGC_SUSET,
+							0,
+							NULL,
+							NULL,
+							NULL
+	);
+
+	DefineCustomIntVariable("aqo.fss_max_items",
+							"Max number of feature subspaces that AQO can operate with.",
+							NULL,
+							&fss_max_items,
+							1000,
+							0, INT_MAX,
+							PGC_SUSET,
+							0,
+							NULL,
+							NULL,
+							NULL
+	);
 
 	prev_shmem_startup_hook						= shmem_startup_hook;
 	shmem_startup_hook							= aqo_init_shmem;
diff --git a/aqo_shared.c b/aqo_shared.c
@@ -26,8 +26,8 @@ shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 AQOSharedState *aqo_state = NULL;
 HTAB *fss_htab = NULL;
 static int aqo_htab_max_items = 1000;
-static int fs_max_items = 1000; /* Max number of different feature spaces in ML model */
-static int fss_max_items = 10000;
+int fs_max_items = 1; /* Max number of different feature spaces in ML model */
+int fss_max_items = 1; /* Max number of different feature subspaces in ML model */
 static uint32 temp_storage_size = 1024 * 1024 * 10; /* Storage size, in bytes */
 static dsm_segment *seg = NULL;
 
@@ -217,30 +217,26 @@ aqo_init_shmem(void)
 
 	info.keysize = sizeof(((StatEntry *) 0)->queryid);
 	info.entrysize = sizeof(StatEntry);
-	stat_htab = ShmemInitHash("AQO Stat HTAB",
-							  fs_max_items, fs_max_items,
+	stat_htab = ShmemInitHash("AQO Stat HTAB", 64, fs_max_items,
 							  &info, HASH_ELEM | HASH_BLOBS);
 
 	/* Init shared memory table for query texts */
 	info.keysize = sizeof(((QueryTextEntry *) 0)->queryid);
 	info.entrysize = sizeof(QueryTextEntry);
-	qtexts_htab = ShmemInitHash("AQO Query Texts HTAB",
-								fs_max_items, fs_max_items,
+	qtexts_htab = ShmemInitHash("AQO Query Texts HTAB", 64, fs_max_items,
 								&info, HASH_ELEM | HASH_BLOBS);
 
 	/* Shared memory hash table for the data */
 	info.keysize = sizeof(data_key);
 	info.entrysize = sizeof(DataEntry);
-	data_htab = ShmemInitHash("AQO Data HTAB",
-							  fss_max_items, fss_max_items,
+	data_htab = ShmemInitHash("AQO Data HTAB", 64, fss_max_items,
 							  &info, HASH_ELEM | HASH_BLOBS);
 
 	/* Shared memory hash table for queries */
 	info.keysize = sizeof(((QueriesEntry *) 0)->queryid);
 	info.entrysize = sizeof(QueriesEntry);
-	queries_htab = ShmemInitHash("AQO Queries HTAB",
-							  fs_max_items, fs_max_items,
-							  &info, HASH_ELEM | HASH_BLOBS);
+	queries_htab = ShmemInitHash("AQO Queries HTAB", 64, fs_max_items,
+								 &info, HASH_ELEM | HASH_BLOBS);
 
 	LWLockRelease(AddinShmemInitLock);
 	LWLockRegisterTranche(aqo_state->lock.tranche, "AQO");
diff --git a/aqo_shared.h b/aqo_shared.h
@@ -50,6 +50,8 @@ extern shmem_startup_hook_type prev_shmem_startup_hook;
 extern AQOSharedState *aqo_state;
 extern HTAB *fss_htab;
 
+extern int fs_max_items; /* Max number of feature spaces that AQO can operate */
+extern int fss_max_items;
 
 extern Size aqo_memsize(void);
 extern void reset_dsm_cache(void);
diff --git a/preprocessing.c b/preprocessing.c
@@ -295,15 +295,34 @@ aqo_planner(Query *parse,
 		 * concurrent addition from another backend we will try to restart
 		 * preprocessing routine.
 		 */
-		aqo_queries_store(query_context.query_hash, query_context.fspace_hash,
+		if (aqo_queries_store(query_context.query_hash, query_context.fspace_hash,
 						  query_context.learn_aqo, query_context.use_aqo,
-						  query_context.auto_tuning);
+						  query_context.auto_tuning))
+		{
+			/*
+			 * Add query text into the ML-knowledge base. Just for further
+			 * analysis. In the case of cached plans we may have NULL query text.
+			 */
+			if (!aqo_qtext_store(query_context.query_hash, query_string))
+			{
+				Assert(0); /* panic only on debug installation */
+				elog(ERROR, "[AQO] Impossible situation was detected. Maybe not enough of shared memory?");
+			}
+		}
+		else
+		{
+			/*
+			 * In the case of problems (shmem overflow, as a typical issue) -
+			 * disable AQO for the query class.
+			 */
+			disable_aqo_for_query();
 
-		/*
-		 * Add query text into the ML-knowledge base. Just for further
-		 * analysis. In the case of cached plans we may have NULL query text.
-		 */
-		aqo_qtext_store(query_context.query_hash, query_string);
+			/*
+			 * Switch AQO to controlled mode. In this mode we wouldn't add new
+			 * query classes, just use and learn on existed set.
+			 */
+			aqo_mode = AQO_MODE_CONTROLLED;
+		}
 	}
 
 	if (force_collect_stat)
diff --git a/storage.c b/storage.c
@@ -204,10 +204,10 @@ add_deactivated_query(uint64 queryid)
 /*
  * Update AQO statistics.
  *
- * Add a record (and replace old, if all stat slots is full) to stat slot for
- * a query class.
+ * Add a record (or update an existed) to stat storage for the query class.
  * Returns a copy of stat entry, allocated in current memory context. Caller is
  * in charge to free this struct after usage.
+ * If stat hash table is full, return NULL and log this fact.
  */
 StatEntry *
 aqo_stat_store(uint64 queryid, bool use_aqo,
@@ -216,16 +216,36 @@ aqo_stat_store(uint64 queryid, bool use_aqo,
 	StatEntry  *entry;
 	bool		found;
 	int			pos;
+	bool		tblOverflow;
+	HASHACTION	action;
 
 	Assert(stat_htab);
 
 	LWLockAcquire(&aqo_state->stat_lock, LW_EXCLUSIVE);
-	entry = (StatEntry *) hash_search(stat_htab, &queryid, HASH_ENTER, &found);
+	tblOverflow = hash_get_num_entries(stat_htab) < fs_max_items ? false : true;
+	action = tblOverflow ? HASH_FIND : HASH_ENTER;
+	entry = (StatEntry *) hash_search(stat_htab, &queryid, action, &found);
 
 	/* Initialize entry on first usage */
 	if (!found)
 	{
-		uint64 qid = entry->queryid;
+		uint64 qid;
+
+		if (action == HASH_FIND)
+		{
+			/*
+			 * Hash table is full. To avoid possible problems - don't try to add
+			 * more, just exit
+			 */
+			LWLockRelease(&aqo_state->stat_lock);
+			ereport(LOG,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("[AQO] Stat storage is full. No more feature spaces can be added."),
+				 errhint("Increase value of aqo.fs_max_items on restart of the instance")));
+			return NULL;
+		}
+
+		qid = entry->queryid;
 		memset(entry, 0, sizeof(StatEntry));
 		entry->queryid = qid;
 	}
@@ -907,6 +927,8 @@ aqo_qtext_store(uint64 queryid, const char *query_string)
 {
 	QueryTextEntry *entry;
 	bool			found;
+	bool			tblOverflow;
+	HASHACTION		action;
 
 	Assert(!LWLockHeldByMe(&aqo_state->qtexts_lock));
 
@@ -916,7 +938,12 @@ aqo_qtext_store(uint64 queryid, const char *query_string)
 	dsa_init();
 
 	LWLockAcquire(&aqo_state->qtexts_lock, LW_EXCLUSIVE);
-	entry = (QueryTextEntry *) hash_search(qtexts_htab, &queryid, HASH_ENTER,
+
+	/* Check hash table overflow */
+	tblOverflow = hash_get_num_entries(qtexts_htab) < fs_max_items ? false : true;
+	action = tblOverflow ? HASH_FIND : HASH_ENTER;
+
+	entry = (QueryTextEntry *) hash_search(qtexts_htab, &queryid, action,
 										   &found);
 
 	/* Initialize entry on first usage */
@@ -925,6 +952,20 @@ aqo_qtext_store(uint64 queryid, const char *query_string)
 		size_t size = strlen(query_string) + 1;
 		char *strptr;
 
+		if (action == HASH_FIND)
+		{
+			/*
+			 * Hash table is full. To avoid possible problems - don't try to add
+			 * more, just exit
+			 */
+			LWLockRelease(&aqo_state->qtexts_lock);
+			ereport(LOG,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("[AQO] Query texts storage is full. No more feature spaces can be added."),
+				 errhint("Increase value of aqo.fs_max_items on restart of the instance")));
+			return false;
+		}
+
 		entry->queryid = queryid;
 		entry->qtext_dp = dsa_allocate(qtext_dsa, size);
 		Assert(DsaPointerIsValid(entry->qtext_dp));
@@ -933,7 +974,7 @@ aqo_qtext_store(uint64 queryid, const char *query_string)
 		aqo_state->qtexts_changed = true;
 	}
 	LWLockRelease(&aqo_state->qtexts_lock);
-	return !found;
+	return true;
 }
 
 Datum
@@ -1089,17 +1130,38 @@ aqo_data_store(uint64 fs, int fss, OkNNrdata *data, List *reloids)
 	char	   *ptr;
 	ListCell   *lc;
 	size_t		size;
+	bool			tblOverflow;
+	HASHACTION		action;
 
 	Assert(!LWLockHeldByMe(&aqo_state->data_lock));
 
 	dsa_init();
 
 	LWLockAcquire(&aqo_state->data_lock, LW_EXCLUSIVE);
-	entry = (DataEntry *) hash_search(data_htab, &key, HASH_ENTER, &found);
+
+	/* Check hash table overflow */
+	tblOverflow = hash_get_num_entries(data_htab) < fss_max_items ? false : true;
+	action = tblOverflow ? HASH_FIND : HASH_ENTER;
+
+	entry = (DataEntry *) hash_search(data_htab, &key, action, &found);
 
 	/* Initialize entry on first usage */
 	if (!found)
 	{
+		if (action == HASH_FIND)
+		{
+			/*
+			 * Hash table is full. To avoid possible problems - don't try to add
+			 * more, just exit
+			 */
+			LWLockRelease(&aqo_state->data_lock);
+			ereport(LOG,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("[AQO] Data storage is full. No more data can be added."),
+				 errhint("Increase value of aqo.fss_max_items on restart of the instance")));
+			return false;
+		}
+
 		entry->cols = data->cols;
 		entry->rows = data->rows;
 		entry->nrels = list_length(reloids);
@@ -1603,11 +1665,13 @@ aqo_queries_remove(PG_FUNCTION_ARGS)
 }
 
 bool
-aqo_queries_store(uint64 queryid, uint64 fs, bool learn_aqo,
-				  bool use_aqo, bool auto_tuning)
+aqo_queries_store(uint64 queryid,
+				  uint64 fs, bool learn_aqo, bool use_aqo, bool auto_tuning)
 {
 	QueriesEntry   *entry;
 	bool			found;
+	bool		tblOverflow;
+	HASHACTION	action;
 
 	Assert(queries_htab);
 
@@ -1616,8 +1680,29 @@ aqo_queries_store(uint64 queryid, uint64 fs, bool learn_aqo,
 		   use_aqo == false && auto_tuning == false));
 
 	LWLockAcquire(&aqo_state->queries_lock, LW_EXCLUSIVE);
-	entry = (QueriesEntry *) hash_search(queries_htab, &queryid, HASH_ENTER,
+
+	/* Check hash table overflow */
+	tblOverflow = hash_get_num_entries(queries_htab) < fs_max_items ? false : true;
+	action = tblOverflow ? HASH_FIND : HASH_ENTER;
+
+	entry = (QueriesEntry *) hash_search(queries_htab, &queryid, action,
 										 &found);
+
+		/* Initialize entry on first usage */
+	if (!found && action == HASH_FIND)
+	{
+		/*
+		 * Hash table is full. To avoid possible problems - don't try to add
+		 * more, just exit
+		 */
+		LWLockRelease(&aqo_state->queries_lock);
+		ereport(LOG,
+			(errcode(ERRCODE_OUT_OF_MEMORY),
+			 errmsg("[AQO] Queries storage is full. No more feature spaces can be added."),
+			 errhint("Increase value of aqo.fs_max_items on restart of the instance")));
+		return false;
+	}
+
 	entry->fs = fs;
 	entry->learn_aqo = learn_aqo;
 	entry->use_aqo = use_aqo;
