diff --git a/src/backend/distributed/transaction/backend_data.c b/src/backend/distributed/transaction/backend_data.c index 7f89b805d..119d5d1c1 100644 --- a/src/backend/distributed/transaction/backend_data.c +++ b/src/backend/distributed/transaction/backend_data.c @@ -25,6 +25,7 @@ #include "distributed/remote_commands.h" #include "distributed/transaction_identifier.h" #include "nodes/execnodes.h" +#include "postmaster/autovacuum.h" /* to access autovacuum_max_workers */ #include "storage/ipc.h" #include "storage/lmgr.h" #include "storage/lwlock.h" @@ -512,6 +513,7 @@ BackendManagementShmemInit(void) if (!alreadyInitialized) { int backendIndex = 0; + int totalProcs = 0; char *trancheName = "Backend Management Tranche"; #if (PG_VERSION_NUM >= 100000) @@ -557,7 +559,8 @@ BackendManagementShmemInit(void) * We also initiate initiatorNodeIdentifier to -1, which can never be * used as a node id. */ - for (backendIndex = 0; backendIndex < TotalProcs; ++backendIndex) + totalProcs = TotalProcCount(); + for (backendIndex = 0; backendIndex < totalProcs; ++backendIndex) { backendManagementShmemData->backends[backendIndex].citusBackend. initiatorNodeIdentifier = -1; @@ -582,14 +585,62 @@ static size_t BackendManagementShmemSize(void) { Size size = 0; + int totalProcs = TotalProcCount(); size = add_size(size, sizeof(BackendManagementShmemData)); - size = add_size(size, mul_size(sizeof(BackendData), TotalProcs)); + size = add_size(size, mul_size(sizeof(BackendData), totalProcs)); return size; } +/* + * TotalProcCount returns the total processes that could run via the current + * postgres server. See the details in the function comments. + * + * There is one thing we should warn the readers. Citus enforces to be loaded + * as the first extension in shared_preload_libraries. However, if any other + * extension overrides MaxConnections, autovacuum_max_workers or + * max_worker_processes, our reasoning in this function may not work as expected. + * Given that it is not a usual pattern for extension, we consider Citus' behaviour + * good enough for now. + */ +int +TotalProcCount(void) +{ + int maxBackends = 0; + int totalProcs = 0; + +#ifdef WIN32 + + /* autovacuum_max_workers is not PGDLLIMPORT, so use a high estimate for windows */ + int estimatedMaxAutovacuumWorkers = 30; + maxBackends = + MaxConnections + estimatedMaxAutovacuumWorkers + 1 + max_worker_processes; +#else + + /* + * We're simply imitating Postgrsql's InitializeMaxBackends(). Given that all + * the items used here PGC_POSTMASTER, should be safe to access them + * anytime during the execution even before InitializeMaxBackends() is called. + */ + maxBackends = MaxConnections + autovacuum_max_workers + 1 + max_worker_processes; +#endif + + /* + * We prefer to maintain space for auxiliary procs or preperad transactions in + * the backend space because they could be blocking processes and our current + * implementation of distributed deadlock detection could process them + * as a regular backend. In the future, we could consider chaning deadlock + * detection algorithm to ignore auxiliary procs or preperad transactions and + * save same space. + */ + totalProcs = maxBackends + NUM_AUXILIARY_PROCS + max_prepared_xacts; + + return totalProcs; +} + + /* * InitializeBackendData initialises MyBackendData to the shared memory segment * belonging to the current backend. diff --git a/src/backend/distributed/transaction/lock_graph.c b/src/backend/distributed/transaction/lock_graph.c index 734b3acac..cee26d0fb 100644 --- a/src/backend/distributed/transaction/lock_graph.c +++ b/src/backend/distributed/transaction/lock_graph.c @@ -389,6 +389,7 @@ BuildLocalWaitGraph(void) WaitGraph *waitGraph = NULL; int curBackend = 0; PROCStack remaining; + int totalProcs = TotalProcCount(); /* * Try hard to avoid allocations while holding lock. Thus we pre-allocate @@ -398,12 +399,12 @@ BuildLocalWaitGraph(void) */ waitGraph = (WaitGraph *) palloc0(sizeof(WaitGraph)); waitGraph->localNodeId = GetLocalGroupId(); - waitGraph->allocatedSize = TotalProcs * 3; + waitGraph->allocatedSize = totalProcs * 3; waitGraph->edgeCount = 0; waitGraph->edges = (WaitEdge *) palloc(waitGraph->allocatedSize * sizeof(WaitEdge)); - remaining.procs = (PGPROC **) palloc(sizeof(PGPROC *) * TotalProcs); - remaining.procAdded = (bool *) palloc0(sizeof(bool *) * TotalProcs); + remaining.procs = (PGPROC **) palloc(sizeof(PGPROC *) * totalProcs); + remaining.procAdded = (bool *) palloc0(sizeof(bool *) * totalProcs); remaining.procCount = 0; LockLockData(); @@ -416,7 +417,7 @@ BuildLocalWaitGraph(void) */ /* build list of starting procs */ - for (curBackend = 0; curBackend < TotalProcs; curBackend++) + for (curBackend = 0; curBackend < totalProcs; curBackend++) { PGPROC *currentProc = &ProcGlobal->allProcs[curBackend]; BackendData currentBackendData; @@ -762,7 +763,7 @@ AddProcToVisit(PROCStack *remaining, PGPROC *proc) return; } - Assert(remaining->procCount < TotalProcs); + Assert(remaining->procCount < TotalProcCount()); 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 4a6eae152..f52b69125 100644 --- a/src/include/distributed/backend_data.h +++ b/src/include/distributed/backend_data.h @@ -22,9 +22,6 @@ #include "storage/s_lock.h" -#define TotalProcs (MaxBackends + NUM_AUXILIARY_PROCS + max_prepared_xacts) - - /* * CitusInitiatedBackend keeps some information about the backends that are * initiated by Citus. @@ -58,6 +55,7 @@ typedef struct BackendData extern void InitializeBackendManagement(void); +extern int TotalProcCount(void); extern void InitializeBackendData(void); extern void LockBackendSharedMemory(LWLockMode lockMode); extern void UnlockBackendSharedMemory(void);