add "clear" functionsadvice2

2 files changed, 102 insertions, 62 deletions

diff --git a/contrib/pg_plan_advice/pg_plan_advice--1.0.sql b/contrib/pg_plan_advice/pg_plan_advice--1.0.sql
index 63a82666df..729950bcce 100644
--- a/contrib/pg_plan_advice/pg_plan_advice--1.0.sql
+++ b/contrib/pg_plan_advice/pg_plan_advice--1.0.sql
@@ -3,6 +3,16 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "CREATE EXTENSION pg_plan_advice" to load this file. \quit
 
+CREATE FUNCTION pg_clear_collected_local_advice()
+RETURNS void
+AS 'MODULE_PATHNAME', 'pg_clear_collected_local_advice'
+LANGUAGE C STRICT;
+
+CREATE FUNCTION pg_clear_collected_shared_advice()
+RETURNS void
+AS 'MODULE_PATHNAME', 'pg_clear_collected_shared_advice'
+LANGUAGE C STRICT;
+
 CREATE FUNCTION pg_get_collected_local_advice(
 	OUT id bigint,
 	OUT userid oid,
diff --git a/contrib/pg_plan_advice/pgpa_collector.c b/contrib/pg_plan_advice/pgpa_collector.c
index 0bfb0fb85b..7fffc55376 100644
--- a/contrib/pg_plan_advice/pgpa_collector.c
+++ b/contrib/pg_plan_advice/pgpa_collector.c
@@ -22,6 +22,8 @@
 #include "utils/builtins.h"
 #include "utils/timestamp.h"
 
+PG_FUNCTION_INFO_V1(pg_clear_collected_local_advice);
+PG_FUNCTION_INFO_V1(pg_clear_collected_shared_advice);
 PG_FUNCTION_INFO_V1(pg_get_collected_local_advice);
 PG_FUNCTION_INFO_V1(pg_get_collected_shared_advice);
 
@@ -106,9 +108,9 @@ static pgpa_collected_advice *pgpa_make_collected_advice(Oid userid,
 														 dsa_area *area,
 														 dsa_pointer *result);
 static void pgpa_store_local_advice(pgpa_collected_advice *ca);
-static void pgpa_trim_local_advice(void);
+static void pgpa_trim_local_advice(int limit);
 static void pgpa_store_shared_advice(dsa_pointer ca_pointer);
-static void pgpa_trim_shared_advice(void);
+static void pgpa_trim_shared_advice(int limit);
 
 /* Helper function to extract the query string from pgpa_collected_advice */
 static inline const char *
@@ -257,59 +259,7 @@ pgpa_store_local_advice(pgpa_collected_advice *ca)
 	++la->next_id;
 
 	/* If we've exceeded the storage limit, discard old data. */
-	pgpa_trim_local_advice();
-}
-
-/*
- * Remove local advice in excess of pg_plan_advice.local_collection_limit.
- */
-static void
-pgpa_trim_local_advice(void)
-{
-	pgpa_local_advice *la = local_collector;
-	uint64		current_count;
-	uint64		trim_count;
-	uint64		total_chunk_count;
-	uint64		trim_chunk_count;
-	uint64		remaining_chunk_count;
-
-	/* If we haven't yet reached the limit, there's nothing to do. */
-	current_count = la->next_id - la->oldest_id;
-	if (current_count < pg_plan_advice_local_collection_limit)
-		return;
-
-	/* Free enough entries to get us back down to the limit. */
-	trim_count = current_count - pg_plan_advice_local_collection_limit;
-	while (trim_count > 0)
-	{
-		uint64		chunk_number;
-		uint64		chunk_offset;
-
-		chunk_number = (la->oldest_id - la->base_id) / ADVICE_CHUNK_SIZE;
-		chunk_offset = (la->oldest_id - la->base_id) % ADVICE_CHUNK_SIZE;
-
-		Assert(la->chunks[chunk_number]->entries[chunk_offset] != NULL);
-		pfree(la->chunks[chunk_number]->entries[chunk_offset]);
-		la->chunks[chunk_number]->entries[chunk_offset] = NULL;
-		++la->oldest_id;
-		--trim_count;
-	}
-
-	/* Free any chunks that are now entirely unused. */
-	trim_chunk_count = (la->oldest_id - la->base_id) / ADVICE_CHUNK_SIZE;
-	for (uint64 n = 0; n < trim_chunk_count; ++n)
-		pfree(la->chunks[n]);
-
-	/* Slide remaining chunk pointers back toward the base of the array. */
-	total_chunk_count = (la->next_id - la->base_id +
-						 ADVICE_CHUNK_SIZE - 1) / ADVICE_CHUNK_SIZE;
-	remaining_chunk_count = total_chunk_count - trim_chunk_count;
-	if (remaining_chunk_count > 0)
-		memmove(&la->chunks[0], &la->chunks[trim_chunk_count],
-				sizeof(pgpa_local_advice_chunk *) * remaining_chunk_count);
-
-	/* Adjust base ID value accordingly. */
-	la->base_id += trim_chunk_count * ADVICE_CHUNK_SIZE;
+	pgpa_trim_local_advice(pg_plan_advice_local_collection_limit);
 }
 
 /*
@@ -396,17 +346,71 @@ pgpa_store_shared_advice(dsa_pointer ca_pointer)
 	++sa->next_id;
 
 	/* If we've exceeded the storage limit, discard old data. */
-	pgpa_trim_shared_advice();
+	pgpa_trim_shared_advice(pg_plan_advice_shared_collection_limit);
 
 	/* Release lock on shared state. */
 	LWLockRelease(&state->lock);
 }
 
 /*
- * Remove shared advice in excess of pg_plan_advice.shared_collection_limit.
+ * Discard collected advice stored in backend-local memory in excess of the
+ * specified limit.
+ */
+static void
+pgpa_trim_local_advice(int limit)
+{
+	pgpa_local_advice *la = local_collector;
+	uint64		current_count;
+	uint64		trim_count;
+	uint64		total_chunk_count;
+	uint64		trim_chunk_count;
+	uint64		remaining_chunk_count;
+
+	/* If we haven't yet reached the limit, there's nothing to do. */
+	current_count = la->next_id - la->oldest_id;
+	if (current_count <= limit)
+		return;
+
+	/* Free enough entries to get us back down to the limit. */
+	trim_count = current_count - limit;
+	while (trim_count > 0)
+	{
+		uint64		chunk_number;
+		uint64		chunk_offset;
+
+		chunk_number = (la->oldest_id - la->base_id) / ADVICE_CHUNK_SIZE;
+		chunk_offset = (la->oldest_id - la->base_id) % ADVICE_CHUNK_SIZE;
+
+		Assert(la->chunks[chunk_number]->entries[chunk_offset] != NULL);
+		pfree(la->chunks[chunk_number]->entries[chunk_offset]);
+		la->chunks[chunk_number]->entries[chunk_offset] = NULL;
+		++la->oldest_id;
+		--trim_count;
+	}
+
+	/* Free any chunks that are now entirely unused. */
+	trim_chunk_count = (la->oldest_id - la->base_id) / ADVICE_CHUNK_SIZE;
+	for (uint64 n = 0; n < trim_chunk_count; ++n)
+		pfree(la->chunks[n]);
+
+	/* Slide remaining chunk pointers back toward the base of the array. */
+	total_chunk_count = (la->next_id - la->base_id +
+						 ADVICE_CHUNK_SIZE - 1) / ADVICE_CHUNK_SIZE;
+	remaining_chunk_count = total_chunk_count - trim_chunk_count;
+	if (remaining_chunk_count > 0)
+		memmove(&la->chunks[0], &la->chunks[trim_chunk_count],
+				sizeof(pgpa_local_advice_chunk *) * remaining_chunk_count);
+
+	/* Adjust base ID value accordingly. */
+	la->base_id += trim_chunk_count * ADVICE_CHUNK_SIZE;
+}
+
+/*
+ * Discard collected advice stored in shared memory in excess of the
+ * specified limit.
  */
 static void
-pgpa_trim_shared_advice(void)
+pgpa_trim_shared_advice(int limit)
 {
 	dsa_area   *area = pg_plan_advice_dsa_area();
 	pgpa_shared_advice *sa = shared_collector;
@@ -419,14 +423,14 @@ pgpa_trim_shared_advice(void)
 
 	/* If we haven't yet reached the limit, there's nothing to do. */
 	current_count = sa->next_id - sa->oldest_id;
-	if (current_count < pg_plan_advice_shared_collection_limit)
+	if (current_count <= limit)
 		return;
 
 	/* Get a pointer to the chunk array. */
 	chunk_array = dsa_get_address(area, sa->chunks);
 
 	/* Free enough entries to get us back down to the limit. */
-	trim_count = current_count - pg_plan_advice_shared_collection_limit;
+	trim_count = current_count - limit;
 	while (trim_count > 0)
 	{
 		uint64		chunk_number;
@@ -462,7 +466,33 @@ pgpa_trim_shared_advice(void)
 }
 
 /*
- * SQL-callable SRF to return locally collected advice
+ * SQL-callable function to discard advice collected in backend-local memory
+ */
+Datum
+pg_clear_collected_local_advice(PG_FUNCTION_ARGS)
+{
+	pgpa_trim_local_advice(0);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * SQL-callable function to discard advice collected in backend-local memory
+ */
+Datum
+pg_clear_collected_shared_advice(PG_FUNCTION_ARGS)
+{
+	pgpa_shared_state *state = pg_plan_advice_attach();
+
+	LWLockAcquire(&state->lock, LW_EXCLUSIVE);
+	pgpa_trim_shared_advice(0);
+	LWLockRelease(&state->lock);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * SQL-callable SRF to return advice collected in backend-local memory
  */
 Datum
 pg_get_collected_local_advice(PG_FUNCTION_ARGS)
@@ -514,7 +544,7 @@ pg_get_collected_local_advice(PG_FUNCTION_ARGS)
 }
 
 /*
- * SQL-callable SRF to return locally collected advice
+ * SQL-callable SRF to return advice collected in shared memory
  */
 Datum
 pg_get_collected_shared_advice(PG_FUNCTION_ARGS)