poc_shared_fileset_cleanup_on_procexit_v1.patch

diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 14e057cfff..99cb43b1e4 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -1945,49 +1945,6 @@ UpdateWorkerStats(XLogRecPtr last_lsn, TimestampTz send_time, bool reply)
 	}
 }
 
-/*
- * Cleanup function.
- *
- * Called on logical replication worker exit.
- */
-static void
-worker_onexit(int code, Datum arg)
-{
-	HASH_SEQ_STATUS status;
-	StreamXidHash *ent;
-	char		path[MAXPGPATH];
-
-	/* nothing to clean */
-	if (xidhash == NULL)
-		return;
-
-	/*
-	 * Scan complete hash and delete the underlying files for the xids.
-	 * Also release the memory for the shared file sets.
-	 */
-	hash_seq_init(&status, xidhash);
-	while ((ent = (StreamXidHash *) hash_seq_search(&status)) != NULL)
-	{
-		changes_filename(path, MyLogicalRepWorker->subid, ent->xid);
-		BufFileDeleteShared(ent->stream_fileset, path);
-		pfree(ent->stream_fileset);
-
-		/*
-		 * We might not have created the subxact fileset if there is no sub
-		 * transaction.
-		 */
-		if (ent->subxact_fileset)
-		{
-			subxact_filename(path, MyLogicalRepWorker->subid, ent->xid);
-			BufFileDeleteShared(ent->subxact_fileset, path);
-			pfree(ent->subxact_fileset);
-		}
-	}
-
-	/* Remove the xid hash */
-	hash_destroy(xidhash);
-}
-
 /*
  * Apply main loop.
  */
@@ -2012,9 +1969,6 @@ LogicalRepApplyLoop(XLogRecPtr last_received)
 													"LogicalStreamingContext",
 													ALLOCSET_DEFAULT_SIZES);
 
-	/* do cleanup on worker exit (e.g. after DROP SUBSCRIPTION) */
-	before_shmem_exit(worker_onexit, (Datum) 0);
-
 	/* mark as idle, before starting to loop */
 	pgstat_report_activity(STATE_IDLE, NULL);
 
@@ -2515,10 +2469,18 @@ subxact_info_write(Oid subid, TransactionId xid)
 	 */
 	if (ent->subxact_fileset == NULL)
 	{
-		ent->subxact_fileset =
-			MemoryContextAlloc(ApplyContext, sizeof(SharedFileSet));
+		MemoryContext oldctx;
 
+		/*
+		 * Shared fileset handle must be allocated in the persistent context.
+		 * Also, SharedFileSetInit allocate the memory for sharefileset list
+		 * so we need to allocate that in the long term meemory context.
+		 */
+		oldctx = MemoryContextSwitchTo(ApplyContext);
+		ent->subxact_fileset = palloc(sizeof(SharedFileSet));
 		SharedFileSetInit(ent->subxact_fileset, NULL);
+		MemoryContextSwitchTo(oldctx);
+
 		fd = BufFileCreateShared(ent->subxact_fileset, path);
 	}
 	else
@@ -2784,13 +2746,15 @@ stream_open_file(Oid subid, TransactionId xid, bool first_segment)
 	 */
 	if (first_segment)
 	{
-		/*
-		 * Shared fileset handle must be allocated in the persistent context.
-		 */
-		SharedFileSet *fileset =
-		MemoryContextAlloc(ApplyContext, sizeof(SharedFileSet));
+		MemoryContext oldctx;
+		SharedFileSet *fileset;
 
+		/* Shared fileset handle must be allocated in the persistent context */
+		oldctx = MemoryContextSwitchTo(ApplyContext);
+		fileset = palloc(sizeof(SharedFileSet));
 		SharedFileSetInit(fileset, NULL);
+		oldctx = MemoryContextSwitchTo(oldctx);
+
 		stream_fd = BufFileCreateShared(fileset, path);
 
 		/* Remember the fileset for the next stream of the same transaction */
diff --git a/src/backend/storage/file/buffile.c b/src/backend/storage/file/buffile.c
index bde6fa1ef3..502875a09c 100644
--- a/src/backend/storage/file/buffile.c
+++ b/src/backend/storage/file/buffile.c
@@ -364,6 +364,9 @@ BufFileDeleteShared(SharedFileSet *fileset, const char *name)
 		CHECK_FOR_INTERRUPTS();
 	}
 
+	/* Unregister the shared fileset */
+	SharedFileSetUnregister(fileset);
+
 	if (!found)
 		elog(ERROR, "could not delete unknown shared BufFile \"%s\"", name);
 }
diff --git a/src/backend/storage/file/sharedfileset.c b/src/backend/storage/file/sharedfileset.c
index c81d298fc3..3361a8274f 100644
--- a/src/backend/storage/file/sharedfileset.c
+++ b/src/backend/storage/file/sharedfileset.c
@@ -25,10 +25,14 @@
 #include "common/hashfn.h"
 #include "miscadmin.h"
 #include "storage/dsm.h"
+#include "storage/ipc.h"
 #include "storage/sharedfileset.h"
 #include "utils/builtins.h"
 
+static List *filesetlist = NIL;
+
 static void SharedFileSetOnDetach(dsm_segment *segment, Datum datum);
+static void SharedFileSetOnProcExit(int status, Datum arg);
 static void SharedFileSetPath(char *path, SharedFileSet *fileset, Oid tablespace);
 static void SharedFilePath(char *path, SharedFileSet *fileset, const char *name);
 static Oid	ChooseTablespace(const SharedFileSet *fileset, const char *name);
@@ -76,6 +80,13 @@ SharedFileSetInit(SharedFileSet *fileset, dsm_segment *seg)
 	/* Register our cleanup callback. */
 	if (seg)
 		on_dsm_detach(seg, SharedFileSetOnDetach, PointerGetDatum(fileset));
+	else
+	{
+		if (filesetlist == NIL)
+			on_proc_exit(SharedFileSetOnProcExit, 0);
+
+		filesetlist = lcons((void *) fileset, filesetlist);
+	}
 }
 
 /*
@@ -213,6 +224,57 @@ SharedFileSetOnDetach(dsm_segment *segment, Datum datum)
 		SharedFileSetDeleteAll(fileset);
 }
 
+/*
+ * Callback function that will be invoked on the process exit.  This will
+ * process the list if all the sharedfileset registered and delete the
+ * underlying files.
+ */
+static void
+SharedFileSetOnProcExit(int status, Datum arg)
+{
+	ListCell *l;
+
+	/* Loop over all the pending shared fileset entry */
+	foreach (l, filesetlist)
+	{
+		SharedFileSet *fileset = (SharedFileSet *) lfirst(l);
+		SharedFileSetDeleteAll(fileset);
+	}
+	filesetlist = NIL;
+}
+
+/*
+ * Unregister the shared fileset entry, registered for cleanup on proc exit.
+ */
+void
+SharedFileSetUnregister(SharedFileSet *input_fileset)
+{
+	bool	found = false;
+	ListCell *l;
+
+	/*
+	 * If the caller is following the dsm based cleanup then we don't
+	 * maintain the filesetlist so return.
+	 */
+	if (filesetlist == NULL)
+		return;
+
+	/* Loop over all the shared fileset entries to find the input fileset */
+	foreach (l, filesetlist)
+	{
+		SharedFileSet *fileset = (SharedFileSet *) lfirst(l);
+
+		/* Remove the entry from the list */
+		if (input_fileset == fileset)
+		{
+			filesetlist = list_delete_cell(filesetlist, l);
+			found = true;
+			break;
+		}
+	}
+	Assert(found);
+}
+
 /*
  * Build the path for the directory holding the files backing a SharedFileSet
  * in a given tablespace.
diff --git a/src/include/storage/sharedfileset.h b/src/include/storage/sharedfileset.h
index b2f4ba4bd8..d5edb600af 100644
--- a/src/include/storage/sharedfileset.h
+++ b/src/include/storage/sharedfileset.h
@@ -42,5 +42,6 @@ extern File SharedFileSetOpen(SharedFileSet *fileset, const char *name,
 extern bool SharedFileSetDelete(SharedFileSet *fileset, const char *name,
 								bool error_on_failure);
 extern void SharedFileSetDeleteAll(SharedFileSet *fileset);
+extern void SharedFileSetUnregister(SharedFileSet *input_fileset);
 
 #endif