fix(proxy): bound connection memory and admission

pkg/proxy/config.go

@@ -48,10 +48,32 @@ var (
 	defaultAuthTimeout = time.Second * 10
 	// The default value of TSL connect timeout.
 	defaultTLSConnectTimeout = time.Second * 10
+	// The default deadline for receiving a client login packet. It applies only
+	// before authentication; established tunnel traffic has no such deadline.
+	defaultClientHandshakeTimeout = time.Second * 10
 	// The default base cooldown of the CN health circuit breaker.
 	defaultCNHealthCheckBaseCooldown = time.Second * 5
 	// The default max cooldown of the CN health circuit breaker.
 	defaultCNHealthCheckMaxCooldown = time.Second * 30
+	// Default connection budgets bound Proxy memory independently of client
+	// cleanup behavior. Per-tenant headroom prevents one tenant from consuming
+	// all capacity while still allowing large connection pools.
+	defaultMaxConnections          = 10000
+	defaultMaxConnectionsPerTenant = 8000
+	// Protocol memory covers the shared allocator used by client and backend
+	// MySQL protocol sessions. Login packets are retained for migration, so
+	// their independent per-connection bound is intentionally small.
+	defaultProtocolMemoryLimit        = toml.ByteSize(1 << 30)
+	defaultClientHandshakePacketLimit = toml.ByteSize(64 << 10)
+	// A protocol 4.1 SSLRequest contains only the 32-byte fixed response
+	// prefix. A usable final login must additionally carry at least a one-byte
+	// username, its NUL terminator, and one auth length/terminator byte. Keep the
+	// configured minimum tied to a complete login rather than the shorter
+	// intermediate TLS request.
+	protocol41SSLRequestPayloadSize   = toml.ByteSize(32)
+	minimumClientHandshakePacketLimit = protocol41SSLRequestPayloadSize + 3
+	// The packet header itself can encode at most (1<<24)-1 payload bytes.
+	maximumClientHandshakePacketLimit = toml.ByteSize((1 << 24) - 1)
 )
 
 type RebalancePolicy int
@@ -92,6 +114,9 @@ type Config struct {
 	AuthTimeout toml.Duration `toml:"auth-timeout" user_setting:"advanced"`
 	// TLSConnectTimeout is the timeout duration when TLS connect to server.
 	TLSConnectTimeout toml.Duration `toml:"tls-connect-timeout" user_setting:"advanced"`
+	// ClientHandshakeTimeout bounds how long an unauthenticated client may hold
+	// a Proxy connection slot while sending its login packet.
+	ClientHandshakeTimeout toml.Duration `toml:"client-handshake-timeout" user_setting:"advanced"`
 
 	// Default is false. With true. Server will support tls.
 	// This value should be ths same with all CN servers, and the name
@@ -112,6 +137,18 @@ type Config struct {
 	InternalCIDRs []string `toml:"internal-cidrs"`
 	// ConnCacheEnabled indicates if the connection cache feature is enabled.
 	ConnCacheEnabled bool `toml:"conn-cache-enabled"`
+	// MaxConnections bounds live client connections owned by one Proxy.
+	MaxConnections int `toml:"max-connections" user_setting:"advanced"`
+	// MaxConnectionsPerTenant bounds one tenant's live client connections.
+	// It must not exceed MaxConnections.
+	MaxConnectionsPerTenant int `toml:"max-connections-per-tenant" user_setting:"advanced"`
+	// ProtocolMemoryLimit bounds buffers allocated through the Proxy's shared
+	// MySQL protocol allocator, including the login packet retained for
+	// connection migration.
+	ProtocolMemoryLimit toml.ByteSize `toml:"protocol-memory-limit" user_setting:"advanced"`
+	// ClientHandshakePacketLimit bounds the login packet retained for routing
+	// and migration for the lifetime of a client connection.
+	ClientHandshakePacketLimit toml.ByteSize `toml:"client-handshake-packet-limit" user_setting:"advanced"`
 
 	// CNHealthCheckDisabled disables the CN health circuit breaker. By default
 	// the breaker is enabled: it temporarily skips CN servers that fail to
@@ -228,6 +265,9 @@ func (c *Config) FillDefault() {
 	if c.TLSConnectTimeout.Duration == 0 {
 		c.TLSConnectTimeout.Duration = defaultTLSConnectTimeout
 	}
+	if c.ClientHandshakeTimeout.Duration == 0 {
+		c.ClientHandshakeTimeout.Duration = defaultClientHandshakeTimeout
+	}
 	if c.RebalanceInterval.Duration == 0 {
 		c.RebalanceInterval.Duration = defaultRebalanceInterval
 	}
@@ -260,11 +300,78 @@ func (c *Config) FillDefault() {
 	if c.CNHealthCheckFailThreshold < 1 {
 		c.CNHealthCheckFailThreshold = defaultCNHealthFailThreshold
 	}
+	if c.MaxConnections == 0 {
+		c.MaxConnections = defaultMaxConnections
+	}
+	if c.MaxConnectionsPerTenant == 0 {
+		c.MaxConnectionsPerTenant = min(defaultMaxConnectionsPerTenant, c.MaxConnections)
+	}
+	if c.ProtocolMemoryLimit == 0 {
+		c.ProtocolMemoryLimit = defaultProtocolMemoryLimit
+	}
+	if c.ClientHandshakePacketLimit == 0 {
+		c.ClientHandshakePacketLimit = defaultClientHandshakePacketLimit
+	}
 }
 
 // Validate validates the configuration of proxy server.
 func (c *Config) Validate() error {
 	noReport := errutil.ContextWithNoReport(context.Background(), true)
+	if c.MaxConnections < 0 {
+		return moerr.NewInternalError(noReport, "proxy max-connections must be positive")
+	}
+	if c.ClientHandshakeTimeout.Duration < 0 {
+		return moerr.NewInternalError(noReport, "proxy client-handshake-timeout must be positive")
+	}
+	if c.MaxConnectionsPerTenant < 0 {
+		return moerr.NewInternalError(noReport, "proxy max-connections-per-tenant must be positive")
+	}
+	if c.MaxConnections > 0 && c.MaxConnectionsPerTenant > c.MaxConnections {
+		return moerr.NewInternalError(noReport,
+			"proxy max-connections-per-tenant must not exceed max-connections")
+	}
+	if c.ProtocolMemoryLimit > toml.ByteSize(^uint64(0)>>1) {
+		return moerr.NewInternalError(noReport, "proxy protocol-memory-limit is too large")
+	}
+	if c.ClientHandshakePacketLimit > 0 &&
+		c.ClientHandshakePacketLimit < minimumClientHandshakePacketLimit {
+		return moerr.NewInternalError(noReport,
+			"proxy client-handshake-packet-limit is smaller than a complete protocol 4.1 login")
+	}
+	if c.ClientHandshakePacketLimit > maximumClientHandshakePacketLimit {
+		return moerr.NewInternalError(noReport,
+			"proxy client-handshake-packet-limit exceeds the MySQL packet payload limit")
+	}
+	if c.ProtocolMemoryLimit > 0 && c.MaxConnections > 0 {
+		cachedSessions := uint64(0)
+		// Plugin routing and connection-cache reuse are mutually exclusive in
+		// newProxyHandler, so reserve cached session buffers only when the cache
+		// can actually be instantiated.
+		if c.ConnCacheEnabled && c.Plugin == nil {
+			cachedSessions = defaultMaxNumTotal
+		}
+		maxConnections := uint64(c.MaxConnections)
+		if maxConnections > (^uint64(0)-cachedSessions)/2 {
+			return moerr.NewInternalError(noReport, "proxy max-connections is too large")
+		}
+		fixedSessions := maxConnections*2 + cachedSessions
+		if fixedSessions > ^uint64(0)/proxyIOSessionBufferSize {
+			return moerr.NewInternalError(noReport, "proxy max-connections is too large")
+		}
+		minimum := fixedSessions * proxyIOSessionBufferSize
+		handshakePacketLimit := uint64(c.ClientHandshakePacketLimit)
+		if handshakePacketLimit == 0 {
+			handshakePacketLimit = uint64(defaultClientHandshakePacketLimit)
+		}
+		if maxConnections > (^uint64(0)-minimum)/handshakePacketLimit {
+			return moerr.NewInternalError(noReport, "proxy max-connections is too large")
+		}
+		minimum += maxConnections * handshakePacketLimit
+		if uint64(c.ProtocolMemoryLimit) < minimum {
+			return moerr.NewInternalError(noReport,
+				"proxy protocol-memory-limit is smaller than configured retained protocol buffers")
+		}
+	}
 	if c.Plugin != nil {
 		if c.Plugin.Backend == "" {
 			return moerr.NewInternalError(noReport, "proxy plugin backend must be set")

pkg/proxy/config_test.go

@@ -18,6 +18,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/matrixorigin/matrixone/pkg/util/toml"
 	"github.com/stretchr/testify/require"
 )
 
@@ -33,6 +34,17 @@ func TestFillDefault(t *testing.T) {
 	require.Equal(t, defaultCNHealthCheckBaseCooldown, c.CNHealthCheckBaseCooldown.Duration)
 	require.Equal(t, defaultCNHealthCheckMaxCooldown, c.CNHealthCheckMaxCooldown.Duration)
 	require.Equal(t, defaultCNHealthFailThreshold, c.CNHealthCheckFailThreshold)
+	require.Equal(t, defaultClientHandshakeTimeout, c.ClientHandshakeTimeout.Duration)
+	require.Equal(t, defaultMaxConnections, c.MaxConnections)
+	require.Equal(t, defaultMaxConnectionsPerTenant, c.MaxConnectionsPerTenant)
+	require.Equal(t, defaultProtocolMemoryLimit, c.ProtocolMemoryLimit)
+	require.Equal(t, defaultClientHandshakePacketLimit, c.ClientHandshakePacketLimit)
+	require.NoError(t, c.Validate(), "default limits must form a valid retained-memory budget")
+
+	c = Config{MaxConnections: 1000}
+	c.FillDefault()
+	require.Equal(t, 1000, c.MaxConnections)
+	require.Equal(t, 1000, c.MaxConnectionsPerTenant)
 }
 
 func TestValidate(t *testing.T) {
@@ -43,6 +55,96 @@ func TestValidate(t *testing.T) {
 	}{{
 		name: "empty",
 		cfg:  Config{},
+	}, {
+		name: "negative client handshake timeout",
+		cfg: Config{
+			ClientHandshakeTimeout: toml.Duration{Duration: -time.Second},
+		},
+		wantErr: true,
+	}, {
+		name: "negative global connection limit",
+		cfg: Config{
+			MaxConnections: -1,
+		},
+		wantErr: true,
+	}, {
+		name: "negative tenant connection limit",
+		cfg: Config{
+			MaxConnectionsPerTenant: -1,
+		},
+		wantErr: true,
+	}, {
+		name: "tenant limit exceeds global limit",
+		cfg: Config{
+			MaxConnections:          10,
+			MaxConnectionsPerTenant: 11,
+		},
+		wantErr: true,
+	}, {
+		name: "protocol memory below fixed buffers",
+		cfg: Config{
+			MaxConnections:          10,
+			MaxConnectionsPerTenant: 10,
+			ProtocolMemoryLimit:     1,
+		},
+		wantErr: true,
+	}, {
+		name: "protocol memory below retained handshake budget",
+		cfg: Config{
+			MaxConnections:          10,
+			MaxConnectionsPerTenant: 10,
+			ProtocolMemoryLimit: toml.ByteSize(
+				10*(2*proxyIOSessionBufferSize+64) - 1,
+			),
+			ClientHandshakePacketLimit: 64,
+		},
+		wantErr: true,
+	}, {
+		name: "protocol memory exactly covers retained buffers",
+		cfg: Config{
+			MaxConnections:          10,
+			MaxConnectionsPerTenant: 10,
+			ProtocolMemoryLimit: toml.ByteSize(
+				10 * (2*proxyIOSessionBufferSize + 64),
+			),
+			ClientHandshakePacketLimit: 64,
+		},
+	}, {
+		name: "connection calculation overflow",
+		cfg: Config{
+			MaxConnections:          int(^uint(0) >> 1),
+			MaxConnectionsPerTenant: int(^uint(0) >> 1),
+			ProtocolMemoryLimit:     1,
+			ConnCacheEnabled:        true,
+		},
+		wantErr: true,
+	}, {
+		name: "handshake packet limit below protocol minimum",
+		cfg: Config{
+			ClientHandshakePacketLimit: minimumClientHandshakePacketLimit - 1,
+		},
+		wantErr: true,
+	}, {
+		name: "handshake packet limit at protocol minimum",
+		cfg: Config{
+			ClientHandshakePacketLimit: minimumClientHandshakePacketLimit,
+		},
+	}, {
+		name: "handshake packet limit above protocol minimum",
+		cfg: Config{
+			ClientHandshakePacketLimit: minimumClientHandshakePacketLimit + 1,
+		},
+	}, {
+		name: "handshake packet limit at protocol maximum",
+		cfg: Config{
+			ClientHandshakePacketLimit: maximumClientHandshakePacketLimit,
+		},
+	}, {
+		name: "handshake packet limit above protocol maximum",
+		cfg: Config{
+			ClientHandshakePacketLimit: maximumClientHandshakePacketLimit + 1,
+		},
+		wantErr: true,
 	}, {
 		name: "plugin enabled but no backend",
 		cfg: Config{
@@ -65,6 +167,21 @@ func TestValidate(t *testing.T) {
 				Timeout: time.Second,
 			},
 		},
+	}, {
+		name: "plugin mode does not reserve disabled connection cache",
+		cfg: Config{
+			MaxConnections:          10,
+			MaxConnectionsPerTenant: 10,
+			ProtocolMemoryLimit: toml.ByteSize(
+				10 * (2*proxyIOSessionBufferSize + 64),
+			),
+			ClientHandshakePacketLimit: 64,
+			ConnCacheEnabled:           true,
+			Plugin: &PluginConfig{
+				Backend: "test",
+				Timeout: time.Second,
+			},
+		},
 	}}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {