Cancel bgwpool workers on receiver shutdown

Author: kelvich

src/bgwpool.c

@@ -12,6 +12,7 @@
 #include "utils/portal.h"
 #include "tcop/pquery.h"
 #include "utils/guc.h"
+#include "tcop/tcopprot.h"
 
 #include "bgwpool.h"
 #include "mm.h"
@@ -25,11 +26,17 @@ static BgwPool* MtmPool;
 void BgwPoolStaticWorkerMainLoop(Datum arg);
 void BgwPoolDynamicWorkerMainLoop(Datum arg);
 
-static void BgwShutdownWorker(int sig)
+static void
+BgwShutdownHandler(int sig)
 {
-	if (MtmPool) { 
-		BgwPoolStop(MtmPool);
-	}
+	Assert(MtmPool != NULL);
+	BgwPoolStop(MtmPool);
+
+	/*
+	 * set ProcDiePending for cases when we are waiting on latch somewhere
+	 * deep inside our execute() function.
+	 */
+	die(sig);
 }
 
 static void BgwPoolMainLoop(BgwPool* pool)
@@ -44,10 +51,9 @@ static void BgwPoolMainLoop(BgwPool* pool)
 	MtmIsLogicalReceiver = true;
 	MtmPool = pool;
 
-	// XXX: fix that
-	pqsignal(SIGINT, BgwShutdownWorker);
-	pqsignal(SIGQUIT, BgwShutdownWorker);
-	pqsignal(SIGTERM, BgwShutdownWorker);
+	pqsignal(SIGINT, StatementCancelHandler);
+	pqsignal(SIGQUIT, BgwShutdownHandler);
+	pqsignal(SIGTERM, BgwShutdownHandler);
 	pqsignal(SIGHUP, PostgresSigHupHandler);
 
     BackgroundWorkerUnblockSignals();
@@ -93,8 +99,13 @@ static void BgwPoolMainLoop(BgwPool* pool)
 			pool->lastPeakTime = 0;
         }
         SpinLockRelease(&pool->lock);
+
+		/* Ignore cancel that arrived before we started current command */
+		QueryCancelPending = false;
+
 		pool->executor(work, size, NULL);
-        pfree(work);
+		pfree(work);
+
         SpinLockAcquire(&pool->lock);
         pool->active -= 1;
 		pool->lastPeakTime = 0;
@@ -107,6 +118,8 @@ static void BgwPoolMainLoop(BgwPool* pool)
 void BgwPoolInit(BgwPool* pool, BgwPoolExecutor executor, char const* dbname,  char const* dbuser, size_t queueSize, size_t nWorkers)
 {
 	MtmPool = pool;
+
+	pool->bgwhandles = (BackgroundWorkerHandle **) ShmemAlloc(MtmMaxWorkers * sizeof(BackgroundWorkerHandle *));
     pool->queue = (char*)ShmemAlloc(queueSize);
 	if (pool->queue == NULL) { 
 		elog(PANIC, "Failed to allocate memory for background workers pool: %lld bytes requested", (long64)queueSize);
@@ -148,7 +161,6 @@ void BgwPoolDynamicWorkerMainLoop(Datum arg)
 
 void BgwPoolStart(BgwPool* pool, char *poolName)
 {
-    int i;
 	BackgroundWorker worker;
 
 	MemSet(&worker, 0, sizeof(BackgroundWorker));
@@ -159,13 +171,6 @@ void BgwPoolStart(BgwPool* pool, char *poolName)
 	worker.bgw_restart_time = MULTIMASTER_BGW_RESTART_TIMEOUT;
 
 	strncpy(pool->poolName, poolName, MAX_NAME_LEN);
-
-	for (i = 0; i < pool->nWorkers; i++)
-	{
-		snprintf(worker.bgw_name, BGW_MAXLEN, "%s_worker_%d", pool->poolName, i+1);
-		worker.bgw_main_arg = PointerGetDatum(pool);
-        RegisterBackgroundWorker(&worker);
-    }
 }
 
 size_t BgwPoolGetQueueSize(BgwPool* pool)
@@ -180,25 +185,29 @@ size_t BgwPoolGetQueueSize(BgwPool* pool)
 
 static void BgwStartExtraWorker(BgwPool* pool)
 {
-	if (pool->nWorkers < MtmMaxWorkers) { 
-		timestamp_t now = MtmGetSystemTime();
-		/*if (pool->lastDynamicWorkerStartTime + MULTIMASTER_BGW_RESTART_TIMEOUT*USECS_PER_SEC < now)*/
-		{ 
-			BackgroundWorker worker;
-			BackgroundWorkerHandle* handle;
-			MemSet(&worker, 0, sizeof(BackgroundWorker));
-			worker.bgw_flags = BGWORKER_SHMEM_ACCESS |  BGWORKER_BACKEND_DATABASE_CONNECTION;
-			worker.bgw_start_time = BgWorkerStart_ConsistentState;
-			sprintf(worker.bgw_library_name, "multimaster");
-			sprintf(worker.bgw_function_name, "BgwPoolDynamicWorkerMainLoop");
-			worker.bgw_restart_time = MULTIMASTER_BGW_RESTART_TIMEOUT;
-			snprintf(worker.bgw_name, BGW_MAXLEN, "%s-dynworker-%d", pool->poolName, (int)++pool->nWorkers);
-			worker.bgw_main_arg = PointerGetDatum(pool);
-			pool->lastDynamicWorkerStartTime = now;
-			if (!RegisterDynamicBackgroundWorker(&worker, &handle)) { 
-				elog(WARNING, "Failed to start dynamic background worker");
-			}
-		}
+	BackgroundWorker worker;
+	BackgroundWorkerHandle* handle;
+
+	if (pool->nWorkers >= MtmMaxWorkers)
+		return;
+
+	MemSet(&worker, 0, sizeof(BackgroundWorker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |  BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_ConsistentState;
+	worker.bgw_restart_time = MULTIMASTER_BGW_RESTART_TIMEOUT;
+	worker.bgw_main_arg = PointerGetDatum(pool);
+	sprintf(worker.bgw_library_name, "multimaster");
+	sprintf(worker.bgw_function_name, "BgwPoolDynamicWorkerMainLoop");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "%s-dynworker-%d", pool->poolName, (int) pool->nWorkers + 1);
+
+	pool->lastDynamicWorkerStartTime = MtmGetSystemTime();
+	if (RegisterDynamicBackgroundWorker(&worker, &handle))
+	{
+		pool->bgwhandles[pool->nWorkers++] = handle;
+	}
+	else
+	{
+		elog(WARNING, "Failed to start dynamic background worker");
 	}
 }
 
@@ -253,9 +262,29 @@ void BgwPoolExecute(BgwPool* pool, void* work, size_t size)
 
 void BgwPoolStop(BgwPool* pool)
 {
-    // SpinLockAcquire(&pool->lock);
 	pool->shutdown = true;
-    // SpinLockRelease(&pool->lock);            
 	PGSemaphoreUnlock(pool->available);
 	PGSemaphoreUnlock(pool->overflow);
 }
+
+/*
+ * Tell our lads to cancel currently active transactions.
+ */
+void
+BgwPoolCancel(BgwPool* pool)
+{
+	int		i;
+
+	for (i = 0; i < pool->nWorkers; i++)
+	{
+		BgwHandleStatus		status;
+		pid_t				pid;
+
+		status = GetBackgroundWorkerPid(pool->bgwhandles[i], &pid);
+		if (status == BGWH_STARTED)
+		{
+			Assert(pid > 0);
+			kill(pid, SIGINT);
+		}
+	}
+}

src/include/bgwpool.h

@@ -4,6 +4,7 @@
 #include "storage/s_lock.h"
 #include "storage/spin.h"
 #include "storage/pg_sema.h"
+#include "postmaster/bgworker.h"
 #include "bkb.h" // XXX
 
 #include "mm.h"
@@ -43,6 +44,8 @@ typedef struct
     char   dbname[MAX_DBNAME_LEN];
 	char   dbuser[MAX_DBUSER_LEN];
     char*  queue;
+
+	BackgroundWorkerHandle **bgwhandles;
 } BgwPool;
 
 typedef BgwPool*(*BgwPoolConstructor)(void);
@@ -58,4 +61,7 @@ extern size_t BgwPoolGetQueueSize(BgwPool* pool);
 extern timestamp_t BgwGetLastPeekTime(BgwPool* pool);
 
 extern void BgwPoolStop(BgwPool* pool);
+
+extern void BgwPoolCancel(BgwPool* pool);
+
 #endif

src/pglogical_receiver.c

@@ -470,6 +470,7 @@ pglogical_receiver_main(Datum main_arg)
 	/* Register functions for SIGTERM/SIGHUP management */
 	pqsignal(SIGHUP, receiver_raw_sighup);
 	// pqsignal(SIGTERM, receiver_raw_sigterm);
+	// XXX: also need to unlock overflow semaphore (or rewrite it on latches)
 	pqsignal(SIGTERM, die);
 
 	MtmCreateSpillDirectory(nodeId);
@@ -534,16 +535,7 @@ pglogical_receiver_main(Datum main_arg)
 		mtm_log(MtmReceiverMode, "WalReceiver to %d starts in %s mode",
 				nodeId, MtmReplicationModeName[mode]);
 
-		// if (mode == REPLMODE_RECOVERY)
-		// 	synchronous_twophase = false;
-		// else
-		// 	synchronous_twophase = true;
-		// synchronous_twophase = true;
-
-		if (mode == REPLMODE_EXIT)
-		{
-			break;
-		}
+		Assert(mode == REPLMODE_RECOVERY || mode == REPLMODE_RECOVERED);
 		timeline = Mtm->nodes[nodeId-1].timeline;
 		count = Mtm->recoveryCount;
 
@@ -648,16 +640,17 @@ pglogical_receiver_main(Datum main_arg)
 			if (Mtm->status == MTM_DISABLED || (Mtm->status == MTM_RECOVERY && Mtm->recoverySlot != nodeId))
 			{
 				ereport(LOG, (MTM_ERRMSG("%s: restart WAL receiver because node was switched to %s mode", worker_proc, MtmNodeStatusMnem[Mtm->status])));
-				break;
+				goto OnError;
 			}
+
 			if (count != Mtm->recoveryCount) {
 				ereport(LOG, (MTM_ERRMSG("%s: restart WAL receiver because node was recovered", worker_proc)));
-				break;
+				goto OnError;
 			}
 
 			if (timeline != Mtm->nodes[nodeId-1].timeline) {
 				ereport(LOG, (MTM_ERRMSG("%s: restart WAL receiver because node %d timeline is changed", worker_proc, nodeId)));
-				break;
+				goto OnError;
 			}
 
 			/*
@@ -911,14 +904,21 @@ pglogical_receiver_main(Datum main_arg)
 				goto OnError;
 			}
 		}
-		PQfinish(conn);
-		continue;
+		Assert(false);
+
 
-	  OnError:
+	OnError:
 		ByteBufferReset(&buf);
 		PQfinish(conn);
 		if (Mtm->recoverySlot == nodeId)
 			Mtm->recoverySlot = 0;
+
+		/*
+		 * Some of the workers may stuck on lock and survive until node will come
+		 * back and prepare stuck transaction when it was aborted long time ago.
+		 * Force all workers to cancel stmt to ensure this will not happen.
+		 */
+		BgwPoolCancel(&Mtm->nodes[nodeId - 1].pool);
 		MtmSleep(RECEIVER_SUSPEND_TIMEOUT);
 	}
 	ByteBufferFree(&buf);