more protection from QueryCancel in apply worker

Author: kelvich

src/bgwpool.c

@@ -49,7 +49,6 @@ subscription_change_cb(Datum arg, int cacheid, uint32 hashvalue)
 	receiver_mtm_cfg_valid = false;
 }
 
-
 static void
 BgwPoolMainLoop(BgwPool* pool)
 {
@@ -68,10 +67,10 @@ BgwPoolMainLoop(BgwPool* pool)
 	MtmIsLogicalReceiver = true;
 	MtmPool = pool;
 
-	pqsignal(SIGINT, StatementCancelHandler);
+	pqsignal(SIGINT, ApplyCancelHandler);
 	pqsignal(SIGQUIT, BgwShutdownHandler);
 	pqsignal(SIGTERM, BgwShutdownHandler);
-	pqsignal(SIGHUP, PostgresSigHupHandler);
+	pqsignal(SIGHUP, ApplyCancelHandler);
 
 	// XXX: probably we should add static variable that signalizes that
 	// we are between pool->active += 1 and pool->active -= 1, so if
@@ -105,6 +104,7 @@ BgwPoolMainLoop(BgwPool* pool)
 		if (pool->shutdown)
 		{
 			PGSemaphoreUnlock(pool->available);
+			SpinLockRelease(&pool->lock);
 			break;
 		}
 		size = * (int *) &pool->queue[pool->head];
@@ -113,6 +113,7 @@ BgwPoolMainLoop(BgwPool* pool)
 		work = palloc(size);
 		pool->pending -= 1;
 		pool->active += 1;
+
 		if (pool->lastPeakTime == 0 && pool->active == pool->nWorkers && pool->pending != 0)
 			pool->lastPeakTime = MtmGetSystemTime();
 
@@ -152,9 +153,6 @@ BgwPoolMainLoop(BgwPool* pool)
 
 		SpinLockRelease(&pool->lock);
 
-		/* Ignore cancel that arrived before we started current command */
-		QueryCancelPending = false;
-
 		MtmExecutor(work, size, &ctx);
 		pfree(work);
 
@@ -166,7 +164,6 @@ BgwPoolMainLoop(BgwPool* pool)
 		ConditionVariableBroadcast(&pool->syncpoint_cv);
 	}
 
-	SpinLockRelease(&pool->lock);
 	mtm_log(BgwPoolEvent, "Shutdown background worker %d", MyProcPid);
 }
 

src/include/receiver.h

@@ -21,6 +21,7 @@ extern char const *const MtmReplicationModeName[];
 extern BackgroundWorkerHandle *MtmStartReceiver(int nodeId, Oid db_id, Oid user_id, pid_t monitor_pid);
 
 extern void MtmExecutor(void* work, size_t size, MtmReceiverContext *rctx);
+extern void ApplyCancelHandler(SIGNAL_ARGS);
 extern void MtmUpdateLsnMapping(int node_id, XLogRecPtr end_lsn);
 
 extern void MtmBeginSession(int nodeId);

src/pglogical_apply.c

@@ -76,6 +76,7 @@ typedef struct TupleData
 	bool		changed[MaxTupleAttributeNumber];
 } TupleData;
 
+static bool query_cancel_allowed;
 
 static Relation read_rel(StringInfo s, LOCKMODE mode);
 static void read_tuple_parts(StringInfo s, Relation rel, TupleData *tup);
@@ -121,6 +122,25 @@ MtmHandleApplyError(void)
 	FreeErrorData(edata);
 }
 
+void
+ApplyCancelHandler(SIGNAL_ARGS)
+{
+	int			save_errno = errno;
+
+	/*
+	 * Don't joggle the elbow of proc_exit
+	 */
+	if (!proc_exit_inprogress && query_cancel_allowed)
+	{
+		InterruptPending = true;
+		QueryCancelPending = true;
+	}
+
+	/* If we're still here, waken anything waiting on the process latch */
+	SetLatch(MyLatch);
+
+	errno = save_errno;
+}
 
 bool find_heap_tuple(TupleData *tup, Relation rel, TupleTableSlot *slot, bool lock)
 {
@@ -483,6 +503,10 @@ process_remote_begin(StringInfo s, GlobalTransactionId *gtid)
 	gtid->node = pq_getmsgint(s, 4);
 	gtid->xid = pq_getmsgint64(s);
 
+	InterruptPending = false;
+	QueryCancelPending = false;
+	query_cancel_allowed = true;
+
 	// XXX: get rid of MtmReplicationNodeId
 	MtmReplicationNodeId = gtid->node;
 
@@ -929,6 +953,7 @@ process_remote_commit(StringInfo in, GlobalTransactionId *current_gtid, MtmRecei
 	{
 		case PGLOGICAL_PRECOMMIT_PREPARED:
 		{
+			Assert(!query_cancel_allowed);
 			strncpy(gid, pq_getmsgstring(in), sizeof gid);
 			MtmBeginSession(origin_node);
 
@@ -979,6 +1004,10 @@ process_remote_commit(StringInfo in, GlobalTransactionId *current_gtid, MtmRecei
 
 			CommitTransactionCommand();
 
+			InterruptPending = false;
+			QueryCancelPending = false;
+			query_cancel_allowed = false;
+
 			if (receiver_ctx->parallel_allowed)
 			{
 				mtm_send_xid_reply(current_gtid->xid, origin_node,
@@ -1005,6 +1034,7 @@ process_remote_commit(StringInfo in, GlobalTransactionId *current_gtid, MtmRecei
 		}
 		case PGLOGICAL_COMMIT_PREPARED:
 		{
+			Assert(!query_cancel_allowed);
 			pq_getmsgint64(in); /* csn */
 			strncpy(gid, pq_getmsgstring(in), sizeof gid);
 			StartTransactionCommand();
@@ -1621,6 +1651,7 @@ MtmExecutor(void* work, size_t size, MtmReceiverContext *receiver_ctx)
     }
     PG_CATCH();
     {
+		query_cancel_allowed = false;
 
 		old_context = MemoryContextSwitchTo(MtmApplyContext);
 		MtmHandleApplyError();