From 3fb40b63172202410a266249841433afea8e6263 Mon Sep 17 00:00:00 2001
From: Alyshan Jahani <alyshan@cockroachlabs.com>
Date: Tue, 4 Apr 2023 14:03:40 -0400
Subject: [PATCH] pgwire: add
 server.max_non_root_connections_per_gateway_reason cluster setting

This setting can be used to customize the error message returned when
connections are denied due to a limit specified via
server.max_non_root_connections_per_gateway.

This functionality is required for serverless. Being able to indicate to the
user why a limit is placed on their number of connections will support a better
UX. Specifically, indicating to the user which resource limits they have hit.
Part of: https://cockroachlabs.atlassian.net/browse/CC-9288
---
 pkg/sql/pgwire/conn.go      |  9 ++++++---
 pkg/sql/pgwire/conn_test.go | 20 ++++++++++++++++++++
 pkg/sql/pgwire/server.go    | 10 ++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/pkg/sql/pgwire/conn.go b/pkg/sql/pgwire/conn.go
index 8b58eb410f4e..6160960b0bc9 100644
--- a/pkg/sql/pgwire/conn.go
+++ b/pkg/sql/pgwire/conn.go
@@ -241,10 +241,13 @@ func (c *conn) checkMaxConnections(ctx context.Context, sqlServer *sql.Server) e
 	// First check maxNumNonRootConnections.
 	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.
+		// Check if there is a reason to use in the error message.
+		msg := "cluster connections are limited"
+		if reason := maxNumNonRootConnectionsReason.Get(&sqlServer.GetExecutorConfig().Settings.SV); reason != "" {
+			msg = reason
+		}
 		return c.sendError(ctx, sqlServer.GetExecutorConfig(), errors.WithHintf(
-			pgerror.New(pgcode.TooManyConnections, "cluster connections are limited"),
+			pgerror.Newf(pgcode.TooManyConnections, "%s", msg),
 			"the maximum number of allowed connections is %d",
 			maxNonRootConnectionsValue,
 		))
diff --git a/pkg/sql/pgwire/conn_test.go b/pkg/sql/pgwire/conn_test.go
index 5836b6d50078..94a63c1da814 100644
--- a/pkg/sql/pgwire/conn_test.go
+++ b/pkg/sql/pgwire/conn_test.go
@@ -1946,6 +1946,13 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 		require.NoError(t, err)
 	}
 
+	setMaxNonRootConnectionsReason := func(reason string) {
+		conn, cleanup := openConnWithUserSuccess(rootUser)
+		defer cleanup()
+		_, err := conn.Exec(ctx, "SET CLUSTER SETTING server.max_non_root_connections_per_gateway_reason = $1", reason)
+		require.NoError(t, err)
+	}
+
 	createUser := func(user string, isAdmin bool) {
 		conn, cleanup := openConnWithUserSuccess(rootUser)
 		defer cleanup()
@@ -2053,6 +2060,19 @@ func TestPGWireRejectsNewConnIfTooManyConns(t *testing.T) {
 		nonAdminCleanup()
 		requireConnectionCount(t, 0)
 	})
+
+	t.Run("max_non_root_connections_reason", func(t *testing.T) {
+		setMaxNonRootConnections(0)
+		setMaxNonRootConnectionsReason("foobar")
+		requireConnectionCount(t, 0)
+		_, cleanup, err := openConnWithUser(admin)
+		defer cleanup()
+		require.Error(t, err)
+		var pgErr *pgconn.PgError
+		require.ErrorAs(t, err, &pgErr)
+		require.Equal(t, pgcode.TooManyConnections.String(), pgErr.Code)
+		require.Regexp(t, "foobar", pgErr.Error())
+	})
 }
 
 func TestConnCloseReleasesReservedMem(t *testing.T) {
diff --git a/pkg/sql/pgwire/server.go b/pkg/sql/pgwire/server.go
index 99af0b38c895..838a09ed4e14 100644
--- a/pkg/sql/pgwire/server.go
+++ b/pkg/sql/pgwire/server.go
@@ -105,6 +105,16 @@ var maxNumNonRootConnections = settings.RegisterIntSetting(
 	-1,
 	)
 
+// maxNumNonRootConnectionsReason is used to supplement the error message for connections that denied due to
+// server.max_non_root_connections_per_gateway.
+var maxNumNonRootConnectionsReason = settings.RegisterStringSetting(
+	settings.TenantWritable,
+	"server.max_non_root_connections_per_gateway_reason",
+	"a reason to provide in the error message for connections that are denied due to " +
+		"server.max_external_connections_per_gateway",
+	"",
+	)
+
 const (
 	// ErrSSLRequired is returned when a client attempts to connect to a
 	// secure server in cleartext.