pgwire: add server.cockroach_cloud.max_client_connections_per_gateway…

… cluster setting

This setting specifies a maximum number of non root connections that a server
can have open at any given time.
<0 - Connections are unlimited (existing behavior)
=0 - Non root connections are disabled
>0 - Non root connections are limited
If a new connection would exceed this limit, an error message is
returned as postgres: "cluster connections are limited" with the 53300
error code that corresponds to "too many connections".

This functionality is required for serverless.
Part of: https://cockroachlabs.atlassian.net/browse/CC-9288

Release note: None

pkg/sql/conn_executor.go

@@ -323,7 +323,8 @@ type Server struct {
 
 	mu struct {
 		syncutil.Mutex
-		connectionCount int64
+		connectionCount     int64
+		rootConnectionCount int64
 	}
 }
 
@@ -789,14 +790,30 @@ func (s *Server) IncrementConnectionCount() {
 	s.mu.connectionCount++
 }
 
+// IncrementRootConnectionCount increases both connectionCount and rootConnectionCount by 1.
+func (s *Server) IncrementRootConnectionCount() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.mu.connectionCount++
+	s.mu.rootConnectionCount++
+}
+
 // DecrementConnectionCount decreases connectionCount by 1.
 func (s *Server) DecrementConnectionCount() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	s.mu.connectionCount--
 }
 
-// IncrementConnectionCountIfLessThan increases connectionCount by and returns true if allowedConnectionCount < max,
+// DecrementRootConnectionCount decreases both connectionCount and rootConnectionCount by 1.
+func (s *Server) DecrementRootConnectionCount() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.mu.connectionCount--
+	s.mu.rootConnectionCount--
+}
+
+// IncrementConnectionCountIfLessThan increases connectionCount by 1 and returns true if connectionCount < max,
 // otherwise it does nothing and returns false.
 func (s *Server) IncrementConnectionCountIfLessThan(max int64) bool {
 	s.mu.Lock()
@@ -815,6 +832,13 @@ func (s *Server) GetConnectionCount() int64 {
 	return s.mu.connectionCount
 }
 
+// GetNonRootConnectionCount returns the current number of non root connections.
+func (s *Server) GetNonRootConnectionCount() int64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.mu.connectionCount - s.mu.rootConnectionCount
+}
+
 // ConnectionHandler is the interface between the result of SetupConn
 // and the ServeConn below. It encapsulates the connExecutor and hides
 // it away from other packages.

pkg/sql/pgwire/conn.go

@@ -232,13 +232,34 @@ func (c *conn) sendError(ctx context.Context, execCfg *sql.ExecutorConfig, err e
 }
 
 func (c *conn) checkMaxConnections(ctx context.Context, sqlServer *sql.Server) error {
+	// Root user is not affected by connection limits.
+	if c.sessionArgs.User.IsRootUser() {
+		sqlServer.IncrementRootConnectionCount()
+		return nil
+	}
+
+	// First check maxNumNonRootConnections.
+	// Note(alyshan): maxNumNonRootConnections is used by Cockroach Cloud for limiting connections to
+	// serverless clusters.
+	maxNonRootConnectionsValue := maxNumNonRootConnections.Get(&sqlServer.GetExecutorConfig().Settings.SV)
+	if maxNonRootConnectionsValue >= 0 && sqlServer.GetNonRootConnectionCount() >= maxNonRootConnectionsValue {
+		// TODO(alyshan): Add another cluster setting that supplements the error message, so that
+		// clients can know *why* this ReadOnly setting is set and limiting their connections.
+		return c.sendError(ctx, sqlServer.GetExecutorConfig(), errors.WithHintf(
+			pgerror.New(pgcode.TooManyConnections, "cluster connections are limited"),
+			"the maximum number of allowed connections is %d",
+			maxNonRootConnectionsValue,
+		))
+	}
+
+	// Then check maxNumNonAdminConnections.
 	if c.sessionArgs.IsSuperuser {
-		// This user is a super user and is therefore not affected by connection limits.
+		// This user is a super user and is therefore not affected by maxNumNonAdminConnections.
 		sqlServer.IncrementConnectionCount()
 		return nil
 	}
 
-	maxNumConnectionsValue := maxNumConnections.Get(&sqlServer.GetExecutorConfig().Settings.SV)
+	maxNumConnectionsValue := maxNumNonAdminConnections.Get(&sqlServer.GetExecutorConfig().Settings.SV)
 	if maxNumConnectionsValue < 0 {
 		// Unlimited connections are allowed.
 		sqlServer.IncrementConnectionCount()
@@ -249,7 +270,7 @@ func (c *conn) checkMaxConnections(ctx context.Context, sqlServer *sql.Server) e
 			pgerror.New(pgcode.TooManyConnections, "sorry, too many clients already"),
 			"the maximum number of allowed connections is %d and can be modified using the %s config key",
 			maxNumConnectionsValue,
-			maxNumConnections.Key(),
+			maxNumNonAdminConnections.Key(),
 		))
 	}
 	return nil
@@ -707,7 +728,11 @@ func (c *conn) processCommandsAsync(
 		if retErr = c.checkMaxConnections(ctx, sqlServer); retErr != nil {
 			return
 		}
-		defer sqlServer.DecrementConnectionCount()
+		if c.sessionArgs.User.IsRootUser() {
+			defer sqlServer.DecrementRootConnectionCount()
+		} else {
+			defer sqlServer.DecrementConnectionCount()
+		}
 
 		if retErr = c.authOKMessage(); retErr != nil {
 			return

pkg/sql/pgwire/conn_test.go

@@ -1867,8 +1867,9 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 	defer testServer.Stopper().Stop(ctx)
 
 	// Users.
-	admin := username.RootUser
+	rootUser := username.RootUser
 	nonAdmin := username.TestUser
+	admin := "testadmin"
 
 	// openConnWithUser opens a connection to the testServer for the given user
 	// and always returns an associated cleanup function, even in case of error,
@@ -1879,7 +1880,7 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 			testServer.ServingSQLAddr(),
 			t.Name(),
 			url.UserPassword(user, user),
-			user == admin,
+			user == rootUser,
 		)
 		defer cleanup()
 		conn, err := pgx.Connect(ctx, pgURL.String())
@@ -1923,7 +1924,7 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 	}
 
 	getMaxConnections := func() int {
-		conn, cleanup := openConnWithUserSuccess(admin)
+		conn, cleanup := openConnWithUserSuccess(rootUser)
 		defer cleanup()
 		var maxConnections int
 		err := conn.QueryRow(ctx, "SHOW CLUSTER SETTING server.max_connections_per_gateway").Scan(&maxConnections)
@@ -1932,21 +1933,34 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 	}
 
 	setMaxConnections := func(maxConnections int) {
-		conn, cleanup := openConnWithUserSuccess(admin)
+		conn, cleanup := openConnWithUserSuccess(rootUser)
 		defer cleanup()
 		_, err := conn.Exec(ctx, "SET CLUSTER SETTING server.max_connections_per_gateway = $1", maxConnections)
 		require.NoError(t, err)
 	}
 
-	createUser := func(user string) {
-		conn, cleanup := openConnWithUserSuccess(admin)
+	setMaxNonRootConnections := func(maxConnections int) {
+		conn, cleanup := openConnWithUserSuccess(rootUser)
+		defer cleanup()
+		_, err := conn.Exec(ctx, "SET CLUSTER SETTING server.cockroach_cloud.max_client_connections_per_gateway = $1", maxConnections)
+		require.NoError(t, err)
+	}
+
+	createUser := func(user string, isAdmin bool) {
+		conn, cleanup := openConnWithUserSuccess(rootUser)
 		defer cleanup()
 		_, err := conn.Exec(ctx, fmt.Sprintf("CREATE USER %[1]s WITH PASSWORD '%[1]s'", user))
 		require.NoError(t, err)
+
+		if isAdmin {
+			_, err := conn.Exec(ctx, fmt.Sprintf("grant admin to %[1]s", user))
+			require.NoError(t, err)
+		}
 	}
 
 	// create nonAdmin
-	createUser(nonAdmin)
+	createUser(nonAdmin, false)
+	createUser(admin, true)
 	requireConnectionCount(t, 0)
 
 	// assert default value
@@ -2012,6 +2026,33 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 		nonAdminCleanup2()
 		requireConnectionCount(t, 0)
 	})
+
+	t.Run("0 max_non_root_connections", func(t *testing.T) {
+		setMaxNonRootConnections(0)
+		requireConnectionCount(t, 0)
+		// can connect with root
+		_, rootCleanup := openConnWithUserSuccess(rootUser)
+		requireConnectionCount(t, 1)
+		// can't connect with non root
+		openConnWithUserError(admin)
+		requireConnectionCount(t, 1)
+		rootCleanup()
+		requireConnectionCount(t, 0)
+	})
+
+	t.Run("1 max_non_root_connections", func(t *testing.T) {
+		setMaxNonRootConnections(1)
+		requireConnectionCount(t, 0)
+		// can connect with root
+		_, rootCleanup := openConnWithUserSuccess(rootUser)
+		requireConnectionCount(t, 1)
+		// can connect with non root
+		_, nonAdminCleanup := openConnWithUserSuccess(nonAdmin)
+		requireConnectionCount(t, 2)
+		rootCleanup()
+		nonAdminCleanup()
+		requireConnectionCount(t, 0)
+	})
 }
 
 func TestConnCloseReleasesReservedMem(t *testing.T) {

pkg/sql/pgwire/server.go

@@ -82,7 +82,11 @@ var logSessionAuth = settings.RegisterBoolSetting(
 	"if set, log SQL session login/disconnection events (note: may hinder performance on loaded nodes)",
 	false).WithPublic()
 
-var maxNumConnections = settings.RegisterIntSetting(
+// TODO(alyshan): This setting is enforcing max number of connections with superusers not being affected by
+// the limit. However, admin users connections are counted towards the max count. So we should either update the
+// description to say "the maximum number of connections per gateway ... Superusers are not affected by this limit"
+// or stop counting superuser connections towards the max count.
+var maxNumNonAdminConnections = settings.RegisterIntSetting(
 	settings.TenantWritable,
 	"server.max_connections_per_gateway",
 	"the maximum number of non-superuser SQL connections per gateway allowed at a given time "+
@@ -91,6 +95,21 @@ var maxNumConnections = settings.RegisterIntSetting(
 	-1, // Postgres defaults to 100, but we default to -1 to match our previous behavior of unlimited.
 ).WithPublic()
 
+// Note(alyshan): This setting is not public. It is intended to be used by Cockroach Cloud to limit
+// connections to serverless clusters while still being able to connect from the Cockroach Cloud control plane.
+// This setting may be extended one day to include an arbitrary list of users to exclude from connection limiting.
+// This setting may be removed one day.
+var maxNumNonRootConnections = settings.RegisterIntSetting(
+	settings.TenantWritable,
+	"server.cockroach_cloud.max_client_connections_per_gateway",
+	"this setting is intended to be used by Cockroach Cloud for limiting connections to serverless clusters. "+
+		"The maximum number of SQL connections per gateway allowed at a given time "+
+		"(note: this will only limit future connection attempts and will not affect already established connections). "+
+		"Negative values result in unlimited number of connections. Cockroach Cloud internal users (including root user) "+
+		"are not affected by this limit.",
+	-1,
+)
+
 const (
 	// ErrSSLRequired is returned when a client attempts to connect to a
 	// secure server in cleartext.