From 5880ecc6802cc6c901ed308b33abc4d64a8bee91 Mon Sep 17 00:00:00 2001
From: ivyazmitinov <ivanvyazmitinov@gmail.com>
Date: Wed, 18 Oct 2023 16:18:11 +0700
Subject: [PATCH] Introduce new GUC

---
 .../connection/shared_connection_stats.c      |  4 +--
 .../operations/worker_node_manager.c          |  2 +-
 src/backend/distributed/shared_library_init.c | 28 +++++++++++++++++++
 .../distributed/shared_connection_stats.h     |  2 +-
 src/include/distributed/worker_manager.h      |  2 +-
 5 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/backend/distributed/connection/shared_connection_stats.c b/src/backend/distributed/connection/shared_connection_stats.c
index 55df40680..249dfbd4e 100644
--- a/src/backend/distributed/connection/shared_connection_stats.c
+++ b/src/backend/distributed/connection/shared_connection_stats.c
@@ -672,7 +672,7 @@ SharedConnectionStatsShmemSize(void)
 
     size = add_size(size, workerNodeConnHashSize);
 
-    Size workerNodeDatabaseConnSize = hash_estimate_size(DatabasesPerWorker,
+    Size workerNodeDatabaseConnSize = hash_estimate_size(MaxWorkerNodesTracked * MaxDatabasesPerWorkerNodesTracked,
                                                          sizeof(SharedWorkerNodeDatabaseConnStatsHashEntry));
 
     size = add_size(size, workerNodeDatabaseConnSize);
@@ -741,7 +741,7 @@ SharedConnectionStatsShmemInit(void)
     sharedWorkerNodeDatabaseConnStatsHashInfo.hash = SharedWorkerNodeDatabaseHashHash;
     sharedWorkerNodeDatabaseConnStatsHashInfo.match = SharedWorkerNodeDatabaseHashCompare;
 
-    int sharedWorkerNodeDatabaseConnStatsHashSize = MaxWorkerNodesTracked * DatabasesPerWorker;
+    int sharedWorkerNodeDatabaseConnStatsHashSize = MaxWorkerNodesTracked * MaxDatabasesPerWorkerNodesTracked;
     SharedWorkerNodeDatabaseConnStatsHash =
             ShmemInitHash("Shared Conn Per Database. Stats Hash",
                           sharedWorkerNodeDatabaseConnStatsHashSize,
diff --git a/src/backend/distributed/operations/worker_node_manager.c b/src/backend/distributed/operations/worker_node_manager.c
index 32eb28b94..53109e324 100644
--- a/src/backend/distributed/operations/worker_node_manager.c
+++ b/src/backend/distributed/operations/worker_node_manager.c
@@ -39,7 +39,7 @@
 /* Config variables managed via guc.c */
 char *WorkerListFileName;
 int MaxWorkerNodesTracked = 2048;    /* determines worker node hash table size */
-int DatabasesPerWorker = 1; /* determine database per worker hash table size */
+int MaxDatabasesPerWorkerNodesTracked = 1; /* determine database per worker hash table size */
 
 
 /* Local functions forward declarations */
diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c
index 45e212e8b..97f78defd 100644
--- a/src/backend/distributed/shared_library_init.c
+++ b/src/backend/distributed/shared_library_init.c
@@ -1993,6 +1993,20 @@ RegisterCitusConfigVariables(void)
 		GUC_SUPERUSER_ONLY,
 		NULL, NULL, MaxSharedPoolSizeGucShowHook);
 
+    DefineCustomRealVariable(
+            "citus.shared_pool_size_maintenance_quota",
+            gettext_noop("Sets the maximum number of connections allowed per worker node "
+                         "across all the backends from this node. Setting to -1 disables "
+                         "connections throttling. Setting to 0 makes it auto-adjust, meaning "
+                         "equal to max_connections on the coordinator."),
+            gettext_noop("As a rule of thumb, the value should be at most equal to the "
+                         "max_connections on the remote nodes."),
+            &SharedPoolSizeMaintenanceQuota,
+            0.1, 0, 1,
+            PGC_SIGHUP,
+            GUC_SUPERUSER_ONLY,
+            NULL, NULL, MaxSharedPoolSizeGucShowHook);
+
 	DefineCustomIntVariable(
 		"citus.max_worker_nodes_tracked",
 		gettext_noop("Sets the maximum number of worker nodes that are tracked."),
@@ -2012,6 +2026,20 @@ RegisterCitusConfigVariables(void)
 		GUC_STANDARD,
 		NULL, NULL, NULL);
 
+    DefineCustomIntVariable(
+            "citus.max_databases_per_worker_tracked",
+            gettext_noop("Sets the amount of databases per worker tracked."),
+            gettext_noop(
+                    "This configuration value compliments the citus.max_worker_nodes_tracked."
+                    "It should be used when there are more then one database with Citus in cluster,"
+                    "and, effectively, limits the size of the hash table with connections per worker + database."
+                    "Currently, it does not affect the connection management logic and serves only statistical purposes."),
+            &MaxDatabasesPerWorkerNodesTracked,
+            1, 1, INT_MAX,
+            PGC_POSTMASTER,
+            GUC_STANDARD,
+            NULL, NULL, NULL);
+
 	DefineCustomIntVariable(
 		"citus.metadata_sync_interval",
 		gettext_noop("Sets the time to wait between metadata syncs."),
diff --git a/src/include/distributed/shared_connection_stats.h b/src/include/distributed/shared_connection_stats.h
index 817f80b54..5a0be6ea2 100644
--- a/src/include/distributed/shared_connection_stats.h
+++ b/src/include/distributed/shared_connection_stats.h
@@ -25,7 +25,7 @@ enum SharedPoolCounterMode
 };
 
 extern int MaxSharedPoolSize;
-extern float MaintenanceSharedPoolSizePercent;
+extern double SharedPoolSizeMaintenanceQuota;
 extern int LocalSharedPoolSize;
 extern int MaxClientConnections;
 
diff --git a/src/include/distributed/worker_manager.h b/src/include/distributed/worker_manager.h
index aad2f157c..4e1287e07 100644
--- a/src/include/distributed/worker_manager.h
+++ b/src/include/distributed/worker_manager.h
@@ -59,7 +59,7 @@ typedef struct WorkerNode
 
 /* Config variables managed via guc.c */
 extern int MaxWorkerNodesTracked;
-extern int DatabasesPerWorker;
+extern int MaxDatabasesPerWorkerNodesTracked;
 extern char *WorkerListFileName;
 extern char *CurrentCluster;