Fix release of memory of dropped in-memory tables

Use object access hook instead of event trigger.  That makes it possible to
handle DROP SCHEMA and DROP DATABASE.

Author: Alexander Korotkov

contrib/in_memory/Makefile

@@ -60,5 +60,5 @@ isolationcheck: | submake-isolation submake-in_memory temp-install
 
 temp-install: EXTRA_INSTALL=contrib/in_memory
 
-standby-check: submake-in_memory temp-install
+prove-check: submake-in_memory temp-install
 	$(prove_check)

contrib/in_memory/fdw.c

@@ -88,6 +88,7 @@ is_valid_option(const char *option, Oid context)
 }
 
 static void in_memory_fdw_validate_index(List* index_params, bool pk);
+static void drop_fdw_by_descr(SharedDescr *descr);
 
 Datum
 in_memory_fdw_validator(PG_FUNCTION_ARGS)
@@ -1659,20 +1660,32 @@ void
 drop_fdw(Oid reloid)
 {
 	SharedDescr			   *descr;
-	int						i;
-	DropFDWUndoStackItem   *item;
 
 	descr = find_shared_descr(MyDatabaseId, reloid);
 	if (descr)
-	{
-		for (i = 0; i < descr->ntrees; i++)
-			drop_btree(descr->tree[i].root, descr->tree[i].meta);
-		item = MemoryContextAlloc(CurTransactionContext, sizeof(*item));
-		item->header.cb = drop_fdw_undo_callback;
-		item->datoid = descr->key.datoid;
-		item->reloid = descr->key.reloid;
-		add_undo_item(&item->header);
-	}
+		drop_fdw_by_descr(descr);
+}
+
+void
+drop_database_fdw(Oid datoid)
+{
+	iterate_database_fdws(datoid, drop_fdw_by_descr);
+}
+
+static void
+drop_fdw_by_descr(SharedDescr *descr)
+{
+	int						i;
+	DropFDWUndoStackItem   *item;
+
+	for (i = 0; i < descr->ntrees; i++)
+		drop_btree(descr->tree[i].root, descr->tree[i].meta);
+
+	item = MemoryContextAlloc(CurTransactionContext, sizeof(*item));
+	item->header.cb = drop_fdw_undo_callback;
+	item->datoid = descr->key.datoid;
+	item->reloid = descr->key.reloid;
+	add_undo_item(&item->header);
 }
 
 bool

contrib/in_memory/fdw.h

@@ -25,6 +25,7 @@ extern void im_fdw_init_root(void *buf, bool found);
 extern bool check_fdw_is_in_memory_relid(Oid relid);
 extern void truncate_fdw(Oid relid);
 extern void drop_fdw(Oid relid);
+extern void drop_database_fdw(Oid datoid);
 extern void drop_fdw_descr(Oid datoid, Oid reloid);
 extern void free_fdw(Oid datoid, Oid reloid);
 extern void invalidate_private_fdw_descr(Oid reloid);

contrib/in_memory/fdw_descr.c

@@ -296,6 +296,34 @@ find_shared_descr(Oid datoid, Oid reloid)
 									 COMMITSEQNO_INPROGRESS);
 }
 
+/*
+ * Iterate over all present FDWs in given database.
+ */
+void
+iterate_database_fdws(Oid datoid, DescrCallback callback)
+{
+	BTreeIterator	   *iter;
+	DescrKey			key;
+	SharedDescr		   *descr;
+
+	key.datoid = datoid;
+	key.reloid = (Oid) 0;
+
+	iter = make_btree_iterator(&im_fdw_meta_cache_tree, (Pointer) &key,
+							   BTreeKey, COMMITSEQNO_INPROGRESS);
+
+	while ((descr = (SharedDescr *) btree_iterate(iter)) != NULL)
+	{
+		if (descr->key.datoid != datoid)
+			break;
+
+		callback(descr);
+
+		pfree(descr);
+	}
+	free_btree_iterator(iter);
+}
+
 void
 drop_fdw_descr(Oid datoid, Oid reloid)
 {

contrib/in_memory/fdw_internal.h

@@ -191,4 +191,8 @@ typedef struct IMFdwPlanState
 	List		   *local_conds;
 } IMFdwPlanState;
 
+typedef void (*DescrCallback) (SharedDescr *descr);
+
+extern void iterate_database_fdws(Oid datoid, DescrCallback callback);
+
 #endif /* __IN_MEMORY_FDW_INTERNAL_H__ */

contrib/in_memory/in_memory.c

@@ -17,6 +17,8 @@
 #include "access/heapam.h"
 #include "access/htup_details.h"
 #include "catalog/namespace.h"
+#include "catalog/objectaccess.h"
+#include "catalog/pg_database.h"
 #include "commands/event_trigger.h"
 #include "commands/tablecmds.h"
 #include "executor/spi.h"
@@ -50,6 +52,7 @@ void		_PG_init(void);
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 static ProcessUtility_hook_type next_ProcessUtility_hook = NULL;
 static logical_message_redo_hook_type prev_logical_message_redo_hook = NULL;
+static object_access_hook_type prev_object_access_hook = NULL;
 
 static Pointer				shared_segment = NULL;
 static bool					shared_segment_initialized = false;
@@ -81,6 +84,8 @@ static void in_memory_utility_command(PlannedStmt *pstmt,
 									  char *completionTag);
 static void in_memory_logical_message_hook(xl_logical_message *xlrec);
 static void in_memory_relcache_hook(Datum arg, Oid relid);
+static void in_memory_object_access_hook(ObjectAccessType access, Oid classId,
+										 Oid objectId, int subId, void *arg);
 
 PG_FUNCTION_INFO_V1(in_memory_fsm_test);
 PG_FUNCTION_INFO_V1(in_memory_page_stats);
@@ -176,6 +181,8 @@ _PG_init(void)
 	ProcessUtility_hook = in_memory_utility_command;
 	prev_logical_message_redo_hook = logical_message_redo_hook;
 	logical_message_redo_hook = in_memory_logical_message_hook;
+	prev_object_access_hook = object_access_hook;
+	object_access_hook = in_memory_object_access_hook;
 	RegisterXactCallback(undo_xact_callback, NULL);
 	RegisterSubXactCallback(undo_subxact_callback, NULL);
 	CacheRegisterRelcacheCallback(in_memory_relcache_hook, PointerGetDatum(NULL));
@@ -515,48 +522,29 @@ in_memory_relcache_hook(Datum arg, Oid relid)
 	invalidate_private_fdw_descr(relid);
 }
 
-Datum
-in_memory_event_trigger(PG_FUNCTION_ARGS)
+static void
+in_memory_object_access_hook(ObjectAccessType access, Oid classId,
+							 Oid objectId, int subId, void *arg)
 {
-	EventTriggerData   *trigdata;
-	int					ret;
-	uint64				proc;
-
-	if (!CALLED_AS_EVENT_TRIGGER(fcinfo))  /* internal error */
-		elog(ERROR, "in_memory_event_trigger: not fired by event trigger manager");
-
-	trigdata = (EventTriggerData *) fcinfo->context;
-
-	if (strcmp(trigdata->event, "sql_drop"))
-		PG_RETURN_NULL();
-
-	SPI_connect();
-	ret = SPI_execute(
-		"SELECT "
-			"objid "
-		"FROM "
-			"pg_event_trigger_dropped_objects() "
-		"WHERE classid = 'pg_catalog.pg_class'::regclass;", true, 0);
-	proc = SPI_processed;
-	elog(LOG, "%u", (uint32) proc);
-
-	if (ret > 0 && SPI_tuptable != NULL)
+	if (access == OAT_DROP && classId == RelationRelationId)
 	{
-		TupleDesc		tupdesc = SPI_tuptable->tupdesc;
-		SPITupleTable  *tuptable = SPI_tuptable;
-		uint64			i;
-		bool			isnull;
-
-		for (i = 0; i < proc; i++)
-		{
-			HeapTuple		tuple = tuptable->vals[i];
-			Oid				reloid;
+		drop_fdw(objectId);
+	}
 
-			reloid = DatumGetObjectId(SPI_getbinval(tuple, tupdesc, 1, &isnull));
-			drop_fdw(reloid);
-		}
+	if (access == OAT_DROP && classId == DatabaseRelationId)
+	{
+		drop_database_fdw(objectId);
 	}
-	SPI_finish();
 
+	if (prev_object_access_hook)
+		prev_object_access_hook(access, classId, objectId, subId, arg);
+}
+
+/*
+ * No longer used, but reserved for future usage.
+ */
+Datum
+in_memory_event_trigger(PG_FUNCTION_ARGS)
+{
 	PG_RETURN_NULL();
 }

contrib/in_memory/t/002_drop.pl

@@ -0,0 +1,38 @@
+# Test in_memory tables memory releazed on drop.
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More tests => 2;
+
+my $node;
+
+# Initialize node
+$node = get_new_node('master');
+$node->init(allows_streaming => 1);
+$node->append_conf('postgresql.conf', "shared_preload_libraries = 'in_memory'");
+$node->start;
+
+# Create in_memory table master and fill it with some data
+$node->safe_psql("postgres", "CREATE EXTENSION in_memory;");
+$node->safe_psql("postgres", "CREATE DATABASE test;");
+
+# Make table, and then drop extension
+$node->safe_psql("test", "CREATE EXTENSION in_memory;");
+$node->safe_psql("test", "CREATE FOREIGN TABLE im_test (id int8 not null, value text) SERVER in_memory;");
+$node->safe_psql("test", "INSERT INTO im_test SELECT id, 'val' || id FROM generate_series(1, 10000) id;");
+$node->safe_psql("test", "DROP EXTENSION in_memory CASCADE;");
+
+# Check that memory is releazed
+my $result = $node->safe_psql("postgres", "SELECT busy_pages FROM in_memory_page_stats();");
+is($result, "2", "in_memory table memory releazed after extension drop");
+
+# Make table again, and then drop database
+$node->safe_psql("test", "CREATE EXTENSION in_memory;");
+$node->safe_psql("test", "CREATE FOREIGN TABLE im_test (id int8 not null, value text) SERVER in_memory;");
+$node->safe_psql("test", "INSERT INTO im_test SELECT id, 'val' || id FROM generate_series(1, 10000) id;");
+$node->safe_psql("postgres", "DROP DATABASE test;");
+
+# Check that memory is releazed
+$result = $node->safe_psql("postgres", "SELECT busy_pages FROM in_memory_page_stats();");
+is($result, "2", "in_memory table memory releazed after database drop");