[#23770] YSQL: Deterministically populate catalog cache in tests with…

… Connection Manager enabled

Summary:
Some tests allow catalog cache population (via SELECT, INSERT type queries) before running subsequent queries dependent on this population. When these tests are run with connection manager such that a pool of server connections are warmed up and randomly allocated to logical connections, it becomes possible that the same logical connection may not be attached to the server that had its cache populated.

To solve this problem, we introduce a new mode where connections are instead allocated on a round robin basis so that the cache populating query can be run a fixed number of times to consistently have the entire pool's catalog cache populated.

Jira: DB-12674

Test Plan: Jenkins: enable connection manager, all tests

Reviewers: stiwary, mkumar, skumar

Reviewed By: skumar

Subscribers: hsunder, yql, mkumar, skumar

Differential Revision: https://phorge.dev.yugabyte.com/D37725

java/yb-pgsql/src/test/java/org/yb/pgsql/BasePgSQLTest.java

@@ -178,6 +178,17 @@ public class BasePgSQLTest extends BaseMiniClusterTest {
         "variables at the beginning of transaction boundaries, causing erroneous results in " +
         "the test, leading to failure.";
 
+  // Warmup modes for Connection Manager during test runs.
+  protected static enum ConnectionManagerWarmupMode {
+    NONE,
+    RANDOM,
+    ROUND_ROBIN
+  }
+
+  protected static ConnectionManagerWarmupMode warmupMode = ConnectionManagerWarmupMode.RANDOM;
+
+  protected static final int CONN_MGR_WARMUP_BACKEND_COUNT = 3;
+
   // CQL and Redis settings, will be reset before each test via resetSettings method.
   protected boolean startCqlProxy = false;
   protected boolean startRedisProxy = false;
@@ -307,8 +318,6 @@ protected void customizeMiniClusterBuilder(MiniYBClusterBuilder builder) {
       builder.enableYsqlConnMgr(true);
       builder.addCommonTServerFlag("ysql_conn_mgr_stats_interval",
         Integer.toString(CONNECTIONS_STATS_UPDATE_INTERVAL_SECS));
-      builder.addCommonTServerFlag(
-        "TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach", "true");
     }
   }
 
@@ -590,6 +599,25 @@ protected void closeControlConnOnReloadConfig(int attempts) {
     }
   }
 
+  protected
+  void setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode wm) throws Exception {
+    if (!isTestRunningWithConnectionManager()) {
+      return;
+    }
+
+    Map<String, String> tsFlagMap = getTServerFlags();
+    tsFlagMap.put("TEST_ysql_conn_mgr_dowarmup_all_pools_mode",
+      warmupMode.toString().toLowerCase());
+    warmupMode = wm;
+    Map<String, String> masterFlagMap = getMasterFlags();
+    restartClusterWithFlags(masterFlagMap, tsFlagMap);
+  }
+
+  protected boolean isConnMgrWarmupRoundRobinMode() {
+    return isTestRunningWithConnectionManager() &&
+      warmupMode == ConnectionManagerWarmupMode.ROUND_ROBIN;
+  }
+
   protected void recreateWithYsqlVersion(YsqlSnapshotVersion version) throws Exception {
     destroyMiniCluster();
     pgInitialized = false;
@@ -2216,6 +2244,16 @@ protected Long getNumStorageRoundtrips(Statement stmt, String query) throws Exce
   protected long getNumDocdbRequests(Statement stmt, String query) throws Exception {
     // Executing query once just in case if master catalog cache is not refreshed
     stmt.execute(query);
+
+    // Execute query twice more to deterministically populate all caches if
+    // connection manager is enabled in round robin allocation mode.
+    // Additionally execute it once more to allow rotation onto a new physical
+    // connection before executing the subsequent queries.
+    if (isConnMgrWarmupRoundRobinMode()) {
+      for (int i = 0; i < CONN_MGR_WARMUP_BACKEND_COUNT; i++) {
+        stmt.execute(query);
+      }
+    }
     Long rpc_count_before =
       getTServerMetric("handler_latency_yb_tserver_PgClientService_Perform").count;
     stmt.execute(query);

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgConstraintCache.java

@@ -46,6 +46,7 @@ public class TestPgConstraintCache extends BasePgSQLTest {
    */
   @Test
   public void groupByTest() throws Exception {
+    setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
     try (Statement stmt = connection.createStatement()) {
       // Set up the table and data
       stmt.execute("CREATE TABLE t(a INT PRIMARY KEY, b INT, c INT);");
@@ -69,6 +70,14 @@ public void groupByTest() throws Exception {
         assertEquals("a", groupKeyArray.getString(0));
       }
 
+      // Run the query a few times to warm up the caches of all backends when
+      // Connection Manager is in round-robin warmup mode
+      if (isConnMgrWarmupRoundRobinMode()) {
+        for (int i = 0; i < CONN_MGR_WARMUP_BACKEND_COUNT; i++) {
+          stmt.execute(query);
+        }
+      }
+
       // Run the query again and check that we're not doing catalog read requests
       rs = stmt.executeQuery(query);
       while (rs.next()) {

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgCostModelSeekNextEstimation.java

@@ -306,11 +306,20 @@ protected Map<String, String> getTServerFlags() {
   // in Nov/2023.
   @Test
   public void testSeekNextEstimationIndexScan() throws Exception {
+    setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
     boolean isConnMgr = isTestRunningWithConnectionManager();
     try (Statement stmt = this.connection2.createStatement()) {
+      // Warmup the cache when Connection Manager is enabled.
+      // Additionally warmup all backends in round-robin mode.
       if (isConnMgr) {
         testExplainDebug(stmt, String.format("/*+IndexScan(%s)*/ SELECT * "
         + "FROM %s WHERE k1 IN (4, 8)", T1_NAME, T1_NAME), null);
+        if (isConnMgrWarmupRoundRobinMode()) {
+          for (int i = 0; i < CONN_MGR_WARMUP_BACKEND_COUNT - 1; i++) {
+            testExplainDebug(stmt, String.format("/*+IndexScan(%s)*/ SELECT * "
+            + "FROM %s WHERE k1 IN (4, 8)", T1_NAME, T1_NAME), null);
+          }
+        }
       }
       testSeekAndNextEstimationIndexScanHelper(stmt, String.format("/*+IndexScan(%s)*/ SELECT * "
         + "FROM %s WHERE k1 IN (4, 8)", T1_NAME, T1_NAME),
@@ -423,12 +432,21 @@ public void testSeekNextEstimationIndexScan() throws Exception {
   @Test
   public void testSeekNextEstimationBitmapScan() throws Exception {
     assumeTrue("BitmapScan has much fewer nexts in fastdebug (#22052)", TestUtils.isReleaseBuild());
+    setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
     boolean isConnMgr = isTestRunningWithConnectionManager();
     try (Statement stmt = this.connection2.createStatement()) {
       stmt.execute("SET work_mem TO '1GB'"); /* avoid getting close to work_mem */
+      // Warmup the cache when Connection Manager is enabled.
+      // Additionally warmup all backends in round-robin mode.
       if (isConnMgr) {
         testExplainDebug(stmt, String.format("/*+BitmapScan(%s)*/ SELECT * "
         + "FROM %s WHERE k1 IN (4, 8)", T1_NAME, T1_NAME), null);
+        if (isConnMgrWarmupRoundRobinMode()) {
+          for (int i = 0; i < CONN_MGR_WARMUP_BACKEND_COUNT - 1; i++) {
+            testExplainDebug(stmt, String.format("/*+BitmapScan(%s)*/ SELECT * "
+            + "FROM %s WHERE k1 IN (4, 8)", T1_NAME, T1_NAME), null);
+          }
+        }
       }
       testSeekAndNextEstimationBitmapScanHelper(stmt, String.format("/*+BitmapScan(%s)*/ SELECT * "
         + "FROM %s WHERE k1 IN (4, 8)", T1_NAME, T1_NAME),

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgExplainAnalyze.java

@@ -96,6 +96,7 @@ public void testExplainNoTiming(String query, Checker checker) throws Exception
 
   @Test
   public void testSeqScan() throws Exception {
+    setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
     try (Statement stmt = connection.createStatement()) {
       Checker checker = makeTopLevelBuilder()
           .storageReadRequests(Checkers.equal(5))
@@ -120,6 +121,14 @@ public void testSeqScan() throws Exception {
       // Warm up the cache to get catalog requests out of the way.
       testExplain(String.format("SELECT * FROM %s", TABLE_NAME), null);
 
+      // Additionally warmup the cache for all backends when Connection Manager
+      // is in round-robin warmup mode.
+      if (isConnMgrWarmupRoundRobinMode()) {
+        for(int i = 0; i < CONN_MGR_WARMUP_BACKEND_COUNT; i++) {
+            testExplain(String.format("SELECT * FROM %s", TABLE_NAME), null);
+        }
+      }
+
       // Seq Scan (ybc_fdw ForeignScan)
       testExplain(String.format("SELECT * FROM %s", TABLE_NAME), checker);
 
@@ -333,13 +342,22 @@ public void testEmptyNestedLoop() throws Exception {
 
   @Test
   public void testInsertValues() throws Exception {
+    setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
     try (Statement stmt = connection.createStatement()) {
       // reduce the batch size to avoid 0 wait time
       stmt.execute("SET ysql_session_max_batch_size = 4");
 
       // Warm up the cache to get catalog requests out of the way.
       testExplain(String.format("SELECT * FROM %s", TABLE_NAME), null);
 
+      // Additionally warmup the cache for all backends when Connection Manager
+      // is in round-robin warmup mode.
+      if (isConnMgrWarmupRoundRobinMode()) {
+        for(int i = 0; i < CONN_MGR_WARMUP_BACKEND_COUNT; i++) {
+            testExplain(String.format("SELECT * FROM %s", TABLE_NAME), null);
+        }
+      }
+
       ObjectChecker planChecker =
           makePlanBuilder()
               .nodeType(NODE_MODIFY_TABLE)

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgSelect.java

@@ -1548,6 +1548,7 @@ public void testIsNotNull() throws Exception {
 
   @Test
   public void testInequalitiesRangePartitioned() throws Exception {
+      setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
       String query = "CREATE TABLE sample (key int, val int, primary key(key asc) ) " +
                      "SPLIT AT VALUES ((65535), (2000000000), (2100000000) )";
       try (Statement statement = connection.createStatement()) {
@@ -1606,6 +1607,7 @@ public void testInequalitiesRangePartitioned() throws Exception {
 
   @Test
   public void testINQueriesRangePartitioned() throws Exception {
+      setConnMgrWarmupModeAndRestartCluster(ConnectionManagerWarmupMode.ROUND_ROBIN);
 
       // Creating a table with the same schema that is created for the previous test
       // testInequalitiesRangePartitioned.

java/yb-ysql-conn-mgr/src/test/java/org/yb/ysqlconnmgr/BaseYsqlConnMgr.java

@@ -55,7 +55,7 @@ protected void customizeMiniClusterBuilder(MiniYBClusterBuilder builder) {
     builder.addCommonTServerFlag("ysql_conn_mgr_dowarmup", "false");
     if (warmup_random_mode) {
       builder.addCommonTServerFlag(
-      "TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach", "true");
+      "TEST_ysql_conn_mgr_dowarmup_all_pools_mode", "random");
     }
   }
 
@@ -65,7 +65,7 @@ protected ConnectionBuilder getConnectionBuilder() {
 
   protected void disableWarmupRandomMode(MiniYBClusterBuilder builder) {
     builder.addCommonTServerFlag(
-        "TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach", "false");
+        "TEST_ysql_conn_mgr_dowarmup_all_pools_mode", "none");
     warmup_random_mode = false;
     return;
   }

src/odyssey/sources/router.c

@@ -599,18 +599,26 @@ od_router_status_t od_router_attach(od_router_t *router,
 	od_server_t *server;
 	int busyloop_sleep = 0;
 	int busyloop_retry = 0;
-	const char *is_warmup_needed = getenv("YB_YSQL_CONN_MGR_DOWARMUP_ALL_POOLS_RANDOM_ATTACH");
-	bool enable_warmup_and_random_allot = false;
-	if (is_warmup_needed != NULL && strcmp(is_warmup_needed, "true") == 0)
-		enable_warmup_and_random_allot = true;
+
+	const char *is_warmup_needed_flag = getenv("YB_YSQL_CONN_MGR_DOWARMUP_ALL_POOLS_MODE");
+	bool is_warmup_needed = false;
+	bool random_allot = false;
+
+	is_warmup_needed = is_warmup_needed_flag != NULL && strcmp(is_warmup_needed_flag, "none") != 0;
+	random_allot = is_warmup_needed && strcmp(is_warmup_needed_flag, "random") == 0;
 
 	for (;;) {
 
-		if (enable_warmup_and_random_allot)
+		if (is_warmup_needed)
 		{
-			server = od_server_pool_idle_random(&route->server_pool);
+			if (random_allot)
+				server = yb_od_server_pool_idle_random(&route->server_pool);
+			else /* round_robin allotment */
+				server = yb_od_server_pool_idle_last(&route->server_pool);
+
 			if (server &&
-				(od_server_pool_total(&route->server_pool) >= route->rule->min_pool_size))
+			    (od_server_pool_total(&route->server_pool) >=
+			     route->rule->min_pool_size))
 				goto attach;
 		}
 		else
@@ -691,7 +699,7 @@ od_router_status_t od_router_attach(od_router_t *router,
 
 	/* create new server object */
 	bool created_atleast_one = false;
-	while (enable_warmup_and_random_allot &&
+	while (is_warmup_needed &&
 		  (od_server_pool_total(&route->server_pool) < route->rule->min_pool_size))
 	{
 		server = od_server_allocate(

src/odyssey/sources/server_pool.h

@@ -129,7 +129,20 @@ OD_SERVER_POOL_NEXT_DECLARE(pg, od_server_t)
 OD_SERVER_POOL_NEXT_DECLARE(ldap, od_ldap_server_t)
 #endif
 
-static inline od_server_t *od_server_pool_idle_random (od_server_pool_t *pool)
+static inline od_server_t *yb_od_server_pool_idle_last(od_server_pool_t *pool)
+{
+	od_list_t *target = &pool->idle;
+	od_server_t *server = NULL;
+	od_list_t *i, *n;
+	int len = pool->count_idle;
+	if (len == 0)
+		return NULL;
+	server = od_container_of(target->prev, od_server_t, link);
+
+	return server;
+}
+
+static inline od_server_t *yb_od_server_pool_idle_random (od_server_pool_t *pool)
 {
 	od_list_t *target = &pool->idle;
 	od_server_t *server;

src/postgres/src/backend/commands/dbcommands.c

@@ -1105,7 +1105,7 @@ dropdb(const char *dbname, bool missing_ok, bool force)
 			yb_net_client_connections +=
 				yb_num_logical_conn - yb_num_physical_conn_from_ysqlconnmgr;
 
-			if (YbIsYsqlConnMgrWarmupRandomEnabled())
+			if (YbIsYsqlConnMgrWarmupModeEnabled())
 				yb_net_client_connections = yb_num_logical_conn;
 		}
 
@@ -1308,7 +1308,7 @@ RenameDatabase(const char *oldname, const char *newname)
 			yb_net_client_connections +=
 				yb_num_logical_conn - yb_num_physical_conn_from_ysqlconnmgr;
 
-			if (YbIsYsqlConnMgrWarmupRandomEnabled())
+			if (YbIsYsqlConnMgrWarmupModeEnabled())
 				yb_net_client_connections = yb_num_logical_conn;
 		}
 

src/postgres/src/backend/utils/misc/pg_yb_utils.c

@@ -4995,9 +4995,9 @@ bool YbUseFastBackwardScan() {
   return *(YBCGetGFlags()->ysql_use_fast_backward_scan);
 }
 
-bool YbIsYsqlConnMgrWarmupRandomEnabled()
+bool YbIsYsqlConnMgrWarmupModeEnabled()
 {
-	return *(YBCGetGFlags()->TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach);
+	return strcmp(YBCGetGFlags()->TEST_ysql_conn_mgr_dowarmup_all_pools_mode, "none") != 0;
 }
 
 /* Used in YB to check if an attribute is a key column. */

src/postgres/src/backend/utils/misc/yb_ysql_conn_mgr_helper.c

@@ -701,7 +701,7 @@ SetLogicalClientUserDetailsIfValid(const char *rolename, bool *is_superuser,
 		yb_net_client_connections +=
 		yb_num_logical_conn - yb_num_physical_conn_from_ysqlconnmgr;
 
-		if (YbIsYsqlConnMgrWarmupRandomEnabled())
+		if (YbIsYsqlConnMgrWarmupModeEnabled())
 			yb_net_client_connections = yb_num_logical_conn;
 	}
 

src/postgres/src/include/pg_yb_utils.h

@@ -1171,7 +1171,7 @@ extern YbReadTimePointHandle YbBuildCurrentReadTimePointHandle();
 
 extern bool YbUseFastBackwardScan();
 
-extern bool YbIsYsqlConnMgrWarmupRandomEnabled();
+extern bool YbIsYsqlConnMgrWarmupModeEnabled();
 
 bool YbIsAttrPrimaryKeyColumn(Relation rel, AttrNumber attnum);
 

src/yb/yql/pggate/ybc_pg_typedefs.h

@@ -394,7 +394,7 @@ typedef struct PgGFlagsAccessor {
   const bool*     TEST_ysql_hide_catalog_version_increment_log;
   const bool*     TEST_generate_ybrowid_sequentially;
   const bool*     ysql_use_fast_backward_scan;
-  const bool*     TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach;
+  const char*     TEST_ysql_conn_mgr_dowarmup_all_pools_mode;
 } YBCPgGFlagsAccessor;
 
 typedef struct YbTablePropertiesData {

src/yb/yql/pggate/ybc_pggate.cc

@@ -91,13 +91,13 @@ DEFINE_UNKNOWN_bool(ysql_disable_server_file_access, false,
 
 DEFINE_NON_RUNTIME_bool(ysql_enable_profile, false, "Enable PROFILE feature.");
 
-DEFINE_test_flag(bool, ysql_conn_mgr_dowarmup_all_pools_random_attach, false,
+DEFINE_test_flag(string, ysql_conn_mgr_dowarmup_all_pools_mode, "random",
   "Enable precreation of server connections in every pool in Ysql Connection Manager and "
-  " randomly attach any idle server connection to client to serve it's queries. "
+  "choose the mode of attachment of idle server connections to clients to serve their queries. "
   "ysql_conn_mgr_dowarmup is responsible for creating server connections only in "
   "yugabyte (user), yugabyte (database) pool during the initialization of connection "
-  "manager process. Whereas this flag will create max of ysql_conn_mgr_min_conns_per_db"
-  " and 3 number of server connections in any pool whenever there is a requirement to create "
+  "manager process. This flag will create max(ysql_conn_mgr_min_conns_per_db, "
+  "3) number of server connections in any pool whenever there is a requirement to create the "
   "first backend process in that particular pool.");
 
 // This gflag should be deprecated but kept to avoid breaking some customer
@@ -1966,8 +1966,8 @@ const YBCPgGFlagsAccessor* YBCGetGFlags() {
       .TEST_generate_ybrowid_sequentially =
           &FLAGS_TEST_generate_ybrowid_sequentially,
       .ysql_use_fast_backward_scan = &FLAGS_use_fast_backward_scan,
-      .TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach =
-          &FLAGS_TEST_ysql_conn_mgr_dowarmup_all_pools_random_attach,
+      .TEST_ysql_conn_mgr_dowarmup_all_pools_mode =
+          FLAGS_TEST_ysql_conn_mgr_dowarmup_all_pools_mode.c_str(),
   };
   // clang-format on
   return &accessor;