[#18886] YSQL: Supporting login limit per user with Ysql Connection M…

…anager

Summary:
This diff adds the support for login limit connections per user in ysql connection manager.

The only known limitation of this implementation is that the client has to wait for few seconds (by default 10seconds) for connection manager stats to get updated which is used to calculate current number of connections made by a particular user. Track Here: [[ #21645 | GH #21645 ]]

One of the existing java unit tests (`org.yb.pgsql.TestPgAuthorization#testAttributes`) which tests connection limit per user is not fixed as part of this diff, as it still fails while dropping the user which is not supported in ysql connection manager at the time of creating this diff. (Tracked by [[ #21862 | GH #21862 ]]  ). Added a TODO comment to fix conn per limit user part of the test.

Jira: DB-7747

Test Plan:
 Added java unit test:
```./yb_build.sh --enable-ysql-conn-mgr-test --java-test org.yb.ysqlconnmgr.TestUserLoginLimit```

Reviewers: rbarigidad

Reviewed By: rbarigidad

Subscribers: rbarigidad, nkumar, yql

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

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

@@ -494,7 +494,11 @@ public void testAttributes() throws Exception {
 
       ConnectionBuilder limitRoleUserConnBldr = getConnectionBuilder().withUser("limit_role");
       try (Connection ignored1 = limitRoleUserConnBldr.connect()) {
+        // TODO(GH #18886) While running the test on ysql conn mgr port, add a wait for
+        // ysql conn mgr stats to get updated to check if conn limit is exceeded or not.
         try (Connection connection2 = limitRoleUserConnBldr.connect()) {
+          // TODO(GH #18886) While running the test on ysql conn mgr port, add a wait for
+          // ysql conn mgr stats to get updated to check if conn limit is exceeded or not.
           // Third concurrent connection causes error.
           try (Connection ignored3 = limitRoleUserConnBldr.connect()) {
             fail("Expected third login attempt to fail");
@@ -507,7 +511,8 @@ public void testAttributes() throws Exception {
 
           // Close second connection.
           connection2.close();
-
+          // TODO(GH #18886) While running the test on ysql conn mgr port, add a wait for
+          // ysql conn mgr stats to get updated to check if conn limit is exceeded or not.
           // New connection now succeeds.
           try (Connection ignored2 = limitRoleUserConnBldr.connect()) {
             // No-op.

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

@@ -0,0 +1,92 @@
+// Copyright (c) YugabyteDB, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations
+// under the License.
+//
+package org.yb.ysqlconnmgr;
+
+import static org.yb.AssertionWrappers.*;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.hamcrest.CoreMatchers;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.yb.minicluster.MiniYBClusterBuilder;
+import org.yb.pgsql.BasePgSQLTest;
+import org.yb.pgsql.ConnectionBuilder;
+import org.yb.pgsql.ConnectionEndpoint;
+
+@RunWith(value = YBTestRunnerYsqlConnMgr.class)
+public class TestUserLoginLimit extends BaseYsqlConnMgr {
+
+    @Override
+    protected void customizeMiniClusterBuilder(MiniYBClusterBuilder builder) {
+        super.customizeMiniClusterBuilder(builder);
+
+        builder.addCommonTServerFlag("ysql_conn_mgr_stats_interval",
+            Integer.toString(BasePgSQLTest.CONNECTIONS_STATS_UPDATE_INTERVAL_SECS));
+    }
+
+    @Test
+    public void testUserLoginLimit() throws Exception {
+        try (Connection connection = getConnectionBuilder()
+                    .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                    .withUser("yugabyte")
+                    .withPassword("yugabyte")
+                    .connect();
+            Statement statement = connection.createStatement()) {
+
+            statement.execute("CREATE ROLE limit_role LOGIN CONNECTION LIMIT 2");
+
+            ConnectionBuilder limitRoleUserConnBldr =
+                                getConnectionBuilder()
+                                .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                                .withUser("limit_role");
+            try (Connection ignored1 = limitRoleUserConnBldr.connect()) {
+                BasePgSQLTest.waitForStatsToGetUpdated();
+                try (Connection connection2 = limitRoleUserConnBldr.connect()) {
+                    BasePgSQLTest.waitForStatsToGetUpdated();
+                    // Third concurrent connection causes error.
+                    try (Connection ignored3 = limitRoleUserConnBldr.connect()) {
+                        fail("Expected third login attempt to fail");
+                    } catch (SQLException sqle) {
+                        assertThat(
+                            sqle.getMessage(),
+                            CoreMatchers.containsString("too many connections for " +
+                                    "role \"limit_role\"")
+                        );
+                    }
+
+                    // Close second connection.
+                    connection2.close();
+                    BasePgSQLTest.waitForStatsToGetUpdated();
+
+                    // New connection now succeeds.
+                    try (Connection ignored2 = limitRoleUserConnBldr.connect()) {
+                        // No-op.
+                    }
+                }
+            }
+
+        } catch (Exception e) {
+            LOG.error("Allowing unexpected number of connections than what limit has been " +
+                "set for given user: ", e);
+            fail ("Creating unexpected number of connections than what limit has been set for " +
+                "given user");
+        }
+
+    }
+}

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

@@ -41,6 +41,7 @@
 #include "libpq/pqformat.h"
 #include "pg_yb_utils.h"
 #include "storage/dsm_impl.h"
+#include "storage/procarray.h"
 #include "utils/guc.h"
 #include "utils/syscache.h"
 
@@ -635,13 +636,15 @@ YbHandleSetSessionParam(int yb_client_id)
  * Checks done in this function:
  *  		1. Does the role exist.
  * 			2. Is the role permitted to login.
+ * 			3. Check whether connection limit is exceeded for the role
  */
 static int8_t
 SetLogicalClientUserDetailsIfValid(const char *rolename, bool *is_superuser,
 						  Oid *roleid)
 {
 	HeapTuple	roleTup;
 	Form_pg_authid rform;
+	int yb_net_client_connections = 0;
 	char	   *rname;
 
 	/* TODO(janand) GH #19951 Do we need support for initializing via OID */
@@ -679,9 +682,39 @@ SetLogicalClientUserDetailsIfValid(const char *rolename, bool *is_superuser,
 	}
 
 	/*
-	 * TODO(janand) GH #18886 Add support for "too many connections for role"
-	 * error.
-	 */
+	* yb_num_logical_conn: Stores count for all client connections made to conn mgr.
+	* yb_num_physical_conn_from_ysqlconnmgr: Stores physical connection count created from
+	* conn mgr to yb/database.
+	* CountUserBackends: Function returns total number of backend connections made by given
+	* user(roleid). It will be sum of physical connections from connection manager and direct
+	* connections to yb/database. 
+	*/
+
+	uint32_t yb_num_logical_conn = 0,
+				 yb_num_physical_conn_from_ysqlconnmgr = 0;
+
+		yb_net_client_connections = CountUserBackends(*roleid);
+
+		if (IsYugaByteEnabled() &&
+		YbGetNumYsqlConnMgrConnections(NULL, rname, &yb_num_logical_conn,
+									   &yb_num_physical_conn_from_ysqlconnmgr)) {
+			yb_net_client_connections +=
+			yb_num_logical_conn - yb_num_physical_conn_from_ysqlconnmgr;
+
+		}
+
+	if (rform->rolconnlimit >= 0 &&
+			!rform->rolsuper &&
+			yb_net_client_connections + 1 > rform->rolconnlimit)
+	{
+		YbSendFatalForLogicalConnectionPacket();
+		ereport(WARNING,
+				(errcode(ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION),
+				 errmsg("too many connections for role \"%s\"", rname)));
+		ReleaseSysCache(roleTup);
+		return -1;
+	}
+
 	ReleaseSysCache(roleTup);
 	return 0;
 }