From f0fd108b415abf7bc425d6ca94e8f2f5bf3a827f Mon Sep 17 00:00:00 2001
From: Amul Sul <sulamul@gmail.com>
Date: Fri, 6 Jan 2017 18:13:13 +0530
Subject: [PATCH 1/2] Changes to Peter Eisentraut's bgsession v2 patch

---
 src/backend/tcop/bgsession.c     | 26 +++++++++-----------------
 src/include/tcop/bgsession.h     | 19 +++++++++++++++----
 src/pl/plpython/plpy_bgsession.c |  2 +-
 3 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/src/backend/tcop/bgsession.c b/src/backend/tcop/bgsession.c
index 486819b..26dee53 100644
--- a/src/backend/tcop/bgsession.c
+++ b/src/backend/tcop/bgsession.c
@@ -60,14 +60,11 @@
 #include "nodes/pg_list.h"
 #include "pgstat.h"
 #include "postmaster/bgworker.h"
-#include "storage/dsm.h"
-#include "storage/shm_mq.h"
 #include "storage/shm_toc.h"
 #include "tcop/bgsession.h"
 #include "tcop/tcopprot.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
-#include "utils/resowner.h"
 
 /* Table-of-contents constants for our dynamic shared memory segment. */
 #define BGSESSION_MAGIC					0x50674267
@@ -89,16 +86,6 @@ typedef struct background_session_fixed_data
 	int sec_context;
 } background_session_fixed_data;
 
-struct BackgroundSession
-{
-	ResourceOwner resowner;
-	dsm_segment *seg;
-	BackgroundWorkerHandle *worker_handle;
-	shm_mq_handle *command_qh;
-	shm_mq_handle *response_qh;
-	int		transaction_status;
-};
-
 struct BackgroundSessionPreparedStatement
 {
 	BackgroundSession *session;
@@ -117,11 +104,10 @@ static void invalid_protocol_message(char msgtype) pg_attribute_noreturn();
 
 
 BackgroundSession *
-BackgroundSessionStart(void)
+BackgroundSessionStart(MemoryContext context)
 {
 	ResourceOwner oldowner;
 	BackgroundWorker worker;
-	pid_t		pid;
 	BackgroundSession *session;
 	shm_toc_estimator e;
 	Size		segsize;
@@ -135,8 +121,9 @@ BackgroundSessionStart(void)
 	BgwHandleStatus bgwstatus;
 	StringInfoData msg;
 	char		msgtype;
+	MemoryContext   oldcontext;
 
-	session = palloc(sizeof(*session));
+	session = MemoryContextAlloc(context, sizeof(*session));;
 
 	session->resowner = ResourceOwnerCreate(NULL, "background session");
 
@@ -188,8 +175,10 @@ BackgroundSessionStart(void)
 	shm_toc_insert(toc, BGSESSION_KEY_RESPONSE_QUEUE, response_mq);
 	shm_mq_set_receiver(response_mq, MyProc);
 
+	oldcontext = MemoryContextSwitchTo(context);
 	session->command_qh = shm_mq_attach(command_mq, seg, NULL);
 	session->response_qh = shm_mq_attach(response_mq, seg, NULL);
+	MemoryContextSwitchTo(oldcontext);
 
 	worker.bgw_flags =
 		BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION;
@@ -200,16 +189,19 @@ BackgroundSessionStart(void)
 	worker.bgw_main_arg = UInt32GetDatum(dsm_segment_handle(seg));
 	worker.bgw_notify_pid = MyProcPid;
 
+	oldcontext = MemoryContextSwitchTo(context);
 	if (!RegisterDynamicBackgroundWorker(&worker, &session->worker_handle))
 		ereport(ERROR,
 				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
 				 errmsg("could not register background process"),
 				 errhint("You might need to increase max_worker_processes.")));
+	MemoryContextSwitchTo(oldcontext);
 
 	shm_mq_set_handle(session->command_qh, session->worker_handle);
 	shm_mq_set_handle(session->response_qh, session->worker_handle);
 
-	bgwstatus = WaitForBackgroundWorkerStartup(session->worker_handle, &pid);
+	bgwstatus = WaitForBackgroundWorkerStartup(session->worker_handle,
+											   &session->pid);
 	if (bgwstatus != BGWH_STARTED)
 		ereport(ERROR,
 				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
diff --git a/src/include/tcop/bgsession.h b/src/include/tcop/bgsession.h
index 71415a6..65dfb56 100644
--- a/src/include/tcop/bgsession.h
+++ b/src/include/tcop/bgsession.h
@@ -3,13 +3,24 @@
 
 #include "access/tupdesc.h"
 #include "nodes/pg_list.h"
-
-struct BackgroundSession;
-typedef struct BackgroundSession BackgroundSession;
+#include "storage/dsm.h"
+#include "storage/shm_mq.h"
+#include "utils/resowner.h"
 
 struct BackgroundSessionPreparedStatement;
 typedef struct BackgroundSessionPreparedStatement BackgroundSessionPreparedStatement;
 
+typedef struct BackgroundSession
+{
+	pid_t pid;
+	ResourceOwner resowner;
+	dsm_segment *seg;
+	BackgroundWorkerHandle *worker_handle;
+	shm_mq_handle *command_qh;
+	shm_mq_handle *response_qh;
+	int		transaction_status;
+} BackgroundSession;
+
 typedef struct BackgroundSessionResult
 {
 	TupleDesc	tupdesc;
@@ -17,7 +28,7 @@ typedef struct BackgroundSessionResult
 	const char *command;
 } BackgroundSessionResult;
 
-BackgroundSession *BackgroundSessionStart(void);
+BackgroundSession *BackgroundSessionStart(MemoryContext oldcontext);
 void BackgroundSessionEnd(BackgroundSession *session);
 
 void BackgroundSessionSend(BackgroundSession *session, const char *sql);
diff --git a/src/pl/plpython/plpy_bgsession.c b/src/pl/plpython/plpy_bgsession.c
index 68f7207..7e6e906 100644
--- a/src/pl/plpython/plpy_bgsession.c
+++ b/src/pl/plpython/plpy_bgsession.c
@@ -110,7 +110,7 @@ PLyBackgroundSession_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 	PyObject   *result = type->tp_alloc(type, 0);
 	PLyBackgroundSession_Object *bgsession = (PLyBackgroundSession_Object *) result;
 
-	bgsession->bgsession = BackgroundSessionStart();
+	bgsession->bgsession = BackgroundSessionStart(CurrentMemoryContext);
 
 	return result;
 }
-- 
2.6.2