From a94184fff8801825f0bb6a3d990574b44a02fc1f Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onder@citusdata.com>
Date: Fri, 14 Sep 2018 14:51:35 +0300
Subject: [PATCH] Prevent overflow of memory accesses during deadlock detection

In the distributed deadlock detection design, we concluded that prepared transactions
cannot be part of a distributed deadlock. The idea is that (a) when the transaction
is prepared it already acquires all the locks, so cannot be part of a deadlock
(b) even if some other processes blocked on the prepared transaction,  prepared transactions
 would eventually be committed (or rollbacked) and the system will continue operating.

With the above in mind, we probably had a mistake in terms of memory allocations. For each
backend initialized, we keep a `BackendData` struct. The bug we've introduced is that, we
assumed there would only be `MaxBackend` number of backends. However, `MaxBackends` doesn't
include the prepared transactions and axuliary processes. When you check Postgres' InitProcGlobal`
you'd see that `TotalProcs = MaxBackends + NUM_AUXILIARY_PROCS + max_prepared_xacts;`

This commit aligns with total procs processed with that.
---
 src/backend/distributed/transaction/backend_data.c |  8 +++++---
 src/backend/distributed/transaction/lock_graph.c   | 10 +++++-----
 src/include/distributed/backend_data.h             |  4 ++++
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/backend/distributed/transaction/backend_data.c b/src/backend/distributed/transaction/backend_data.c
index b78e694d9..1f30878b9 100644
--- a/src/backend/distributed/transaction/backend_data.c
+++ b/src/backend/distributed/transaction/backend_data.c
@@ -541,9 +541,11 @@ BackendManagementShmemInit(void)
 
 		/*
 		 * We need to init per backend's spinlock before any backend
-		 * starts its execution.
+		 * starts its execution. Note that we initialize TotalProcs (e.g., not
+		 * MaxBackends) since some of the blocking processes could be prepared
+		 * transactions, which aren't covered by MaxBackends.
 		 */
-		for (backendIndex = 0; backendIndex < MaxBackends; ++backendIndex)
+		for (backendIndex = 0; backendIndex < TotalProcs; ++backendIndex)
 		{
 			SpinLockInit(&backendManagementShmemData->backends[backendIndex].mutex);
 		}
@@ -568,7 +570,7 @@ BackendManagementShmemSize(void)
 	Size size = 0;
 
 	size = add_size(size, sizeof(BackendManagementShmemData));
-	size = add_size(size, mul_size(sizeof(BackendData), MaxBackends));
+	size = add_size(size, mul_size(sizeof(BackendData), TotalProcs));
 
 	return size;
 }
diff --git a/src/backend/distributed/transaction/lock_graph.c b/src/backend/distributed/transaction/lock_graph.c
index d3b51b3e5..734b3acac 100644
--- a/src/backend/distributed/transaction/lock_graph.c
+++ b/src/backend/distributed/transaction/lock_graph.c
@@ -398,12 +398,12 @@ BuildLocalWaitGraph(void)
 	 */
 	waitGraph = (WaitGraph *) palloc0(sizeof(WaitGraph));
 	waitGraph->localNodeId = GetLocalGroupId();
-	waitGraph->allocatedSize = MaxBackends * 3;
+	waitGraph->allocatedSize = TotalProcs * 3;
 	waitGraph->edgeCount = 0;
 	waitGraph->edges = (WaitEdge *) palloc(waitGraph->allocatedSize * sizeof(WaitEdge));
 
-	remaining.procs = (PGPROC **) palloc(sizeof(PGPROC *) * MaxBackends);
-	remaining.procAdded = (bool *) palloc0(sizeof(bool *) * MaxBackends);
+	remaining.procs = (PGPROC **) palloc(sizeof(PGPROC *) * TotalProcs);
+	remaining.procAdded = (bool *) palloc0(sizeof(bool *) * TotalProcs);
 	remaining.procCount = 0;
 
 	LockLockData();
@@ -416,7 +416,7 @@ BuildLocalWaitGraph(void)
 	 */
 
 	/* build list of starting procs */
-	for (curBackend = 0; curBackend < MaxBackends; curBackend++)
+	for (curBackend = 0; curBackend < TotalProcs; curBackend++)
 	{
 		PGPROC *currentProc = &ProcGlobal->allProcs[curBackend];
 		BackendData currentBackendData;
@@ -762,7 +762,7 @@ AddProcToVisit(PROCStack *remaining, PGPROC *proc)
 		return;
 	}
 
-	Assert(remaining->procCount < MaxBackends);
+	Assert(remaining->procCount < TotalProcs);
 
 	remaining->procs[remaining->procCount++] = proc;
 	remaining->procAdded[proc->pgprocno] = true;
diff --git a/src/include/distributed/backend_data.h b/src/include/distributed/backend_data.h
index 5388f8d42..8a4d64080 100644
--- a/src/include/distributed/backend_data.h
+++ b/src/include/distributed/backend_data.h
@@ -13,6 +13,7 @@
 #define BACKEND_DATA_H
 
 
+#include "access/twophase.h"
 #include "datatype/timestamp.h"
 #include "distributed/transaction_identifier.h"
 #include "nodes/pg_list.h"
@@ -21,6 +22,9 @@
 #include "storage/s_lock.h"
 
 
+#define TotalProcs (MaxBackends + NUM_AUXILIARY_PROCS + max_prepared_xacts)
+
+
 /*
  * Each backend's active distributed transaction information is tracked via
  * BackendData in shared memory.