From 6a657ab2940e4114b2781f7024ed0c410c1870c7 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 06:03:12 +0800 Subject: [PATCH 01/15] fix(proxy): bound connection memory and admission --- pkg/proxy/client_conn.go | 177 ++++++++++++++++++--- pkg/proxy/client_conn_test.go | 235 ++++++++++++++++++++++++++++ pkg/proxy/config.go | 86 +++++++++++ pkg/proxy/config_test.go | 56 +++++++ pkg/proxy/connection_limit.go | 147 ++++++++++++++++++ pkg/proxy/connection_limit_test.go | 238 +++++++++++++++++++++++++++++ pkg/proxy/handler.go | 80 ++++++++-- pkg/proxy/handler_test.go | 33 ++++ pkg/proxy/handshake.go | 39 ++++- pkg/proxy/ppv2.go | 32 +++- pkg/proxy/router.go | 12 +- pkg/proxy/server.go | 4 +- pkg/proxy/server_conn.go | 28 +++- 13 files changed, 1121 insertions(+), 46 deletions(-) create mode 100644 pkg/proxy/connection_limit.go create mode 100644 pkg/proxy/connection_limit_test.go diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 7780ff99e0419..a094e7ea13be4 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -116,6 +116,49 @@ type migration struct { setVarStmts []string } +type clientConnOption func(*clientConn) + +// absoluteReadDeadlineConn prevents a relative timeout set by an upper layer +// from extending an already established lifecycle deadline. It remains usable +// after the handshake; disable makes subsequent deadline updates transparent. +type absoluteReadDeadlineConn struct { + net.Conn + deadline time.Time + enabled atomic.Bool +} + +func newAbsoluteReadDeadlineConn(conn net.Conn, deadline time.Time) *absoluteReadDeadlineConn { + c := &absoluteReadDeadlineConn{ + Conn: conn, + deadline: deadline, + } + c.enabled.Store(true) + return c +} + +func (c *absoluteReadDeadlineConn) SetReadDeadline(deadline time.Time) error { + if c.enabled.Load() && (deadline.IsZero() || deadline.After(c.deadline)) { + deadline = c.deadline + } + return c.Conn.SetReadDeadline(deadline) +} + +func (c *absoluteReadDeadlineConn) disable() { + c.enabled.Store(false) +} + +func withClientConnAllocator(allocator frontend.Allocator) clientConnOption { + return func(c *clientConn) { + c.sessionAllocator = allocator + } +} + +func withClientConnAdmission(lease *connectionLease) clientConnOption { + return func(c *clientConn) { + c.admission = lease + } +} + // clientConn is the connection between proxy and client. type clientConn struct { ctx context.Context @@ -134,6 +177,9 @@ type clientConn struct { connID uint32 // clientInfo is the information of the client. clientInfo clientInfo + // proxyHeaderReceived prevents repeated PROXY headers from resetting the + // handshake read path or growing its recursion depth. + proxyHeaderReceived bool // haKeeperClient is the client of HAKeeper. haKeeperClient logservice.ClusterHAKeeperClient // moCluster is the CN server cache, which used to filter CN servers @@ -147,6 +193,9 @@ type clientConn struct { tlsConfig *tls.Config // tlsConnectTimeout is the TLS connect timeout value. tlsConnectTimeout time.Duration + // clientHandshakeTimeout bounds only the unauthenticated login read. Its + // deadline is cleared before the connection enters the tunnel data path. + clientHandshakeTimeout time.Duration // ipNetList is the list of ip net, which is parsed from CIDRs. ipNetList []*net.IPNet // queryClient is used to send query request to CN servers. @@ -160,6 +209,12 @@ type clientConn struct { sc ServerConn // connCache is the cache of the connections. connCache ConnCache + // sessionAllocator is shared across the Proxy instead of creating one + // ManagedAllocator for every client connection. + sessionAllocator frontend.Allocator + // admission owns this connection's global and, after login parsing, + // per-tenant capacity. The handler releases it on every exit path. + admission *connectionLease // quit tracks quit/cleanup status. It is shared by quit event path and // EOF/connection-end fallback path. quit struct { @@ -192,6 +247,7 @@ func newClientConn( ipNetList []*net.IPNet, qc qclient.QueryClient, connCache ConnCache, + options ...clientConnOption, ) (ClientConn, error) { var originIP net.IP var port int @@ -215,9 +271,18 @@ func newClientConn( }, ipNetList: ipNetList, // set the connection timeout value. - tlsConnectTimeout: cfg.TLSConnectTimeout.Duration, - queryClient: qc, - connCache: connCache, + tlsConnectTimeout: cfg.TLSConnectTimeout.Duration, + clientHandshakeTimeout: cfg.ClientHandshakeTimeout.Duration, + queryClient: qc, + connCache: connCache, + } + if c.clientHandshakeTimeout == 0 { + c.clientHandshakeTimeout = defaultClientHandshakeTimeout + } + for _, option := range options { + if option != nil { + option(c) + } } c.connID, err = c.genConnID() if err != nil { @@ -230,11 +295,31 @@ func newClientConn( fp.SetDefaultValues() pu := config.NewParameterUnit(&fp, nil, nil, nil) frontend.InitServerLevelVars(cfg.UUID) - frontend.SetSessionAlloc(cfg.UUID, frontend.NewSessionAllocator(pu)) - ios, err := frontend.NewIOSession(c.RawConn(), pu, cfg.UUID) + allocator := c.sessionAllocator + if allocator == nil { + allocator = frontend.NewSessionAllocator(pu) + } + handshakePacketLimit := cfg.ClientHandshakePacketLimit + if handshakePacketLimit == 0 { + handshakePacketLimit = defaultClientHandshakePacketLimit + } + ios, err := frontend.NewIOSessionWithOptions( + c.RawConn(), + pu, + cfg.UUID, + frontend.WithIOSessionBufferSize(proxyIOSessionBufferSize), + frontend.WithIOSessionAllowedPacketSize(int(handshakePacketLimit)), + frontend.WithIOSessionAllocator(allocator), + ) if err != nil { return nil, err } + owned := true + defer func() { + if owned { + _ = ios.Close() + } + }() c.mysqlProto = frontend.NewMysqlClientProtocol(cfg.UUID, c.connID, ios, 0, &fp) if cfg.TLSEnabled { tlsConfig, err := frontend.ConstructTLSConfig( @@ -245,6 +330,7 @@ func newClientConn( c.tlsConfig = tlsConfig } c.migration.setVarStmtMap = make(map[string]struct{}) + owned = false return c, nil } @@ -279,9 +365,30 @@ func (c *clientConn) GetTenant() Tenant { return EmptyTenant } +func rewriteProxyError(err error) (uint16, string, string) { + errorCode, sqlState, msg := frontend.RewriteError(err, "") + if errors.Is(err, errProxyConnectionLimit) { + definition := moerr.MysqlErrorMsgRefer[moerr.ER_CON_COUNT_ERROR] + errorCode = definition.ErrorCode + sqlState = definition.SqlStates[0] + msg = definition.ErrorMsgOrFormat + } else if errors.Is(err, frontend.ErrPacketTooLarge) { + definition := moerr.MysqlErrorMsgRefer[moerr.ER_SERVER_NET_PACKET_TOO_LARGE] + errorCode = definition.ErrorCode + sqlState = definition.SqlStates[0] + msg = definition.ErrorMsgOrFormat + } + return errorCode, sqlState, msg +} + +func isProxyAdmissionError(err error) bool { + return errors.Is(err, errProxyConnectionLimit) || + errors.Is(err, frontend.ErrPacketTooLarge) +} + // SendErrToClient implements the ClientConn interface. func (c *clientConn) SendErrToClient(err error) { - errorCode, sqlState, msg := frontend.RewriteError(err, "") + errorCode, sqlState, msg := rewriteProxyError(err) p := c.mysqlProto.MakeErrPayload(errorCode, sqlState, msg) if err := c.mysqlProto.WritePacket(p); err != nil { c.log.Error("failed to send access error to client", zap.Error(err)) @@ -307,7 +414,17 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { if errors.Is(err, io.EOF) { return nil, err } - c.log.Error("failed to handle Handshake response", zap.Error(err)) + if isProxyAdmissionError(err) { + c.log.Debug("client handshake rejected", zap.Error(err)) + } else { + c.log.Error("failed to handle Handshake response", zap.Error(err)) + } + return nil, err + } + // goetty implements read timeouts with a net.Conn deadline, which + // otherwise survives the handshake and would abort a legitimate long + // query later in the raw tunnel. + if err := c.RawConn().SetReadDeadline(time.Time{}); err != nil { return nil, err } } @@ -776,22 +893,40 @@ func (c *clientConn) sendPacketToClient(r []byte, sc ServerConn) error { // readPacket reads MySQL packets from clients. It is mainly used in // handshake phase. func (c *clientConn) readPacket() (*frontend.Packet, error) { - msg, err := c.conn.Read(goetty.ReadOptions{}) - if err != nil { - return nil, err - } - if proxyAddr, ok := msg.(*ProxyAddr); ok { - if proxyAddr.SourceAddress != nil { - c.clientInfo.originIP = proxyAddr.SourceAddress - c.clientInfo.originPort = proxyAddr.SourcePort + return c.readPacketBefore(time.Now().Add(c.clientHandshakeTimeout)) +} + +func (c *clientConn) readPacketBefore(deadline time.Time) (*frontend.Packet, error) { + for { + remaining := time.Until(deadline) + if remaining <= 0 { + return nil, context.DeadlineExceeded } - return c.readPacket() - } - packet, ok := msg.(*frontend.Packet) - if !ok { - return nil, moerr.NewInternalError(c.ctx, "message is not a Packet") + msg, err := c.conn.Read(goetty.ReadOptions{Timeout: remaining}) + if err != nil { + var packetErr *mysqlPacketDecodeError + if errors.As(err, &packetErr) { + c.mysqlProto.SetSequenceID(packetErr.sequenceID + 1) + } + return nil, err + } + if proxyAddr, ok := msg.(*ProxyAddr); ok { + if c.proxyHeaderReceived { + return nil, moerr.NewInvalidInputNoCtx("duplicate PROXY protocol header") + } + c.proxyHeaderReceived = true + if proxyAddr.SourceAddress != nil { + c.clientInfo.originIP = proxyAddr.SourceAddress + c.clientInfo.originPort = proxyAddr.SourcePort + } + continue + } + packet, ok := msg.(*frontend.Packet) + if !ok { + return nil, moerr.NewInternalError(c.ctx, "message is not a Packet") + } + return packet, nil } - return packet, nil } // nextClientConnID increases baseConnID by 1 and returns the result. diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index d20868bc8d404..8d5f80705298f 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -18,6 +18,7 @@ import ( "context" "encoding/binary" "fmt" + "io" "net" "strings" "sync" @@ -49,6 +50,25 @@ type mockNetConn struct { c net.Conn } +type readDeadlineTrackingConn struct { + net.Conn + mu sync.Mutex + deadline time.Time +} + +func (c *readDeadlineTrackingConn) SetReadDeadline(deadline time.Time) error { + c.mu.Lock() + c.deadline = deadline + c.mu.Unlock() + return c.Conn.SetReadDeadline(deadline) +} + +func (c *readDeadlineTrackingConn) readDeadline() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.deadline +} + func newMockNetConn( localIP string, localPort int, remoteIP string, remotePort int, c net.Conn, ) *mockNetConn { @@ -489,6 +509,28 @@ func TestNewClientConn(t *testing.T) { require.NotNil(t, cc.RawConn()) } +func TestNewClientConnReleasesIOSessionWhenTLSConfigFails(t *testing.T) { + s := goetty.NewIOSession(goetty.WithSessionConn(1, + newMockNetConn("127.0.0.1", 30001, + "127.0.0.1", 30010, nil)), + goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec()))) + allocator := frontend.NewLeakCheckAllocator() + cfg := &Config{ + TLSEnabled: true, + TLSCAFile: "testdata/does-not-exist-ca.pem", + TLSCertFile: "testdata/does-not-exist-cert.pem", + TLSKeyFile: "testdata/does-not-exist-key.pem", + } + cc, err := newClientConn( + context.Background(), cfg, runtime.DefaultRuntime().Logger(), + newCounterSet(), s, nil, nil, nil, nil, nil, nil, nil, + withClientConnAllocator(allocator), + ) + require.Error(t, err) + require.Nil(t, cc) + require.True(t, allocator.CheckBalance()) +} + func makeClientHandshakeResp() []byte { payload := make([]byte, 200) pos := 0 @@ -587,6 +629,199 @@ func TestClientConn_ConnectToBackend(t *testing.T) { }) } +func TestClientConnHandshakeAdmission(t *testing.T) { + limiter := newConnectionLimiter(3, 1) + handleHandshake := func(lease *connectionLease) error { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + client.admission = lease + + local, remote := net.Pipe() + defer remote.Close() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(makeClientHandshakeResp()) + }() + err := client.handleHandshakeResp() + <-writeDone + return err + } + + first, ok := limiter.acquire() + require.True(t, ok) + require.NoError(t, handleHandshake(first)) + + second, ok := limiter.acquire() + require.True(t, ok) + require.ErrorIs(t, handleHandshake(second), errProxyConnectionLimit) + second.release() + + first.release() + third, ok := limiter.acquire() + require.True(t, ok) + require.NoError(t, handleHandshake(third)) + third.release() + require.Equal(t, 0, limiter.total) + require.Empty(t, limiter.byTenant) +} + +func TestClientConnHandshakeTimeout(t *testing.T) { + t.Run("silent client times out", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + client.clientHandshakeTimeout = 20 * time.Millisecond + + local, remote := net.Pipe() + defer remote.Close() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + + err := client.handleHandshakeResp() + require.Error(t, err) + var netErr net.Error + require.ErrorAs(t, err, &netErr) + require.True(t, netErr.Timeout()) + }) + + t.Run("fragmented client cannot renew the absolute deadline", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + client.clientHandshakeTimeout = 100 * time.Millisecond + + local, remote := net.Pipe() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + payload := makeClientHandshakeResp() + ticker := time.NewTicker(20 * time.Millisecond) + defer ticker.Stop() + for i := range payload { + <-ticker.C + _ = remote.SetWriteDeadline(time.Now().Add(40 * time.Millisecond)) + if _, err := remote.Write(payload[i : i+1]); err != nil { + return + } + } + }() + + start := time.Now() + errC := make(chan error, 1) + go func() { + errC <- client.handleHandshakeResp() + }() + select { + case err := <-errC: + require.Error(t, err) + var netErr net.Error + require.ErrorAs(t, err, &netErr) + require.True(t, netErr.Timeout()) + require.Less(t, time.Since(start), 8*client.clientHandshakeTimeout) + case <-time.After(10 * client.clientHandshakeTimeout): + t.Fatal("fragmented handshake outlived its absolute deadline") + } + require.NoError(t, remote.Close()) + <-writeDone + }) + + t.Run("successful handshake clears read deadline", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + client.clientHandshakeTimeout = time.Second + + local, remote := net.Pipe() + defer remote.Close() + tracked := &readDeadlineTrackingConn{Conn: local} + client.conn.UseConn(tracked) + client.mysqlProto.UseConn(tracked) + + peerDone := make(chan struct{}) + go func() { + defer close(peerDone) + header := make([]byte, 4) + if _, err := io.ReadFull(remote, header); err != nil { + return + } + payloadLength := int(header[0]) | int(header[1])<<8 | int(header[2])<<16 + payload := make([]byte, payloadLength) + if _, err := io.ReadFull(remote, payload); err != nil { + return + } + _, _ = remote.Write(makeClientHandshakeResp()) + }() + + _, err := client.BuildConnWithServer("") + require.Error(t, err) // no router is configured after the handshake + <-peerDone + require.Same(t, tracked, client.RawConn()) + require.True(t, tracked.readDeadline().IsZero()) + }) +} + +func TestOversizedHandshakeErrorPreservesPacketSequence(t *testing.T) { + for _, sequenceID := range []uint8{1, 2} { + t.Run(fmt.Sprintf("sequence-%d", sequenceID), func(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + session := goetty.NewIOSession( + goetty.WithSessionConn(1, local), + goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(16), + ))), + ) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cc, err := newClientConn( + ctx, + &Config{ClientHandshakePacketLimit: 16}, + runtime.DefaultRuntime().Logger(), + newCounterSet(), + session, + nil, nil, nil, nil, nil, nil, nil, + ) + require.NoError(t, err) + defer cc.Close() + client := cc.(*clientConn) + client.mysqlProto.SetSequenceID(sequenceID) + + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write([]byte{17, 0, 0, sequenceID}) + }() + _, err = client.readPacketBefore(time.Now().Add(time.Second)) + <-writeDone + require.ErrorIs(t, err, frontend.ErrPacketTooLarge) + + sendDone := make(chan struct{}) + go func() { + defer close(sendDone) + client.SendErrToClient(err) + }() + header := make([]byte, 4) + _, err = io.ReadFull(remote, header) + require.NoError(t, err) + payloadLength := int(header[0]) | int(header[1])<<8 | int(header[2])<<16 + payload := make([]byte, payloadLength) + _, err = io.ReadFull(remote, payload) + require.NoError(t, err) + <-sendDone + require.Equal(t, sequenceID+1, header[3]) + require.Equal(t, byte(0xff), payload[0]) + require.Equal(t, moerr.ER_SERVER_NET_PACKET_TOO_LARGE, binary.LittleEndian.Uint16(payload[1:3])) + }) + } +} + func TestClientConn_ReadPacket(t *testing.T) { defer leaktest.AfterTest(t)() diff --git a/pkg/proxy/config.go b/pkg/proxy/config.go index 78790a24f8ffb..83c527c98a95a 100644 --- a/pkg/proxy/config.go +++ b/pkg/proxy/config.go @@ -48,10 +48,27 @@ 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 SSL request is exactly 32 bytes. The packet header itself + // can encode at most (1<<24)-1 payload bytes. + minimumClientHandshakePacketLimit = toml.ByteSize(32) + maximumClientHandshakePacketLimit = toml.ByteSize((1 << 24) - 1) ) type RebalancePolicy int @@ -92,6 +109,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 +132,17 @@ 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. + 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 +259,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 +294,63 @@ 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 protocol 4.1 handshake") + } + 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 { + maxConnectionsForCalculation := (^uint64(0)/proxyIOSessionBufferSize - defaultMaxNumTotal) / 2 + if uint64(c.MaxConnections) > maxConnectionsForCalculation { + return moerr.NewInternalError(noReport, "proxy max-connections is too large") + } + fixedSessions := uint64(c.MaxConnections) * 2 + if c.ConnCacheEnabled { + fixedSessions += defaultMaxNumTotal + } + minimum := fixedSessions * proxyIOSessionBufferSize + if uint64(c.ProtocolMemoryLimit) < minimum { + return moerr.NewInternalError(noReport, + "proxy protocol-memory-limit is smaller than configured fixed session buffers") + } + } if c.Plugin != nil { if c.Plugin.Backend == "" { return moerr.NewInternalError(noReport, "proxy plugin backend must be set") diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index cbeff84822794..1db0ab122b422 100644 --- a/pkg/proxy/config_test.go +++ b/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,16 @@ 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) + + 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 +54,51 @@ 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: "handshake packet limit below protocol minimum", + cfg: Config{ + ClientHandshakePacketLimit: minimumClientHandshakePacketLimit - 1, + }, + wantErr: true, + }, { + name: "handshake packet limit above protocol maximum", + cfg: Config{ + ClientHandshakePacketLimit: maximumClientHandshakePacketLimit + 1, + }, + wantErr: true, }, { name: "plugin enabled but no backend", cfg: Config{ diff --git a/pkg/proxy/connection_limit.go b/pkg/proxy/connection_limit.go new file mode 100644 index 0000000000000..603ea2b4389a1 --- /dev/null +++ b/pkg/proxy/connection_limit.go @@ -0,0 +1,147 @@ +// Copyright 2021 - 2026 Matrix Origin +// +// 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 proxy + +import ( + "encoding/binary" + "net" + "strings" + "sync" + "time" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" +) + +var errProxyConnectionLimit = moerr.NewInvalidInputNoCtx("proxy connection limit exceeded") + +// connectionLimiter provides two-stage admission. The global slot is acquired +// before allocating per-connection protocol state. The tenant slot is bound +// only after the MySQL login packet reveals the tenant. +type connectionLimiter struct { + mu sync.Mutex + maxTotal int + maxPerTenant int + total int + byTenant map[Tenant]int +} + +type connectionLease struct { + limiter *connectionLimiter + tenant Tenant + bound bool + released bool +} + +func newConnectionLimiter(maxTotal, maxPerTenant int) *connectionLimiter { + return &connectionLimiter{ + maxTotal: maxTotal, + maxPerTenant: maxPerTenant, + byTenant: make(map[Tenant]int), + } +} + +func (l *connectionLimiter) acquire() (*connectionLease, bool) { + if l == nil { + return nil, false + } + l.mu.Lock() + defer l.mu.Unlock() + if l.total >= l.maxTotal { + return nil, false + } + l.total++ + return &connectionLease{limiter: l}, true +} + +func (l *connectionLease) bindTenant(tenant Tenant) bool { + if l == nil || l.limiter == nil { + return false + } + tenant = Tenant(strings.ToLower(string(tenant))) + limiter := l.limiter + limiter.mu.Lock() + defer limiter.mu.Unlock() + if l.released { + return false + } + if l.bound { + return l.tenant == tenant + } + if limiter.byTenant[tenant] >= limiter.maxPerTenant { + return false + } + limiter.byTenant[tenant]++ + l.tenant = tenant + l.bound = true + return true +} + +// release is idempotent so error cleanup and normal cleanup may safely +// converge without counter underflow. +func (l *connectionLease) release() { + if l == nil || l.limiter == nil { + return + } + limiter := l.limiter + limiter.mu.Lock() + defer limiter.mu.Unlock() + if l.released { + return + } + l.released = true + limiter.total-- + if !l.bound { + return + } + remaining := limiter.byTenant[l.tenant] - 1 + if remaining == 0 { + delete(limiter.byTenant, l.tenant) + } else { + limiter.byTenant[l.tenant] = remaining + } +} + +// writeConnectionLimitError writes a protocol-level MySQL error before any +// handshake allocation. The write is bounded because this path exists to shed +// load and must never become another unbounded resource wait. +func writeConnectionLimitError(conn net.Conn) { + if conn == nil { + return + } + definition := moerr.MysqlErrorMsgRefer[moerr.ER_CON_COUNT_ERROR] + message := []byte(definition.ErrorMsgOrFormat) + payloadLength := 1 + 2 + 1 + 5 + len(message) + packet := make([]byte, 4+payloadLength) + packet[0] = byte(payloadLength) + packet[1] = byte(payloadLength >> 8) + packet[2] = byte(payloadLength >> 16) + packet[3] = 0 + packet[4] = 0xff + binary.LittleEndian.PutUint16(packet[5:7], definition.ErrorCode) + packet[7] = '#' + copy(packet[8:13], definition.SqlStates[0]) + copy(packet[13:], message) + + if err := conn.SetWriteDeadline(time.Now().Add(time.Second)); err != nil { + return + } + for len(packet) > 0 { + n, err := conn.Write(packet) + if err != nil || n == 0 { + return + } + packet = packet[n:] + } +} diff --git a/pkg/proxy/connection_limit_test.go b/pkg/proxy/connection_limit_test.go new file mode 100644 index 0000000000000..c0eae4dce3f66 --- /dev/null +++ b/pkg/proxy/connection_limit_test.go @@ -0,0 +1,238 @@ +// Copyright 2021 - 2026 Matrix Origin +// +// 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 proxy + +import ( + "encoding/binary" + "fmt" + "io" + "net" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/frontend" +) + +type deadlineErrorConn struct { + net.Conn +} + +func (c *deadlineErrorConn) SetWriteDeadline(time.Time) error { + return net.ErrClosed +} + +func TestConnectionLimiter(t *testing.T) { + t.Run("global limit and idempotent release", func(t *testing.T) { + limiter := newConnectionLimiter(2, 2) + first, ok := limiter.acquire() + require.True(t, ok) + second, ok := limiter.acquire() + require.True(t, ok) + _, ok = limiter.acquire() + require.False(t, ok) + + first.release() + first.release() + third, ok := limiter.acquire() + require.True(t, ok) + second.release() + third.release() + require.Equal(t, 0, limiter.total) + }) + + t.Run("tenant rejection retains global lease until cleanup", func(t *testing.T) { + limiter := newConnectionLimiter(3, 1) + first, ok := limiter.acquire() + require.True(t, ok) + require.True(t, first.bindTenant("tenant-a")) + require.True(t, first.bindTenant("TENANT-A")) + require.False(t, first.bindTenant("tenant-b")) + + second, ok := limiter.acquire() + require.True(t, ok) + require.False(t, second.bindTenant("tenant-a")) + require.Equal(t, 2, limiter.total) + require.Equal(t, 1, limiter.byTenant[Tenant("tenant-a")]) + + second.release() + first.release() + require.Equal(t, 0, limiter.total) + require.Empty(t, limiter.byTenant) + }) + + t.Run("concurrent acquire never exceeds global bound", func(t *testing.T) { + const limit = 8 + limiter := newConnectionLimiter(limit, limit) + var active atomic.Int64 + var peak atomic.Int64 + var wg sync.WaitGroup + var attempted sync.WaitGroup + start := make(chan struct{}) + release := make(chan struct{}) + attempted.Add(128) + for range 128 { + wg.Add(1) + go func() { + defer wg.Done() + <-start + lease, ok := limiter.acquire() + if !ok { + attempted.Done() + return + } + current := active.Add(1) + for { + old := peak.Load() + if current <= old || peak.CompareAndSwap(old, current) { + break + } + } + attempted.Done() + <-release + active.Add(-1) + lease.release() + }() + } + close(start) + attempted.Wait() + close(release) + wg.Wait() + require.LessOrEqual(t, peak.Load(), int64(limit)) + require.Equal(t, 0, limiter.total) + }) + + t.Run("concurrent tenant binding never exceeds tenant bound", func(t *testing.T) { + const ( + connections = 128 + tenantLimit = 8 + ) + limiter := newConnectionLimiter(connections, tenantLimit) + leases := make([]*connectionLease, connections) + for i := range leases { + var ok bool + leases[i], ok = limiter.acquire() + require.True(t, ok) + } + + start := make(chan struct{}) + var admitted atomic.Int64 + var wg sync.WaitGroup + for _, lease := range leases { + wg.Add(1) + go func(lease *connectionLease) { + defer wg.Done() + <-start + if lease.bindTenant("TENANT-A") { + admitted.Add(1) + } + }(lease) + } + close(start) + wg.Wait() + + require.Equal(t, int64(tenantLimit), admitted.Load()) + require.Equal(t, tenantLimit, limiter.byTenant[Tenant("tenant-a")]) + for _, lease := range leases { + lease.release() + } + require.Equal(t, 0, limiter.total) + require.Empty(t, limiter.byTenant) + }) + + t.Run("bind racing with release closes exactly once", func(t *testing.T) { + limiter := newConnectionLimiter(1, 1) + for range 256 { + lease, ok := limiter.acquire() + require.True(t, ok) + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + _ = lease.bindTenant("tenant-a") + }() + go func() { + defer wg.Done() + <-start + lease.release() + }() + close(start) + wg.Wait() + lease.release() + require.Equal(t, 0, limiter.total) + require.Empty(t, limiter.byTenant) + } + }) +} + +func TestRewriteProxyError(t *testing.T) { + t.Run("connection limit", func(t *testing.T) { + err := fmt.Errorf("wrapped: %w", errProxyConnectionLimit) + code, state, message := rewriteProxyError(err) + definition := moerr.MysqlErrorMsgRefer[moerr.ER_CON_COUNT_ERROR] + require.Equal(t, definition.ErrorCode, code) + require.Equal(t, definition.SqlStates[0], state) + require.Equal(t, definition.ErrorMsgOrFormat, message) + require.True(t, isProxyAdmissionError(err)) + }) + + t.Run("packet too large", func(t *testing.T) { + err := fmt.Errorf("wrapped: %w", frontend.ErrPacketTooLarge) + code, state, message := rewriteProxyError(err) + definition := moerr.MysqlErrorMsgRefer[moerr.ER_SERVER_NET_PACKET_TOO_LARGE] + require.Equal(t, definition.ErrorCode, code) + require.Equal(t, definition.SqlStates[0], state) + require.Equal(t, definition.ErrorMsgOrFormat, message) + require.True(t, isProxyAdmissionError(err)) + }) + + require.False(t, isProxyAdmissionError(moerr.NewInternalErrorNoCtx("other"))) +} + +func TestWriteConnectionLimitError(t *testing.T) { + server, client := net.Pipe() + defer client.Close() + done := make(chan struct{}) + go func() { + defer close(done) + writeConnectionLimitError(server) + _ = server.Close() + }() + + header := make([]byte, 4) + _, err := io.ReadFull(client, header) + require.NoError(t, err) + payloadLength := int(header[0]) | int(header[1])<<8 | int(header[2])<<16 + require.Zero(t, header[3]) + payload := make([]byte, payloadLength) + _, err = io.ReadFull(client, payload) + require.NoError(t, err) + require.Equal(t, byte(0xff), payload[0]) + require.Equal(t, moerr.ER_CON_COUNT_ERROR, binary.LittleEndian.Uint16(payload[1:3])) + require.Equal(t, "#08004", string(payload[3:9])) + require.Equal(t, "Too many connections", string(payload[9:])) + <-done + + server, client = net.Pipe() + defer server.Close() + defer client.Close() + writeConnectionLimitError(&deadlineErrorConn{Conn: server}) +} diff --git a/pkg/proxy/handler.go b/pkg/proxy/handler.go index 97398ec3fe50b..d53fc8d7bc86f 100644 --- a/pkg/proxy/handler.go +++ b/pkg/proxy/handler.go @@ -30,6 +30,8 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/morpc" "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/common/stopper" + moconfig "github.com/matrixorigin/matrixone/pkg/config" + "github.com/matrixorigin/matrixone/pkg/frontend" "github.com/matrixorigin/matrixone/pkg/logservice" "github.com/matrixorigin/matrixone/pkg/queryservice/client" v2 "github.com/matrixorigin/matrixone/pkg/util/metric/v2" @@ -61,6 +63,11 @@ type handler struct { queryClient client.QueryClient // connCache is the cache of server connections. connCache ConnCache + // sessionAllocator owns all MySQL protocol buffers retained by this Proxy. + // Sharing it avoids one ManagedAllocator (and its arenas) per connection. + sessionAllocator frontend.Allocator + // connectionLimiter bounds aggregate and per-tenant live connections. + connectionLimiter *connectionLimiter } var ErrNoAvailableCNServers = moerr.NewInternalErrorNoCtx("no available CN servers") @@ -75,6 +82,17 @@ func newProxyHandler( haKeeperClient logservice.ProxyHAKeeperClient, test bool, ) (*handler, error) { + protocolMemoryLimit := cfg.ProtocolMemoryLimit + if protocolMemoryLimit == 0 { + protocolMemoryLimit = defaultProtocolMemoryLimit + } + frontendParameters := &moconfig.FrontendParameters{ + GuestMmuLimitation: int64(protocolMemoryLimit), + } + sessionAllocator := frontend.NewSessionAllocator( + moconfig.NewParameterUnit(frontendParameters, nil, nil, nil), + ) + // Create the MO cluster. mc := clusterservice.NewMOCluster(cfg.UUID, haKeeperClient, cfg.Cluster.RefreshInterval.Duration) rt.SetGlobalVariables(runtime.ClusterService, mc) @@ -99,9 +117,11 @@ func newProxyHandler( var ru Router if test { - ru = newRouter(mc, re, re.connManager, false) + ru = newRouter(mc, re, re.connManager, false, + withSessionAllocator(sessionAllocator)) } else { routerOpts := []routeOption{ + withSessionAllocator(sessionAllocator), withConnectTimeout(cfg.ConnectTimeout.Duration), withAuthTimeout(cfg.AuthTimeout.Duration), withCNHealthCheckCooldown( @@ -145,19 +165,32 @@ func newProxyHandler( if err != nil { return nil, err } + maxConnections := cfg.MaxConnections + if maxConnections == 0 { + maxConnections = defaultMaxConnections + } + maxConnectionsPerTenant := cfg.MaxConnectionsPerTenant + if maxConnectionsPerTenant == 0 { + maxConnectionsPerTenant = min(defaultMaxConnectionsPerTenant, maxConnections) + } h := &handler{ - ctx: ctx, - logger: rt.Logger(), - config: cfg, - stopper: st, - moCluster: mc, - counterSet: cs, - router: ru, - rebalancer: re, - haKeeperClient: haKeeperClient, - ipNetList: ipNetList, - sqlWorker: sw, - queryClient: qc, + ctx: ctx, + logger: rt.Logger(), + config: cfg, + stopper: st, + moCluster: mc, + counterSet: cs, + router: ru, + rebalancer: re, + haKeeperClient: haKeeperClient, + ipNetList: ipNetList, + sqlWorker: sw, + queryClient: qc, + sessionAllocator: sessionAllocator, + connectionLimiter: newConnectionLimiter( + maxConnections, + maxConnectionsPerTenant, + ), } if h.config.ConnCacheEnabled && h.config.Plugin == nil { var cacheOpts []connCacheOption @@ -177,11 +210,22 @@ func (h *handler) handle(c goetty.IOSession) error { h.logger.Info("new connection comes", zap.Uint64("session ID", c.ID())) v2.ProxyConnectAcceptedCounter.Inc() h.counterSet.connAccepted.Add(1) + defer func() { + v2.ProxyConnectClosedCounter.Inc() + }() + + admission, ok := h.connectionLimiter.acquire() + if !ok { + v2.ProxyConnectRejectCounter.Inc() + writeConnectionLimitError(c.RawConn()) + return nil + } + defer admission.release() + v2.ProxyConnectionsCurrentGauge.Inc() h.counterSet.connTotal.Add(1) defer func() { v2.ProxyConnectionsCurrentGauge.Dec() - v2.ProxyConnectClosedCounter.Inc() h.counterSet.connTotal.Add(-1) }() @@ -209,6 +253,8 @@ func (h *handler) handle(c goetty.IOSession) error { h.ipNetList, h.queryClient, h.connCache, + withClientConnAllocator(h.sessionAllocator), + withClientConnAdmission(admission), ) if err != nil { h.logger.Error("failed to create client conn", zap.Error(err)) @@ -224,6 +270,12 @@ func (h *handler) handle(c goetty.IOSession) error { if isConnEndErr(err) { return nil } + if isProxyAdmissionError(err) { + v2.ProxyConnectRejectCounter.Inc() + h.logger.Debug("connection rejected during handshake", zap.Error(err)) + cc.SendErrToClient(err) + return nil + } h.logger.Error("failed to create server conn", zap.Error(err)) h.counterSet.updateWithErr(err) cc.SendErrToClient(err) diff --git a/pkg/proxy/handler_test.go b/pkg/proxy/handler_test.go index 0be0d3b999570..1a947f0173b23 100644 --- a/pkg/proxy/handler_test.go +++ b/pkg/proxy/handler_test.go @@ -35,6 +35,7 @@ import ( "testing" "time" + "github.com/fagongzi/goetty/v2" "github.com/go-sql-driver/mysql" "github.com/lni/goutils/leaktest" "github.com/prometheus/client_golang/prometheus/testutil" @@ -301,6 +302,38 @@ func TestHandlerCurrentConnections(t *testing.T) { }) } +func TestHandlerGlobalAdmissionRejectionIsHandled(t *testing.T) { + limiter := newConnectionLimiter(1, 1) + lease, ok := limiter.acquire() + require.True(t, ok) + defer lease.release() + + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + session := goetty.NewIOSession(goetty.WithSessionConn(1, serverConn)) + defer session.Close() + h := &handler{ + logger: runtime.DefaultRuntime().Logger(), + counterSet: newCounterSet(), + connectionLimiter: limiter, + } + + readDone := make(chan struct{}) + go func() { + defer close(readDone) + header := make([]byte, 4) + if _, err := io.ReadFull(clientConn, header); err != nil { + return + } + payloadLength := int(header[0]) | int(header[1])<<8 | int(header[2])<<16 + payload := make([]byte, payloadLength) + _, _ = io.ReadFull(clientConn, payload) + }() + + require.NoError(t, h.handle(session)) + <-readDone +} + func TestHandler_HandleErr(t *testing.T) { defer leaktest.AfterTest(t)() currentBefore := testutil.ToFloat64(metricv2.ProxyConnectionsCurrentGauge) diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index 9e47a65fab031..595a5f28f8161 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -19,6 +19,7 @@ import ( "context" "crypto/tls" "encoding/binary" + "time" "github.com/fagongzi/goetty/v2" @@ -34,8 +35,28 @@ func (c *clientConn) writeInitialHandshake() error { // handleHandshakeResp receives login information from client and saves it // in proxy end. func (c *clientConn) handleHandshakeResp() error { + deadline := time.Now().Add(c.clientHandshakeTimeout) + rawConn := c.conn.RawConn() + deadlineConn := newAbsoluteReadDeadlineConn(rawConn, deadline) + c.conn.UseConn(deadlineConn) + c.mysqlProto.UseConn(deadlineConn) + defer func() { + deadlineConn.disable() + // A non-TLS handshake still uses deadlineConn directly, so remove the + // now-disabled wrapper instead of retaining one object per connection. + // After a TLS upgrade, tls.Conn owns the wrapped transport and the + // disabled wrapper remains a transparent part of that stack. + if c.conn.RawConn() == deadlineConn { + c.conn.UseConn(rawConn) + c.mysqlProto.UseConn(rawConn) + } + }() + return c.handleHandshakeRespBefore(deadline) +} + +func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { // The proxy reads login request from client. - pack, err := c.readPacket() + pack, err := c.readPacketBefore(deadline) if err != nil { return err } @@ -52,16 +73,19 @@ func (c *clientConn) handleHandshakeResp() error { return err } if ssl { - if err = c.upgradeToTLS(); err != nil { + if err = c.upgradeToTLS(deadline); err != nil { return err } - return c.handleHandshakeResp() + return c.handleHandshakeRespBefore(deadline) } // parse tenant information from client login request. if err := c.clientInfo.parse(c.mysqlProto.GetUserName()); err != nil { return err } + if c.admission != nil && !c.admission.bindTenant(c.clientInfo.Tenant) { + return errProxyConnectionLimit + } li := &c.clientInfo.labelInfo c.clientInfo.labelInfo = newLabelInfo(c.clientInfo.Tenant, li.Labels) @@ -74,14 +98,19 @@ func (c *clientConn) handleHandshakeResp() error { } // upgradeToTLS upgrades the connection to TLS connection. -func (c *clientConn) upgradeToTLS() error { +func (c *clientConn) upgradeToTLS(handshakeDeadline time.Time) error { if c.tlsConfig == nil { return moerr.NewInternalErrorNoCtx("TLS config is invalid") } + remaining := time.Until(handshakeDeadline) + if remaining <= 0 { + return context.DeadlineExceeded + } + timeout := min(c.tlsConnectTimeout, remaining) // TLS handshake packet from client might have been read into the buffer, use a wrapped conn to // avoid losing handshake packets. tlsConn := tls.Server(c.conn.(goetty.BufferedIOSession).BufferedConn(), c.tlsConfig) - ctx, cancel := context.WithTimeoutCause(context.Background(), c.tlsConnectTimeout, moerr.CauseUpgradeToTLS) + ctx, cancel := context.WithTimeoutCause(context.Background(), timeout, moerr.CauseUpgradeToTLS) defer cancel() if err := tlsConn.HandshakeContext(ctx); err != nil { err = moerr.AttachCause(ctx, err) diff --git a/pkg/proxy/ppv2.go b/pkg/proxy/ppv2.go index 185fbef8062c6..233e6dd4860a2 100644 --- a/pkg/proxy/ppv2.go +++ b/pkg/proxy/ppv2.go @@ -17,12 +17,14 @@ package proxy import ( "bytes" "encoding/binary" + "errors" "io" "net" "github.com/fagongzi/goetty/v2/buf" "github.com/fagongzi/goetty/v2/codec" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/frontend" ) const ( @@ -55,6 +57,22 @@ type proxyProtocolCodec struct { codec.Codec } +// mysqlPacketDecodeError retains the header information needed to produce a +// correctly sequenced MySQL error after the underlying codec rejects a payload +// before constructing a Packet. +type mysqlPacketDecodeError struct { + cause error + sequenceID uint8 +} + +func (e *mysqlPacketDecodeError) Error() string { + return e.cause.Error() +} + +func (e *mysqlPacketDecodeError) Unwrap() error { + return e.cause +} + // ProxyHeaderV2 is the structure of the Proxy Protocol version 2 header. type ProxyHeaderV2 struct { Signature [12]byte @@ -80,7 +98,19 @@ func (c *proxyProtocolCodec) Decode(in *buf.ByteBuf) (interface{}, bool, error) if ok { return proxyAddr, ok, nil } - return c.Codec.Decode(in) + var sequenceID uint8 + hasPacketHeader := in.Readable() >= frontend.PacketHeaderLength + if hasPacketHeader { + sequenceID = in.PeekN(0, frontend.PacketHeaderLength)[3] + } + value, complete, err := c.Codec.Decode(in) + if hasPacketHeader && errors.Is(err, frontend.ErrPacketTooLarge) { + err = &mysqlPacketDecodeError{ + cause: err, + sequenceID: sequenceID, + } + } + return value, complete, err } // Encode implements the Codec interface. diff --git a/pkg/proxy/router.go b/pkg/proxy/router.go index b9ab0750a2f58..c82686187416f 100644 --- a/pkg/proxy/router.go +++ b/pkg/proxy/router.go @@ -116,6 +116,10 @@ type router struct { healthDisabled bool // healthOpts are applied when the default health checker is built. healthOpts []cnHealthOption + + // sessionAllocator is shared by all client and backend protocol sessions + // owned by this Proxy instance. + sessionAllocator frontend.Allocator } type routeOption func(*router) @@ -134,6 +138,12 @@ func withAuthTimeout(t time.Duration) routeOption { } } +func withSessionAllocator(allocator frontend.Allocator) routeOption { + return func(r *router) { + r.sessionAllocator = allocator + } +} + // withCNHealthCheckDisabled disables the CN health circuit breaker. func withCNHealthCheckDisabled() routeOption { return func(r *router) { @@ -358,7 +368,7 @@ func (r *router) connect( cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, accountHealth bool, ) (ServerConn, []byte, error) { // Creates a server connection. - sc, err := newServerConn(cn, t, r.rebalancer, r.connectTimeout) + sc, err := newServerConn(cn, t, r.rebalancer, r.connectTimeout, r.sessionAllocator) if err != nil { // Connection failed, remove the placeholder that was added in selectOne. r.rebalancer.connManager.selectOneFailed(cn.hash, cn.uuid) diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index 7900fe5dec90b..a56da5abef54f 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -118,7 +118,9 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err goetty.WithAppLogger(s.runtime.Logger().RawLogger()), goetty.WithAppHandleSessionFunc(s.handler.handle), goetty.WithAppSessionOptions( - goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec())), + goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(int(config.ClientHandshakePacketLimit)), + ))), goetty.WithSessionLogger(s.runtime.Logger().RawLogger()), ), ) diff --git a/pkg/proxy/server_conn.go b/pkg/proxy/server_conn.go index d35c90f0c71fb..0d820e45a16b4 100644 --- a/pkg/proxy/server_conn.go +++ b/pkg/proxy/server_conn.go @@ -95,8 +95,18 @@ type serverConn struct { var _ ServerConn = (*serverConn)(nil) +// Proxy only needs a small retained buffer for normal MySQL packets. Larger +// packets use frontend.Conn's dynamic path and are released after use. +const proxyIOSessionBufferSize = 16 * 1024 + // newServerConn creates a connection to CN server. -func newServerConn(cn *CNServer, tun *tunnel, r *rebalancer, timeout time.Duration) (ServerConn, error) { +func newServerConn( + cn *CNServer, + tun *tunnel, + r *rebalancer, + timeout time.Duration, + allocators ...frontend.Allocator, +) (ServerConn, error) { var logger *zap.Logger if r != nil && r.logger != nil { logger = r.logger.RawLogger() @@ -123,8 +133,20 @@ func newServerConn(cn *CNServer, tun *tunnel, r *rebalancer, timeout time.Durati fp.SetDefaultValues() pu := config.NewParameterUnit(&fp, nil, nil, nil) frontend.InitServerLevelVars(cn.uuid) - frontend.SetSessionAlloc(cn.uuid, frontend.NewSessionAllocator(pu)) - ios, err := frontend.NewIOSession(c.RawConn(), pu, cn.uuid) + var allocator frontend.Allocator + if len(allocators) > 0 { + allocator = allocators[0] + } + if allocator == nil { + allocator = frontend.NewSessionAllocator(pu) + } + ios, err := frontend.NewIOSessionWithOptions( + c.RawConn(), + pu, + cn.uuid, + frontend.WithIOSessionBufferSize(proxyIOSessionBufferSize), + frontend.WithIOSessionAllocator(allocator), + ) if err != nil { return nil, err } From 9510a63d465885e74676c180adfcc5544e8dddfe Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 13:35:21 +0800 Subject: [PATCH 02/15] fix(proxy): await fragmented proxy protocol headers --- pkg/proxy/ppv2.go | 48 ++++++++++++++++------ pkg/proxy/ppv2_test.go | 90 ++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 122 insertions(+), 16 deletions(-) diff --git a/pkg/proxy/ppv2.go b/pkg/proxy/ppv2.go index 233e6dd4860a2..bb42e9c9db1af 100644 --- a/pkg/proxy/ppv2.go +++ b/pkg/proxy/ppv2.go @@ -91,10 +91,13 @@ type ProxyAddr struct { // Decode implements the Codec interface. func (c *proxyProtocolCodec) Decode(in *buf.ByteBuf) (interface{}, bool, error) { - proxyAddr, ok, err := parseProxyHeaderV2(in) + proxyAddr, ok, needMore, err := parseProxyHeaderV2(in) if err != nil { return nil, false, err } + if needMore { + return nil, false, nil + } if ok { return proxyAddr, ok, nil } @@ -120,21 +123,40 @@ func (c *proxyProtocolCodec) Encode(data interface{}, out *buf.ByteBuf, writer i // parseProxyHeader read potential proxy protocol v2 header from the stream // ref: https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt -func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, error) { - // Read the Proxy Protocol header. +func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, bool, error) { + // A stream read may stop anywhere in the signature, fixed header, or + // variable-length body. Keep input buffered while every available signature + // byte still matches; only a mismatch proves that this is a MySQL packet and + // allows the caller to delegate to the SQL codec. + prefixLength := min(in.Readable(), len(ProxyProtocolV2Signature)) + if prefixLength > 0 { + prefix := in.PeekN(0, prefixLength) + for i := range prefix { + if prefix[i] != ProxyProtocolV2Signature[i] { + return nil, false, false, nil + } + } + } if in.Readable() < ProxyHeaderLength { - return nil, false, nil + return nil, false, true, nil } headerBytes := in.PeekN(0, ProxyHeaderLength) header := &ProxyHeaderV2{} if err := binary.Read(bytes.NewBuffer(headerBytes), binary.BigEndian, header); err != nil { - return nil, false, nil + return nil, false, false, nil } // verify the signature of the header if string(header.Signature[:]) != ProxyProtocolV2Signature { - return nil, false, nil + return nil, false, false, nil + } + + // Do not consume the fixed header until the complete declared body is + // buffered. Decoders are retried with the same ByteBuf after the next socket + // read, so partial consumption would corrupt the protocol boundary. + if in.Readable() < ProxyHeaderLength+int(header.Length) { + return nil, false, true, nil } // valid proxy header, consume the bytes @@ -152,26 +174,26 @@ func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, error) { // read address from it. body := make([]byte, header.Length) if _, err := io.ReadFull(in, body); err != nil { - return nil, false, moerr.NewInternalErrorNoCtxf("cannot read proxy address, %s", err.Error()) + return nil, false, false, moerr.NewInternalErrorNoCtxf("cannot read proxy address, %s", err.Error()) } bodyBuf := bytes.NewBuffer(body) addr := &ProxyAddr{} switch header.ProtocolFamily { case tcpOverIPv4, udpOverIPv4: if err := readProxyAddr(bodyBuf, ipv4AddrLength, addr); err != nil { - return nil, false, err + return nil, false, false, err } - return addr, true, nil + return addr, true, false, nil case tcpOverIPv6, udpOverIPv6: if err := readProxyAddr(bodyBuf, ipv6AddrLength, addr); err != nil { - return nil, false, err + return nil, false, false, err } - return addr, true, nil + return addr, true, false, nil case unspec, unixStream, unixDatagram: // no address to read - return addr, false, nil + return addr, false, false, nil default: - return nil, false, moerr.NewInternalErrorNoCtxf("unknown protocol family [%x]", header.ProtocolFamily) + return nil, false, false, moerr.NewInternalErrorNoCtxf("unknown protocol family [%x]", header.ProtocolFamily) } } diff --git a/pkg/proxy/ppv2_test.go b/pkg/proxy/ppv2_test.go index 1b1a9b556c93d..ddb93fa94de2d 100644 --- a/pkg/proxy/ppv2_test.go +++ b/pkg/proxy/ppv2_test.go @@ -15,6 +15,7 @@ package proxy import ( + "encoding/binary" "testing" "github.com/fagongzi/goetty/v2/buf" @@ -29,6 +30,87 @@ func TestProxyProtocolOptions(t *testing.T) { } func TestProxyProtocolCodec_Decode(t *testing.T) { + t.Run("fragmented header and body", func(t *testing.T) { + data := buf.NewByteBuf(100) + pp := WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(64 << 10))) + + header := make([]byte, ProxyHeaderLength) + copy(header, ProxyProtocolV2Signature) + header[12] = 0x21 + header[13] = tcpOverIPv4 + binary.BigEndian.PutUint16(header[14:], 12) + body := []byte{ + 10, 11, 12, 13, + 20, 21, 22, 23, + 0x1f, 0x40, + 0x23, 0x28, + } + + fragments := [][]byte{ + header[:4], + header[4:12], + header[12:15], + header[15:], + body[:5], + body[5:], + } + buffered := 0 + for i, fragment := range fragments { + n, err := data.Write(fragment) + require.NoError(t, err) + require.Equal(t, len(fragment), n) + buffered += n + + res, ok, err := pp.Decode(data) + require.NoError(t, err) + if i < len(fragments)-1 { + require.False(t, ok) + require.Nil(t, res) + require.Equal(t, buffered, data.Readable(), "incomplete input must not be consumed") + continue + } + + require.True(t, ok) + addr, ok := res.(*ProxyAddr) + require.True(t, ok) + require.Equal(t, "10.11.12.13", addr.SourceAddress.String()) + require.Equal(t, uint16(8000), addr.SourcePort) + require.Equal(t, "20.21.22.23", addr.TargetAddress.String()) + require.Equal(t, uint16(9000), addr.TargetPort) + require.Zero(t, data.Readable()) + } + }) + + t.Run("all incomplete signature prefixes wait", func(t *testing.T) { + pp := WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(64 << 10))) + for length := 1; length < len(ProxyProtocolV2Signature); length++ { + data := buf.NewByteBuf(100) + _, err := data.Write([]byte(ProxyProtocolV2Signature[:length])) + require.NoError(t, err) + + res, ok, err := pp.Decode(data) + require.NoError(t, err, "prefix length %d", length) + require.False(t, ok, "prefix length %d", length) + require.Nil(t, res, "prefix length %d", length) + require.Equal(t, length, data.Readable(), "prefix length %d", length) + } + }) + + t.Run("mismatched signature delegates to mysql codec", func(t *testing.T) { + data := buf.NewByteBuf(100) + _, err := data.Write([]byte{0x0d, 0x0a, 0x0d, 0x00}) + require.NoError(t, err) + + pp := WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(64 << 10))) + res, ok, err := pp.Decode(data) + require.ErrorIs(t, err, frontend.ErrPacketTooLarge) + require.False(t, ok) + require.Nil(t, res) + }) + t.Run("short header", func(t *testing.T) { data := buf.NewByteBuf(100) n, err := data.Write([]byte("12345")) @@ -170,11 +252,13 @@ func TestProxyProtocolCodec_Decode(t *testing.T) { require.NoError(t, err) require.Equal(t, len(ProxyProtocolV2Signature), n) - // skip 2 bytes - n, err = data.Write([]byte{0, 0}) + // ipv4 with a complete declared body that is too short for its address. + n, err = data.Write([]byte{0, tcpOverIPv4}) require.NoError(t, err) require.Equal(t, 2, n) - data.WriteUint16(33) + data.WriteUint16(4) + _, err = data.Write(make([]byte, 4)) + require.NoError(t, err) pp := &proxyProtocolCodec{} res, ok, err := pp.Decode(data) From e63475e3539e379d6b32053d553c3a365fc61367 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 16:02:29 +0800 Subject: [PATCH 03/15] fix(proxy): close handshake memory ownership --- pkg/proxy/client_conn.go | 19 +++++ pkg/proxy/client_conn_test.go | 129 ++++++++++++++++++++++++++++++++++ pkg/proxy/config.go | 27 +++++-- pkg/proxy/config_test.go | 31 ++++++++ pkg/proxy/handshake.go | 51 ++++++++++++-- pkg/proxy/ppv2.go | 6 +- pkg/proxy/ppv2_test.go | 7 +- 7 files changed, 254 insertions(+), 16 deletions(-) diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index a094e7ea13be4..7b2094c5b19c2 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -173,6 +173,14 @@ type clientConn struct { // handshakePack is a cached info, used in connection migration. // When connection is transferred, we use it to rebuild handshake. handshakePack *frontend.Packet + // handshakePayloadAllocation keeps the exact slice returned by Allocator. + // Packet.Payload may be a shorter view because Allocator only guarantees + // len >= requested size; cleanup must return the original allocation. + handshakePayloadAllocation []byte + // handshakePackRelease releases the retained login payload exactly once. + // The payload is allocated from sessionAllocator and remains immutable after + // the unauthenticated handshake transfers ownership to clientConn. + handshakePackRelease sync.Once // connID records the connection ID. connID uint32 // clientInfo is the information of the client. @@ -299,6 +307,7 @@ func newClientConn( if allocator == nil { allocator = frontend.NewSessionAllocator(pu) } + c.sessionAllocator = allocator handshakePacketLimit := cfg.ClientHandshakePacketLimit if handshakePacketLimit == 0 { handshakePacketLimit = defaultClientHandshakePacketLimit @@ -679,6 +688,16 @@ func (c *clientConn) Close() error { } c.mysqlProto.Close() } + c.handshakePackRelease.Do(func() { + if c.handshakePayloadAllocation == nil || c.sessionAllocator == nil { + return + } + c.sessionAllocator.Free(c.handshakePayloadAllocation) + c.handshakePayloadAllocation = nil + if c.handshakePack != nil { + c.handshakePack.Payload = nil + } + }) return nil } diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 8d5f80705298f..f30009b68651c 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -669,6 +669,135 @@ func TestClientConnHandshakeAdmission(t *testing.T) { require.Empty(t, limiter.byTenant) } +type failAfterAllocator struct { + frontend.Allocator + remaining int +} + +func (a *failAfterAllocator) Alloc(capacity int) ([]byte, error) { + if a.remaining == 0 { + return nil, moerr.NewInternalErrorNoCtx("injected allocator exhaustion") + } + a.remaining-- + return a.Allocator.Alloc(capacity) +} + +func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { + newClient := func(t *testing.T, allocator frontend.Allocator) (*clientConn, goetty.IOSession, net.Conn) { + t.Helper() + local, remote := net.Pipe() + session := goetty.NewIOSession( + goetty.WithSessionConn(1, local), + goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec())), + ) + cc, err := newClientConn( + context.Background(), + &Config{ClientHandshakePacketLimit: defaultClientHandshakePacketLimit}, + runtime.DefaultRuntime().Logger(), + newCounterSet(), + session, + nil, nil, nil, nil, nil, nil, nil, + withClientConnAllocator(allocator), + ) + require.NoError(t, err) + return cc.(*clientConn), session, remote + } + + t.Run("successful login transfers and releases ownership", func(t *testing.T) { + allocator := frontend.NewLeakCheckAllocator() + client, session, remote := newClient(t, allocator) + defer remote.Close() + + response := makeClientHandshakeResp() + payloadLength := int(defaultClientHandshakePacketLimit) - 1 + largeResponse := make([]byte, frontend.PacketHeaderLength+payloadLength) + largeResponse[0] = byte(payloadLength) + largeResponse[1] = byte(payloadLength >> 8) + largeResponse[2] = byte(payloadLength >> 16) + largeResponse[3] = response[3] + copy(largeResponse[frontend.PacketHeaderLength:], response[frontend.PacketHeaderLength:]) + response = largeResponse + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(response) + }() + + require.NoError(t, client.handleHandshakeResp()) + <-writeDone + require.Equal(t, response[frontend.PacketHeaderLength:], client.handshakePack.Payload) + buffered, ok := session.(goetty.BufferedIOSession) + require.True(t, ok) + require.Nil(t, buffered.InBuf().RawBuf(), "handshake-only input buffer must not cross into the tunnel phase") + require.False(t, allocator.CheckBalance(), "live connection owns fixed and handshake buffers") + + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + require.NoError(t, client.Close(), "close must release the retained login exactly once") + require.True(t, allocator.CheckBalance()) + require.NoError(t, session.Close()) + }) + + t.Run("allocator exhaustion keeps cleanup balanced", func(t *testing.T) { + leakCheck := frontend.NewLeakCheckAllocator() + allocator := &failAfterAllocator{Allocator: leakCheck, remaining: 1} + client, session, remote := newClient(t, allocator) + defer remote.Close() + + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(makeClientHandshakeResp()) + }() + + require.ErrorContains(t, client.handleHandshakeResp(), "injected allocator exhaustion") + <-writeDone + require.Nil(t, client.handshakePack) + require.NoError(t, client.Close()) + require.True(t, leakCheck.CheckBalance()) + require.NoError(t, session.Close()) + }) +} + +func TestClientConnRejectsProxyHeaderAfterAddresslessFrame(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + + local, remote := net.Pipe() + defer remote.Close() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + + localHeader := make([]byte, ProxyHeaderLength) + copy(localHeader, ProxyProtocolV2Signature) + localHeader[12] = 0x20 + localHeader[13] = unspec + + addressHeader := make([]byte, ProxyHeaderLength+12) + copy(addressHeader, ProxyProtocolV2Signature) + addressHeader[12] = 0x21 + addressHeader[13] = tcpOverIPv4 + binary.BigEndian.PutUint16(addressHeader[14:], 12) + copy(addressHeader[ProxyHeaderLength:], []byte{ + 10, 0, 0, 1, + 10, 0, 0, 2, + 0x1f, 0x40, + 0x17, 0x71, + }) + + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(append(localHeader, addressHeader...)) + }() + + _, err := client.readPacketBefore(time.Now().Add(time.Second)) + require.ErrorContains(t, err, "duplicate PROXY protocol header") + <-writeDone + require.True(t, client.proxyHeaderReceived) +} + func TestClientConnHandshakeTimeout(t *testing.T) { t.Run("silent client times out", func(t *testing.T) { cc, cleanup := createNewClientConn(t) diff --git a/pkg/proxy/config.go b/pkg/proxy/config.go index 83c527c98a95a..fad3d7147720d 100644 --- a/pkg/proxy/config.go +++ b/pkg/proxy/config.go @@ -138,7 +138,8 @@ type Config struct { // 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. + // 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. @@ -337,18 +338,30 @@ func (c *Config) Validate() error { "proxy client-handshake-packet-limit exceeds the MySQL packet payload limit") } if c.ProtocolMemoryLimit > 0 && c.MaxConnections > 0 { - maxConnectionsForCalculation := (^uint64(0)/proxyIOSessionBufferSize - defaultMaxNumTotal) / 2 - if uint64(c.MaxConnections) > maxConnectionsForCalculation { + cachedSessions := uint64(0) + if c.ConnCacheEnabled { + cachedSessions = defaultMaxNumTotal + } + maxConnections := uint64(c.MaxConnections) + if maxConnections > (^uint64(0)-cachedSessions)/2 { return moerr.NewInternalError(noReport, "proxy max-connections is too large") } - fixedSessions := uint64(c.MaxConnections) * 2 - if c.ConnCacheEnabled { - fixedSessions += defaultMaxNumTotal + 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 fixed session buffers") + "proxy protocol-memory-limit is smaller than configured retained protocol buffers") } } if c.Plugin != nil { diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index 1db0ab122b422..724a1f030a55e 100644 --- a/pkg/proxy/config_test.go +++ b/pkg/proxy/config_test.go @@ -39,6 +39,7 @@ func TestFillDefault(t *testing.T) { 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() @@ -87,6 +88,36 @@ func TestValidate(t *testing.T) { 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{ diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index 595a5f28f8161..0b46be5f92814 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -51,34 +51,75 @@ func (c *clientConn) handleHandshakeResp() error { c.mysqlProto.UseConn(rawConn) } }() - return c.handleHandshakeRespBefore(deadline) + err := c.handleHandshakeRespBefore(deadline) + if err == nil { + // The goetty decoder is used only during the unauthenticated handshake. + // A large login packet grows its input ByteBuf, whose capacity otherwise + // remains attached to every established connection. No later path calls + // IOSession.Read: tunnel traffic reads the raw (possibly TLS) connection. + // Close the phase-owned buffer as soon as the final login is accepted. + if session, ok := c.conn.(goetty.BufferedIOSession); ok { + if input := session.InBuf(); input != nil { + input.Close() + } + } + } + return err } func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { + if c.handshakePack != nil { + return moerr.NewInvalidInputNoCtx("client handshake has already been processed") + } // The proxy reads login request from client. pack, err := c.readPacketBefore(deadline) if err != nil { return err } c.mysqlProto.AddSequenceId(1) - // Save the login packet in client connection, it will be used - // in the future. - c.handshakePack = pack + + // SQLCodec copies the decoded payload out of goetty's input buffer. Move the + // one lifetime-retained copy into the Proxy's shared bounded allocator so + // ProtocolMemoryLimit covers established connections as well as the fixed + // client/backend IO buffers. + payload, err := c.sessionAllocator.Alloc(len(pack.Payload)) + if err != nil { + return err + } + payloadView := payload[:len(pack.Payload)] + copy(payloadView, pack.Payload) + ownedPack := &frontend.Packet{ + Length: pack.Length, + SequenceID: pack.SequenceID, + Payload: payloadView, + } // Parse the login information and returns whether ssl is needed. // Also, we can get connection attributes from client if it sets // some. - ssl, err := c.mysqlProto.HandleHandshake(c.ctx, pack.Payload) + ssl, err := c.mysqlProto.HandleHandshake(c.ctx, ownedPack.Payload) if err != nil { + // HandleHandshake may retain a slice of the payload before a later + // validation fails. Keep clientConn as the cleanup owner until Close. + c.handshakePack = ownedPack + c.handshakePayloadAllocation = payload return err } if ssl { + // An SSLRequest contains no authentication state and is superseded by + // the login packet sent inside TLS, so it must not cross the phase. + c.sessionAllocator.Free(payload) if err = c.upgradeToTLS(deadline); err != nil { return err } return c.handleHandshakeRespBefore(deadline) } + // Ownership transfers only for the final non-SSL login packet. It is used + // to authenticate fresh backend connections during migration. + c.handshakePack = ownedPack + c.handshakePayloadAllocation = payload + // parse tenant information from client login request. if err := c.clientInfo.parse(c.mysqlProto.GetUserName()); err != nil { return err diff --git a/pkg/proxy/ppv2.go b/pkg/proxy/ppv2.go index bb42e9c9db1af..56259e0e41c64 100644 --- a/pkg/proxy/ppv2.go +++ b/pkg/proxy/ppv2.go @@ -190,8 +190,10 @@ func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, bool, error) { } return addr, true, false, nil case unspec, unixStream, unixDatagram: - // no address to read - return addr, false, false, nil + // The frame is still a consumed PROXY protocol message even when it + // carries no IP address. Returning complete=true lets the connection + // state machine record it and reject any later duplicate header. + return addr, true, false, nil default: return nil, false, false, moerr.NewInternalErrorNoCtxf("unknown protocol family [%x]", header.ProtocolFamily) } diff --git a/pkg/proxy/ppv2_test.go b/pkg/proxy/ppv2_test.go index ddb93fa94de2d..60ce1ae73b1cb 100644 --- a/pkg/proxy/ppv2_test.go +++ b/pkg/proxy/ppv2_test.go @@ -152,8 +152,11 @@ func TestProxyProtocolCodec_Decode(t *testing.T) { pp := WithProxyProtocolCodec(frontend.NewSqlCodec()) res, ok, err := pp.Decode(data) require.NoError(t, err) - require.False(t, ok) - require.Nil(t, res) + require.True(t, ok) + addr, ok := res.(*ProxyAddr) + require.True(t, ok) + require.Nil(t, addr.SourceAddress) + require.Zero(t, data.Readable()) }) t.Run("ipv4 address", func(t *testing.T) { From 7a6e5eb735d5005c057063d5970ac84d0459e786 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 18:13:32 +0800 Subject: [PATCH 04/15] fix(proxy): preserve pipelined handshake data --- pkg/proxy/client_conn_test.go | 137 +++++++++++++++++++++++++++++ pkg/proxy/handshake.go | 52 ++++++++--- pkg/proxy/handshake_buffer.go | 110 +++++++++++++++++++++++ pkg/proxy/handshake_buffer_test.go | 88 ++++++++++++++++++ 4 files changed, 377 insertions(+), 10 deletions(-) create mode 100644 pkg/proxy/handshake_buffer.go create mode 100644 pkg/proxy/handshake_buffer_test.go diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index f30009b68651c..2b18163d33db6 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -16,6 +16,7 @@ package proxy import ( "context" + "crypto/tls" "encoding/binary" "fmt" "io" @@ -674,6 +675,18 @@ type failAfterAllocator struct { remaining int } +type failCapacityAllocator struct { + frontend.Allocator + failCapacity int +} + +func (a *failCapacityAllocator) Alloc(capacity int) ([]byte, error) { + if capacity == a.failCapacity { + return nil, moerr.NewInternalErrorNoCtx("injected handoff allocator exhaustion") + } + return a.Allocator.Alloc(capacity) +} + func (a *failAfterAllocator) Alloc(capacity int) ([]byte, error) { if a.remaining == 0 { return nil, moerr.NewInternalErrorNoCtx("injected allocator exhaustion") @@ -757,6 +770,130 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { require.True(t, leakCheck.CheckBalance()) require.NoError(t, session.Close()) }) + + t.Run("coalesced plaintext login hands off trailing packet", func(t *testing.T) { + allocator := frontend.NewLeakCheckAllocator() + client, session, remote := newClient(t, allocator) + defer remote.Close() + + command := []byte{1, 0, 0, 0, 0x0e} + stream := append(append([]byte{}, makeClientHandshakeResp()...), command...) + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(stream) + }() + + require.NoError(t, client.handleHandshakeResp()) + <-writeDone + require.IsType(t, &handshakeBufferedConn{}, client.RawConn()) + require.NoError(t, client.RawConn().SetReadDeadline(time.Now().Add(time.Second))) + got := make([]byte, len(command)) + _, err := io.ReadFull(client.RawConn(), got) + require.NoError(t, err) + require.Equal(t, command, got) + buffered := session.(goetty.BufferedIOSession) + require.Nil(t, buffered.InBuf().RawBuf(), "phase-owned handshake buffer must be released") + + require.NoError(t, client.Close()) + require.NoError(t, session.Close()) + require.True(t, allocator.CheckBalance()) + }) + + t.Run("trailing packet allocator exhaustion keeps cleanup balanced", func(t *testing.T) { + command := []byte{1, 0, 0, 0, 0x0e} + leakCheck := frontend.NewLeakCheckAllocator() + allocator := &failCapacityAllocator{ + Allocator: leakCheck, + failCapacity: len(command), + } + client, session, remote := newClient(t, allocator) + defer remote.Close() + + stream := append(append([]byte{}, makeClientHandshakeResp()...), command...) + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(stream) + }() + + require.ErrorContains(t, client.handleHandshakeResp(), "injected handoff allocator exhaustion") + <-writeDone + require.NoError(t, client.Close()) + require.NoError(t, session.Close()) + require.True(t, leakCheck.CheckBalance()) + }) + + t.Run("coalesced TLS login hands off trailing packet", func(t *testing.T) { + allocator := frontend.NewLeakCheckAllocator() + client, session, remote := newClient(t, allocator) + + cert, err := certGen(t.TempDir()) + require.NoError(t, err) + client.tlsConfig, err = frontend.ConstructTLSConfig( + context.Background(), cert.caFile, cert.certFile, cert.keyFile) + require.NoError(t, err) + client.tlsConnectTimeout = time.Second + + sslRequestPayload := make([]byte, 32) + binary.LittleEndian.PutUint32( + sslRequestPayload, + frontend.DefaultCapability|frontend.CLIENT_SSL, + ) + sslRequest := make([]byte, frontend.PacketHeaderLength+len(sslRequestPayload)) + sslRequest[0] = byte(len(sslRequestPayload)) + sslRequest[3] = 1 + copy(sslRequest[frontend.PacketHeaderLength:], sslRequestPayload) + + command := []byte{1, 0, 0, 0, 0x0e} + stream := append(append([]byte{}, makeClientHandshakeResp()...), command...) + clientDone := make(chan error, 1) + releaseClient := make(chan struct{}) + clientClosed := make(chan struct{}) + var releaseClientOnce sync.Once + releaseClientRead := func() { + releaseClientOnce.Do(func() { close(releaseClient) }) + } + t.Cleanup(func() { + releaseClientRead() + _ = remote.Close() + <-clientClosed + }) + go func() { + defer close(clientClosed) + if _, err := remote.Write(sslRequest); err != nil { + clientDone <- err + return + } + tlsClient := tls.Client(remote, &tls.Config{InsecureSkipVerify: true}) //nolint:gosec + if err := tlsClient.Handshake(); err != nil { + clientDone <- err + return + } + _, err := tlsClient.Write(stream) + clientDone <- err + <-releaseClient + _, _ = io.Copy(io.Discard, tlsClient) + _ = tlsClient.Close() + }() + + require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, <-clientDone) + require.IsType(t, &handshakeBufferedConn{}, client.RawConn()) + require.NoError(t, client.RawConn().SetReadDeadline(time.Now().Add(time.Second))) + got := make([]byte, len(command)) + _, err = io.ReadFull(client.RawConn(), got) + require.NoError(t, err) + require.Equal(t, command, got) + buffered := session.(goetty.BufferedIOSession) + require.Nil(t, buffered.InBuf().RawBuf(), "phase-owned handshake buffer must be released") + + releaseClientRead() + require.NoError(t, client.Close()) + <-clientClosed + require.NoError(t, session.Close()) + require.True(t, allocator.CheckBalance()) + }) } func TestClientConnRejectsProxyHeaderAfterAddresslessFrame(t *testing.T) { diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index 0b46be5f92814..0bf9a8100bba4 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -53,20 +53,52 @@ func (c *clientConn) handleHandshakeResp() error { }() err := c.handleHandshakeRespBefore(deadline) if err == nil { - // The goetty decoder is used only during the unauthenticated handshake. - // A large login packet grows its input ByteBuf, whose capacity otherwise - // remains attached to every established connection. No later path calls - // IOSession.Read: tunnel traffic reads the raw (possibly TLS) connection. - // Close the phase-owned buffer as soon as the final login is accepted. - if session, ok := c.conn.(goetty.BufferedIOSession); ok { - if input := session.InBuf(); input != nil { - input.Close() - } - } + err = c.handoffHandshakeBuffer() } return err } +// handoffHandshakeBuffer closes the phase-owned goetty input buffer without +// dropping bytes that were read past the final login packet. This applies to +// both plaintext and TLS: for TLS, the same ByteBuf is also referenced by the +// transport's BufferedConn and must be drained and reset before it is freed. +func (c *clientConn) handoffHandshakeBuffer() error { + session, ok := c.conn.(goetty.BufferedIOSession) + if !ok { + return nil + } + input := session.InBuf() + if input == nil { + return nil + } + + readable := input.Readable() + if readable == 0 { + input.Reset() + input.Close() + return nil + } + + conn, err := newHandshakeBufferedConn( + c.conn.RawConn(), + input.PeekN(0, readable), + c.sessionAllocator, + ) + if err != nil { + return err + } + + // Ownership moves to conn before the tunnel can observe it. Resetting the + // indices is required because a TLS BufferedConn still holds this ByteBuf; + // a closed buffer with stale non-zero indices would panic on its next Read. + input.Skip(readable) + input.Reset() + input.Close() + c.conn.UseConn(conn) + c.mysqlProto.UseConn(conn) + return nil +} + func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { if c.handshakePack != nil { return moerr.NewInvalidInputNoCtx("client handshake has already been processed") diff --git a/pkg/proxy/handshake_buffer.go b/pkg/proxy/handshake_buffer.go new file mode 100644 index 0000000000000..bb3cafb95868e --- /dev/null +++ b/pkg/proxy/handshake_buffer.go @@ -0,0 +1,110 @@ +// Copyright 2026 Matrix Origin +// +// 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 proxy + +import ( + "net" + "sync" + "sync/atomic" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/frontend" +) + +// handshakeBufferedConn transfers bytes read ahead by the handshake decoder to +// the raw tunnel. The prefix allocation is released as soon as it is consumed, +// or by Close when connection setup fails before the tunnel reads it. +// +// The mutex protects only the local prefix. It is deliberately released before +// reading the underlying connection so Close can still interrupt blocked I/O. +type handshakeBufferedConn struct { + net.Conn + + mu sync.Mutex + prefix []byte + allocation []byte + allocator frontend.Allocator + pending atomic.Bool +} + +func newHandshakeBufferedConn( + conn net.Conn, + source []byte, + allocator frontend.Allocator, +) (*handshakeBufferedConn, error) { + if conn == nil { + return nil, moerr.NewInternalErrorNoCtx("nil connection for handshake buffer handoff") + } + if len(source) == 0 { + return &handshakeBufferedConn{Conn: conn}, nil + } + if allocator == nil { + return nil, moerr.NewInternalErrorNoCtx("nil allocator for handshake buffer handoff") + } + allocation, err := allocator.Alloc(len(source)) + if err != nil { + return nil, err + } + if len(allocation) < len(source) { + allocator.Free(allocation) + return nil, moerr.NewInternalErrorNoCtx("short allocation for handshake buffer handoff") + } + prefix := allocation[:len(source)] + copy(prefix, source) + c := &handshakeBufferedConn{ + Conn: conn, + prefix: prefix, + allocation: allocation, + allocator: allocator, + } + c.pending.Store(true) + return c, nil +} + +func (c *handshakeBufferedConn) Read(dst []byte) (int, error) { + if !c.pending.Load() { + return c.Conn.Read(dst) + } + c.mu.Lock() + if len(c.prefix) > 0 { + n := copy(dst, c.prefix) + c.prefix = c.prefix[n:] + if len(c.prefix) == 0 { + c.releasePrefixLocked() + } + c.mu.Unlock() + return n, nil + } + c.mu.Unlock() + return c.Conn.Read(dst) +} + +func (c *handshakeBufferedConn) Close() error { + c.mu.Lock() + c.releasePrefixLocked() + c.mu.Unlock() + return c.Conn.Close() +} + +func (c *handshakeBufferedConn) releasePrefixLocked() { + if c.allocation == nil { + return + } + c.allocator.Free(c.allocation) + c.prefix = nil + c.allocation = nil + c.allocator = nil + c.pending.Store(false) +} diff --git a/pkg/proxy/handshake_buffer_test.go b/pkg/proxy/handshake_buffer_test.go new file mode 100644 index 0000000000000..70e4c08f2453a --- /dev/null +++ b/pkg/proxy/handshake_buffer_test.go @@ -0,0 +1,88 @@ +// Copyright 2026 Matrix Origin +// +// 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 proxy + +import ( + "io" + "net" + "sync" + "testing" + + "github.com/matrixorigin/matrixone/pkg/frontend" + "github.com/stretchr/testify/require" +) + +func TestHandshakeBufferedConnLifecycle(t *testing.T) { + t.Run("read releases prefix immediately", func(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + allocator := frontend.NewLeakCheckAllocator() + conn, err := newHandshakeBufferedConn(local, []byte("prefix"), allocator) + require.NoError(t, err) + require.False(t, allocator.CheckBalance()) + + got := make([]byte, len("prefix")) + _, err = io.ReadFull(conn, got) + require.NoError(t, err) + require.Equal(t, []byte("prefix"), got) + require.True(t, allocator.CheckBalance()) + require.NoError(t, conn.Close()) + }) + + t.Run("partial read and repeated close release once", func(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + allocator := frontend.NewLeakCheckAllocator() + conn, err := newHandshakeBufferedConn(local, []byte("prefix"), allocator) + require.NoError(t, err) + + got := make([]byte, 2) + _, err = io.ReadFull(conn, got) + require.NoError(t, err) + require.Equal(t, []byte("pr"), got) + require.False(t, allocator.CheckBalance()) + require.NoError(t, conn.Close()) + require.True(t, allocator.CheckBalance()) + _ = conn.Close() + require.True(t, allocator.CheckBalance()) + }) + + t.Run("concurrent read and close keep ownership balanced", func(t *testing.T) { + for range 100 { + local, remote := net.Pipe() + allocator := frontend.NewLeakCheckAllocator() + conn, err := newHandshakeBufferedConn(local, []byte("prefix"), allocator) + require.NoError(t, err) + + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + _, _ = conn.Read(make([]byte, len("prefix"))) + }() + go func() { + defer wg.Done() + <-start + _ = conn.Close() + }() + close(start) + wg.Wait() + require.True(t, allocator.CheckBalance()) + _ = remote.Close() + } + }) +} From bc9478aacda0648b2cfd128dae10bc6597c1cf03 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 18:44:13 +0800 Subject: [PATCH 05/15] fix(proxy): synchronize handshake payload ownership --- pkg/proxy/client_conn.go | 60 ++++++++++++++++++++++++++++------- pkg/proxy/client_conn_test.go | 59 ++++++++++++++++++++++++++++++++++ pkg/proxy/server_conn.go | 9 ++++++ pkg/proxy/server_conn_test.go | 28 ++++++++++++++++ 4 files changed, 145 insertions(+), 11 deletions(-) diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 7b2094c5b19c2..3f83aed861f34 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -181,6 +181,11 @@ type clientConn struct { // The payload is allocated from sessionAllocator and remains immutable after // the unauthenticated handshake transfers ownership to clientConn. handshakePackRelease sync.Once + // handshakePackMu transfers read ownership to every backend handshake. + // Close takes the write side before returning the retained allocation, so a + // concurrent migration cannot observe memory that has already been reused. + handshakePackMu sync.RWMutex + handshakePackClosed bool // connID records the connection ID. connID uint32 // clientInfo is the information of the client. @@ -355,7 +360,14 @@ func (c *clientConn) GetSalt() []byte { // GetHandshakePack implements the ClientConn interface. func (c *clientConn) GetHandshakePack() *frontend.Packet { - return c.handshakePack + c.handshakePackMu.RLock() + defer c.handshakePackMu.RUnlock() + if c.handshakePack == nil { + return nil + } + pack := *c.handshakePack + pack.Payload = append([]byte(nil), c.handshakePack.Payload...) + return &pack } // RawConn implements the ClientConn interface. @@ -489,7 +501,9 @@ func (c *clientConn) sendErr(err error, resp chan<- []byte) { } func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) error { - sc, r, err := c.router.Connect(cn, c.handshakePack, nil) + sc, r, err := c.connectWithHandshakePack(func(pack *frontend.Packet) (ServerConn, []byte, error) { + return c.router.Connect(cn, pack, nil) + }) if err != nil { c.log.Error("failed to connect to backend server", zap.Error(err)) if resp != nil { @@ -689,18 +703,35 @@ func (c *clientConn) Close() error { c.mysqlProto.Close() } c.handshakePackRelease.Do(func() { + c.handshakePackMu.Lock() + defer c.handshakePackMu.Unlock() + c.handshakePackClosed = true if c.handshakePayloadAllocation == nil || c.sessionAllocator == nil { + c.handshakePack = nil return } c.sessionAllocator.Free(c.handshakePayloadAllocation) c.handshakePayloadAllocation = nil - if c.handshakePack != nil { - c.handshakePack.Payload = nil - } + c.handshakePack = nil }) return nil } +// connectWithHandshakePack holds a read lease for the complete synchronous +// Connect call. ServerConn.HandleHandshake guarantees that its worker has +// stopped before Connect returns, including on timeout, so releasing this +// lease is also the terminal point for every payload reader. +func (c *clientConn) connectWithHandshakePack( + connect func(*frontend.Packet) (ServerConn, []byte, error), +) (ServerConn, []byte, error) { + c.handshakePackMu.RLock() + defer c.handshakePackMu.RUnlock() + if c.handshakePackClosed { + return nil, nil, moerr.NewInternalErrorNoCtx("client handshake packet is unavailable") + } + return connect(c.handshakePack) +} + // connectToBackend connect to the real CN server. func (c *clientConn) connectToBackend(prevAdd string) (ServerConn, error) { // Testing path. @@ -795,16 +826,23 @@ skipConnCache: // feedback into the CN health breaker; internal/admin connects use // plain Router.Connect and intentionally do not. if prevAdd == "" { - if rr, ok := c.router.(routeSelectedConnector); ok { - sc, r, err = rr.ConnectRouteSelected(cn, c.handshakePack, c.tun) - } else { - sc, r, err = c.router.Connect(cn, c.handshakePack, c.tun) - } + sc, r, err = c.connectWithHandshakePack( + func(pack *frontend.Packet) (ServerConn, []byte, error) { + if rr, ok := c.router.(routeSelectedConnector); ok { + return rr.ConnectRouteSelected(cn, pack, c.tun) + } + return c.router.Connect(cn, pack, c.tun) + }, + ) } else { // Session transfer / migration must not feed success/failure back // into the global CN breaker. It is control-plane traffic, not a // Route-selected new-session connect. - sc, r, err = c.router.Connect(cn, c.handshakePack, c.tun) + sc, r, err = c.connectWithHandshakePack( + func(pack *frontend.Packet) (ServerConn, []byte, error) { + return c.router.Connect(cn, pack, c.tun) + }, + ) } if err != nil { if isRetryableErr(err) { diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 2b18163d33db6..c67cdac1dd0dc 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -771,6 +771,65 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { require.NoError(t, session.Close()) }) + t.Run("close waits for backend handshake reader", func(t *testing.T) { + allocator := frontend.NewLeakCheckAllocator() + allocation, err := allocator.Alloc(4) + require.NoError(t, err) + copy(allocation, "auth") + client := &clientConn{ + handshakePack: &frontend.Packet{ + Length: 4, + Payload: allocation[:4], + }, + handshakePayloadAllocation: allocation, + sessionAllocator: allocator, + } + + readerEntered := make(chan struct{}) + releaseReader := make(chan struct{}) + readerPayload := make(chan []byte, 1) + readerDone := make(chan struct{}) + go func() { + defer close(readerDone) + _, _, _ = client.connectWithHandshakePack( + func(pack *frontend.Packet) (ServerConn, []byte, error) { + close(readerEntered) + <-releaseReader + readerPayload <- append([]byte(nil), pack.Payload...) + return nil, nil, nil + }, + ) + }() + <-readerEntered + + closeErr := make(chan error, 1) + closeDone := make(chan struct{}) + go func() { + defer close(closeDone) + closeErr <- client.Close() + }() + require.Eventually(t, func() bool { + if client.handshakePackMu.TryRLock() { + client.handshakePackMu.RUnlock() + return false + } + return true + }, time.Second, time.Millisecond, "Close must wait as the exclusive payload owner") + select { + case <-closeDone: + t.Fatal("Close released a handshake payload with an active reader") + default: + } + + close(releaseReader) + <-readerDone + <-closeDone + require.NoError(t, <-closeErr) + require.Equal(t, []byte("auth"), <-readerPayload) + require.Nil(t, client.GetHandshakePack()) + require.True(t, allocator.CheckBalance()) + }) + t.Run("coalesced plaintext login hands off trailing packet", func(t *testing.T) { allocator := frontend.NewLeakCheckAllocator() client, session, remote := newClient(t, allocator) diff --git a/pkg/proxy/server_conn.go b/pkg/proxy/server_conn.go index 0d820e45a16b4..52143df044500 100644 --- a/pkg/proxy/server_conn.go +++ b/pkg/proxy/server_conn.go @@ -222,6 +222,15 @@ func (s *serverConn) HandleHandshake( case ret := <-resultC: return ret.resp, ret.err case <-ctx.Done(): + // A caller may release or reuse handshakeResp as soon as this method + // returns. Close only the transport first, then join the worker before + // returning; calling s.Close here would free mysqlProto buffers while + // the worker may still be using them. net.Conn.Close unblocks both the + // goetty read and frontend write paths. + if conn := s.conn.RawConn(); conn != nil { + _ = conn.Close() + } + <-resultC logutil.Errorf("handshake to cn %s timeout %v, conn ID: %d goId:%d", s.cnServer.addr, timeout, s.connID, goid.Get()) // Return a retryable error with timeout flag set. diff --git a/pkg/proxy/server_conn_test.go b/pkg/proxy/server_conn_test.go index e356486446679..373debe55d63d 100644 --- a/pkg/proxy/server_conn_test.go +++ b/pkg/proxy/server_conn_test.go @@ -705,6 +705,34 @@ func TestServerConn_HandleHandshakeEarlyReadFailureIsNotTimeout(t *testing.T) { require.False(t, isTimeoutErr(err), "early backend close must not be reclassified as timeout/busy") } +func TestServerConn_HandleHandshakeTimeoutStopsWorker(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + require.NoError(t, remote.SetWriteDeadline(time.Now().Add(time.Second))) + session := goetty.NewIOSession( + goetty.WithSessionConn(1, local), + goetty.WithSessionCodec(frontend.NewSqlCodec()), + ) + defer session.Close() + sc := &serverConn{ + cnServer: &CNServer{addr: "pipe"}, + conn: session, + connID: 1, + } + + _, err := sc.HandleHandshake( + &frontend.Packet{Payload: []byte("auth")}, + 20*time.Millisecond, + ) + require.Error(t, err) + require.True(t, isTimeoutErr(err)) + + // HandleHandshake may return only after its worker has stopped. Closing the + // transport is the termination edge that makes that guarantee observable. + _, err = remote.Write([]byte{1}) + require.Error(t, err) +} + func TestFakeCNServer(t *testing.T) { defer leaktest.AfterTest(t) From 78a56c72296ec5b67df5ce8a6d1bbf3ddcb93393 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 22:47:37 +0800 Subject: [PATCH 06/15] fix(proxy): validate complete client handshake phases --- pkg/proxy/client_conn_test.go | 241 ++++++++++++++++++++++++++++++++++ pkg/proxy/config.go | 18 ++- pkg/proxy/config_test.go | 30 +++++ pkg/proxy/handshake.go | 120 +++++++++-------- 4 files changed, 347 insertions(+), 62 deletions(-) diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index c67cdac1dd0dc..75a0d589a212c 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -567,6 +567,247 @@ func makeClientHandshakeResp() []byte { return data } +func makeMinimalClientHandshakeResp(sequenceID byte, capabilities uint32) []byte { + payload := make([]byte, minimumClientHandshakePacketLimit) + binary.LittleEndian.PutUint32(payload, capabilities) + payload[8] = 45 // client charset + payload[protocol41SSLRequestPayloadSize] = 'u' + // The following zero bytes terminate the username and encode an empty auth + // response for CLIENT_SECURE_CONNECTION. + data := make([]byte, frontend.PacketHeaderLength+len(payload)) + data[0] = byte(len(payload)) + data[1] = byte(len(payload) >> 8) + data[2] = byte(len(payload) >> 16) + data[3] = sequenceID + copy(data[frontend.PacketHeaderLength:], payload) + return data +} + +func makeClientSSLRequest(sequenceID byte, capabilities uint32) []byte { + payload := make([]byte, protocol41SSLRequestPayloadSize) + binary.LittleEndian.PutUint32(payload, capabilities|frontend.CLIENT_SSL) + data := make([]byte, frontend.PacketHeaderLength+len(payload)) + data[0] = byte(len(payload)) + data[3] = sequenceID + copy(data[frontend.PacketHeaderLength:], payload) + return data +} + +func TestClientConnHandshakePhases(t *testing.T) { + newClient := func(t *testing.T) (*clientConn, goetty.IOSession, net.Conn, *frontend.LeakCheckAllocator) { + t.Helper() + local, remote := net.Pipe() + session := goetty.NewIOSession( + goetty.WithSessionConn(1, local), + goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(int(minimumClientHandshakePacketLimit)), + ))), + ) + allocator := frontend.NewLeakCheckAllocator() + cc, err := newClientConn( + context.Background(), + &Config{ClientHandshakePacketLimit: minimumClientHandshakePacketLimit}, + runtime.DefaultRuntime().Logger(), + newCounterSet(), + session, + nil, nil, nil, nil, nil, nil, nil, + withClientConnAllocator(allocator), + ) + require.NoError(t, err) + return cc.(*clientConn), session, remote, allocator + } + + assertLogin := func(t *testing.T, client *clientConn) { + t.Helper() + require.Equal(t, frontend.GetDefaultTenant(), string(client.GetTenant())) + require.Equal(t, "u", client.clientInfo.username) + require.NotNil(t, client.handshakePack) + require.Len(t, client.handshakePack.Payload, int(minimumClientHandshakePacketLimit)) + } + + t.Run("plaintext final login", func(t *testing.T) { + client, session, remote, allocator := newClient(t) + t.Cleanup(func() { + _ = client.Close() + _ = session.Close() + _ = remote.Close() + }) + response := makeMinimalClientHandshakeResp( + 1, + frontend.CLIENT_PROTOCOL_41|frontend.CLIENT_SECURE_CONNECTION, + ) + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(response) + }() + + require.NoError(t, client.handleHandshakeResp()) + <-writeDone + assertLogin(t, client) + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + }) + + t.Run("TLS request then final login", func(t *testing.T) { + client, session, remote, allocator := newClient(t) + t.Cleanup(func() { + _ = client.Close() + _ = session.Close() + _ = remote.Close() + }) + cert, err := certGen(t.TempDir()) + require.NoError(t, err) + client.tlsConfig, err = frontend.ConstructTLSConfig( + context.Background(), cert.caFile, cert.certFile, cert.keyFile) + require.NoError(t, err) + client.tlsConnectTimeout = time.Second + + capabilities := uint32(frontend.CLIENT_PROTOCOL_41 | + frontend.CLIENT_SECURE_CONNECTION | + frontend.CLIENT_SSL) + sslRequest := makeClientSSLRequest(1, capabilities) + + clientDone := make(chan error, 1) + go func() { + defer remote.Close() + if _, err := remote.Write(sslRequest); err != nil { + clientDone <- err + return + } + tlsClient := tls.Client(remote, &tls.Config{InsecureSkipVerify: true}) //nolint:gosec + if err := tlsClient.Handshake(); err != nil { + clientDone <- err + return + } + _, err := tlsClient.Write(makeMinimalClientHandshakeResp(2, capabilities)) + clientDone <- err + }() + + require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, <-clientDone) + assertLogin(t, client) + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + }) + + t.Run("reject duplicate TLS request", func(t *testing.T) { + client, session, remote, allocator := newClient(t) + cert, err := certGen(t.TempDir()) + require.NoError(t, err) + client.tlsConfig, err = frontend.ConstructTLSConfig( + context.Background(), cert.caFile, cert.certFile, cert.keyFile) + require.NoError(t, err) + client.tlsConnectTimeout = time.Second + + capabilities := uint32(frontend.CLIENT_PROTOCOL_41 | + frontend.CLIENT_SECURE_CONNECTION | + frontend.CLIENT_SSL) + releaseClient := make(chan struct{}) + var releaseOnce sync.Once + release := func() { + releaseOnce.Do(func() { close(releaseClient) }) + } + t.Cleanup(func() { + release() + _ = client.Close() + _ = session.Close() + _ = remote.Close() + }) + + clientDone := make(chan error, 1) + go func() { + defer remote.Close() + if _, err := remote.Write(makeClientSSLRequest(1, capabilities)); err != nil { + clientDone <- err + return + } + tlsClient := tls.Client(remote, &tls.Config{InsecureSkipVerify: true}) //nolint:gosec + if err := tlsClient.Handshake(); err != nil { + clientDone <- err + return + } + _, err := tlsClient.Write(makeClientSSLRequest(2, capabilities)) + clientDone <- err + <-releaseClient + }() + + serverDone := make(chan error, 1) + go func() { + serverDone <- client.handleHandshakeResp() + }() + select { + case err := <-serverDone: + require.ErrorContains(t, err, "duplicate TLS request") + case <-time.After(time.Second): + t.Fatal("duplicate TLS request was not rejected promptly") + } + require.NoError(t, <-clientDone) + release() + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + }) + + t.Run("TLS negotiation failure releases phase payload", func(t *testing.T) { + client, session, remote, allocator := newClient(t) + cert, err := certGen(t.TempDir()) + require.NoError(t, err) + client.tlsConfig, err = frontend.ConstructTLSConfig( + context.Background(), cert.caFile, cert.certFile, cert.keyFile) + require.NoError(t, err) + client.tlsConnectTimeout = time.Second + t.Cleanup(func() { + _ = client.Close() + _ = session.Close() + _ = remote.Close() + }) + + capabilities := uint32(frontend.CLIENT_PROTOCOL_41 | + frontend.CLIENT_SECURE_CONNECTION | + frontend.CLIENT_SSL) + clientDone := make(chan struct{}) + go func() { + defer close(clientDone) + _, _ = remote.Write(makeClientSSLRequest(1, capabilities)) + // A complete invalid TLS record header makes the server fail without + // relying on a timing-sensitive connection close. + _, _ = remote.Write([]byte{0, 0, 0, 0, 0}) + }() + + require.ErrorContains(t, client.handleHandshakeResp(), "TLS handshake error") + <-clientDone + require.Nil(t, client.handshakePack) + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + }) + + t.Run("final login parse failure transfers cleanup ownership", func(t *testing.T) { + client, session, remote, allocator := newClient(t) + t.Cleanup(func() { + _ = client.Close() + _ = session.Close() + _ = remote.Close() + }) + response := makeMinimalClientHandshakeResp( + 1, + frontend.CLIENT_PROTOCOL_41|frontend.CLIENT_SECURE_CONNECTION, + ) + response[frontend.PacketHeaderLength+8] = 0 // unsupported collation + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + _, _ = remote.Write(response) + }() + + require.ErrorContains(t, client.handleHandshakeResp(), "get collationName") + <-writeDone + require.NotNil(t, client.handshakePack) + require.False(t, allocator.CheckBalance()) + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + }) +} + func TestClientConn_ConnectToBackend(t *testing.T) { defer leaktest.AfterTest(t)() diff --git a/pkg/proxy/config.go b/pkg/proxy/config.go index fad3d7147720d..c3f93ffedf557 100644 --- a/pkg/proxy/config.go +++ b/pkg/proxy/config.go @@ -65,9 +65,14 @@ var ( // their independent per-connection bound is intentionally small. defaultProtocolMemoryLimit = toml.ByteSize(1 << 30) defaultClientHandshakePacketLimit = toml.ByteSize(64 << 10) - // A protocol 4.1 SSL request is exactly 32 bytes. The packet header itself - // can encode at most (1<<24)-1 payload bytes. - minimumClientHandshakePacketLimit = toml.ByteSize(32) + // 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) ) @@ -331,7 +336,7 @@ func (c *Config) Validate() error { if c.ClientHandshakePacketLimit > 0 && c.ClientHandshakePacketLimit < minimumClientHandshakePacketLimit { return moerr.NewInternalError(noReport, - "proxy client-handshake-packet-limit is smaller than a protocol 4.1 handshake") + "proxy client-handshake-packet-limit is smaller than a complete protocol 4.1 login") } if c.ClientHandshakePacketLimit > maximumClientHandshakePacketLimit { return moerr.NewInternalError(noReport, @@ -339,7 +344,10 @@ func (c *Config) Validate() error { } if c.ProtocolMemoryLimit > 0 && c.MaxConnections > 0 { cachedSessions := uint64(0) - if c.ConnCacheEnabled { + // 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) diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index 724a1f030a55e..c21a03d77c329 100644 --- a/pkg/proxy/config_test.go +++ b/pkg/proxy/config_test.go @@ -124,6 +124,21 @@ func TestValidate(t *testing.T) { 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{ @@ -152,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) { diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index 0bf9a8100bba4..f11f195088335 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -103,71 +103,77 @@ func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { if c.handshakePack != nil { return moerr.NewInvalidInputNoCtx("client handshake has already been processed") } - // The proxy reads login request from client. - pack, err := c.readPacketBefore(deadline) - if err != nil { - return err - } - c.mysqlProto.AddSequenceId(1) + tlsEstablished := false + for { + // The proxy reads a protocol phase packet from the client. The first + // packet may be an SSLRequest; the terminal packet must be a login. + pack, err := c.readPacketBefore(deadline) + if err != nil { + return err + } + c.mysqlProto.AddSequenceId(1) + + // SQLCodec copies the decoded payload out of goetty's input buffer. Move + // the one lifetime-retained copy into the Proxy's shared bounded allocator + // so ProtocolMemoryLimit covers established connections as well as the + // fixed client/backend IO buffers. + payload, err := c.sessionAllocator.Alloc(len(pack.Payload)) + if err != nil { + return err + } + payloadView := payload[:len(pack.Payload)] + copy(payloadView, pack.Payload) + ownedPack := &frontend.Packet{ + Length: pack.Length, + SequenceID: pack.SequenceID, + Payload: payloadView, + } - // SQLCodec copies the decoded payload out of goetty's input buffer. Move the - // one lifetime-retained copy into the Proxy's shared bounded allocator so - // ProtocolMemoryLimit covers established connections as well as the fixed - // client/backend IO buffers. - payload, err := c.sessionAllocator.Alloc(len(pack.Payload)) - if err != nil { - return err - } - payloadView := payload[:len(pack.Payload)] - copy(payloadView, pack.Payload) - ownedPack := &frontend.Packet{ - Length: pack.Length, - SequenceID: pack.SequenceID, - Payload: payloadView, - } + // Parse the phase packet and determine whether it requests TLS. + ssl, err := c.mysqlProto.HandleHandshake(c.ctx, ownedPack.Payload) + if err != nil { + // HandleHandshake may retain a slice of the payload before a later + // validation fails. Keep clientConn as the cleanup owner until Close. + c.handshakePack = ownedPack + c.handshakePayloadAllocation = payload + return err + } + if ssl { + // An SSLRequest contains no authentication state and is superseded by + // the login sent inside TLS, so it must not cross the phase. + c.sessionAllocator.Free(payload) + if tlsEstablished { + return moerr.NewInvalidInputNoCtx("duplicate TLS request during client handshake") + } + if err = c.upgradeToTLS(deadline); err != nil { + return err + } + tlsEstablished = true + continue + } - // Parse the login information and returns whether ssl is needed. - // Also, we can get connection attributes from client if it sets - // some. - ssl, err := c.mysqlProto.HandleHandshake(c.ctx, ownedPack.Payload) - if err != nil { - // HandleHandshake may retain a slice of the payload before a later - // validation fails. Keep clientConn as the cleanup owner until Close. + // Ownership transfers only for the final login packet. It is used to + // authenticate fresh backend connections during migration. c.handshakePack = ownedPack c.handshakePayloadAllocation = payload - return err - } - if ssl { - // An SSLRequest contains no authentication state and is superseded by - // the login packet sent inside TLS, so it must not cross the phase. - c.sessionAllocator.Free(payload) - if err = c.upgradeToTLS(deadline); err != nil { + + // Parse tenant information from the final login request. + if err := c.clientInfo.parse(c.mysqlProto.GetUserName()); err != nil { return err } - return c.handleHandshakeRespBefore(deadline) - } - - // Ownership transfers only for the final non-SSL login packet. It is used - // to authenticate fresh backend connections during migration. - c.handshakePack = ownedPack - c.handshakePayloadAllocation = payload - - // parse tenant information from client login request. - if err := c.clientInfo.parse(c.mysqlProto.GetUserName()); err != nil { - return err - } - if c.admission != nil && !c.admission.bindTenant(c.clientInfo.Tenant) { - return errProxyConnectionLimit - } + if c.admission != nil && !c.admission.bindTenant(c.clientInfo.Tenant) { + return errProxyConnectionLimit + } - li := &c.clientInfo.labelInfo - c.clientInfo.labelInfo = newLabelInfo(c.clientInfo.Tenant, li.Labels) + li := &c.clientInfo.labelInfo + c.clientInfo.labelInfo = newLabelInfo(c.clientInfo.Tenant, li.Labels) - c.clientInfo.hash, err = c.clientInfo.getHash() - if err != nil { - return err + c.clientInfo.hash, err = c.clientInfo.getHash() + if err != nil { + return err + } + return nil } - return nil } // upgradeToTLS upgrades the connection to TLS connection. @@ -187,7 +193,7 @@ func (c *clientConn) upgradeToTLS(handshakeDeadline time.Time) error { defer cancel() if err := tlsConn.HandshakeContext(ctx); err != nil { err = moerr.AttachCause(ctx, err) - return moerr.NewInternalErrorf(ctx, "TSL handshake error: %v", err) + return moerr.NewInternalErrorf(ctx, "TLS handshake error: %v", err) } c.conn.UseConn(tlsConn) c.mysqlProto.UseConn(tlsConn) From a4a518d6066ae6b97a911596b85eb6d0f792c27b Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sat, 18 Jul 2026 23:23:30 +0800 Subject: [PATCH 07/15] fix(proxy): bind tenant admission after authentication --- pkg/proxy/client_conn.go | 31 +++++++- pkg/proxy/client_conn_test.go | 132 ++++++++++++++++++++++++++++++++-- pkg/proxy/connection_limit.go | 3 +- pkg/proxy/handshake.go | 4 -- 4 files changed, 157 insertions(+), 13 deletions(-) diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 3f83aed861f34..481a9c533ab2e 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -407,6 +407,17 @@ func isProxyAdmissionError(err error) bool { errors.Is(err, frontend.ErrPacketTooLarge) } +// bindAuthenticatedTenant converts the connection's untrusted login claim +// into authoritative tenant admission only after a backend (or the connection +// cache authenticator) has accepted the credentials. Binding during login +// parsing would let an unauthenticated client reserve another tenant's quota. +func (c *clientConn) bindAuthenticatedTenant() error { + if c.admission == nil || c.admission.bindTenant(c.clientInfo.Tenant) { + return nil + } + return errProxyConnectionLimit +} + // SendErrToClient implements the ClientConn interface. func (c *clientConn) SendErrToClient(err error) { errorCode, sqlState, msg := rewriteProxyError(err) @@ -452,7 +463,11 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { // Step 3, proxy connects to a CN server to build connection. conn, err := c.connectToBackend(prevAddr) if err != nil { - c.log.Error("failed to connect to backend", zap.Error(err)) + if isProxyAdmissionError(err) { + c.log.Debug("backend connection rejected by proxy admission", zap.Error(err)) + } else { + c.log.Error("failed to connect to backend", zap.Error(err)) + } return nil, err } // bind the server connection to the client connection. @@ -758,6 +773,10 @@ func (c *clientConn) connectToBackend(prevAdd string) (ServerConn, error) { } sc = c.connCache.Pop(c.clientInfo.hash, c.connID, c.mysqlProto.GetSalt(), c.mysqlProto.GetAuthResponse()) if sc != nil { + if err := c.bindAuthenticatedTenant(); err != nil { + _ = sc.Close() + return nil, err + } // get the response from the cn server. re := sc.GetConnResponse() if err := c.sendPacketToClient(re, sc); err != nil { @@ -883,6 +902,16 @@ skipConnCache: v2.ProxyConnectCommonFailCounter.Inc() return nil, moerr.NewInternalErrorNoCtx("the response from cn server is not correct") } + // The tenant name parsed from the login packet is only a claim until + // the backend accepts the credentials. Bind its quota after that + // authentication boundary, but before exposing the OK packet to the + // client. Global admission already bounds this pre-auth work. + if isOKPacket(r) { + if err := c.bindAuthenticatedTenant(); err != nil { + _ = sc.Close() + return nil, err + } + } // set the response from the cn server. sc.SetConnResponse(r[4:]) diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 75a0d589a212c..f9f1bf6c00510 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -214,6 +214,7 @@ func (c *mockClientConn) Close() error { return nil } type mockConnCache struct { pushFn func(cacheKey, ServerConn) bool + popFn func(cacheKey, uint32, []byte, []byte) ServerConn } func (m *mockConnCache) Push(key cacheKey, sc ServerConn) bool { @@ -223,9 +224,14 @@ func (m *mockConnCache) Push(key cacheKey, sc ServerConn) bool { return true } -func (m *mockConnCache) Pop(cacheKey, uint32, []byte, []byte) ServerConn { return nil } -func (m *mockConnCache) Count() int { return 0 } -func (m *mockConnCache) Close() error { return nil } +func (m *mockConnCache) Pop(key cacheKey, connID uint32, salt, authResp []byte) ServerConn { + if m.popFn != nil { + return m.popFn(key, connID, salt, authResp) + } + return nil +} +func (m *mockConnCache) Count() int { return 0 } +func (m *mockConnCache) Close() error { return nil } type killTestRouter struct { connectFn func(*CNServer, *frontend.Packet, *tunnel) (ServerConn, []byte, error) @@ -871,7 +877,7 @@ func TestClientConn_ConnectToBackend(t *testing.T) { }) } -func TestClientConnHandshakeAdmission(t *testing.T) { +func TestClientConnHandshakeDoesNotTrustClaimedTenant(t *testing.T) { limiter := newConnectionLimiter(3, 1) handleHandshake := func(lease *connectionLease) error { cc, cleanup := createNewClientConn(t) @@ -899,13 +905,29 @@ func TestClientConnHandshakeAdmission(t *testing.T) { second, ok := limiter.acquire() require.True(t, ok) - require.ErrorIs(t, handleHandshake(second), errProxyConnectionLimit) - second.release() + require.NoError(t, handleHandshake(second), + "an unauthenticated tenant claim must not consume authoritative tenant capacity") + require.Empty(t, limiter.byTenant) + + // Only the connection whose credentials were accepted by a backend may + // consume the tenant slot. A second authenticated connection is rejected. + require.NoError(t, (&clientConn{ + admission: first, + clientInfo: clientInfo{labelInfo: labelInfo{Tenant: "tenant1"}}, + }).bindAuthenticatedTenant()) + require.ErrorIs(t, (&clientConn{ + admission: second, + clientInfo: clientInfo{labelInfo: labelInfo{Tenant: "tenant1"}}, + }).bindAuthenticatedTenant(), errProxyConnectionLimit) first.release() third, ok := limiter.acquire() require.True(t, ok) - require.NoError(t, handleHandshake(third)) + require.NoError(t, (&clientConn{ + admission: third, + clientInfo: clientInfo{labelInfo: labelInfo{Tenant: "tenant1"}}, + }).bindAuthenticatedTenant()) + second.release() third.release() require.Equal(t, 0, limiter.total) require.Empty(t, limiter.byTenant) @@ -1539,6 +1561,7 @@ type shortHandshakeServerConn struct { *mockServerConn closeCount int setConnResponseCount int + connResponse []byte } func (s *shortHandshakeServerConn) Close() error { @@ -1550,6 +1573,10 @@ func (s *shortHandshakeServerConn) SetConnResponse([]byte) { s.setConnResponseCount++ } +func (s *shortHandshakeServerConn) GetConnResponse() []byte { + return s.connResponse +} + type shortHandshakeRouter struct { testRouter response []byte @@ -1773,6 +1800,97 @@ func Test_connectToBackend_ShortHandshakeResponse(t *testing.T) { } } +func Test_connectToBackend_BindsOnlyAuthenticatedTenant(t *testing.T) { + newClient := func(t *testing.T, limiter *connectionLimiter) (*clientConn, *connectionLease, *writeCountingConn, func()) { + t.Helper() + cc, cleanup := createNewClientConn(t) + client := cc.(*clientConn) + client.clientInfo.Tenant = "tenant1" + lease, ok := limiter.acquire() + require.True(t, ok) + client.admission = lease + local, remote := net.Pipe() + writer := &writeCountingConn{Conn: local} + client.mysqlProto.UseConn(writer) + return client, lease, writer, func() { + _ = remote.Close() + _ = local.Close() + lease.release() + cleanup() + } + } + + t.Run("fresh backend success binds before client OK", func(t *testing.T) { + limiter := newConnectionLimiter(2, 1) + client, _, writer, cleanup := newClient(t, limiter) + defer cleanup() + serverConn := &shortHandshakeServerConn{mockServerConn: newMockServerConn(nil)} + client.router = &shortHandshakeRouter{response: makeOKPacket(8), sc: serverConn} + + got, err := client.connectToBackend("") + require.NoError(t, err) + require.Same(t, serverConn, got) + require.Equal(t, 1, limiter.byTenant[Tenant("tenant1")]) + require.Equal(t, 1, writer.writeCount) + }) + + t.Run("fresh backend rejection closes before client OK", func(t *testing.T) { + limiter := newConnectionLimiter(3, 1) + occupied, ok := limiter.acquire() + require.True(t, ok) + require.True(t, occupied.bindTenant("tenant1")) + defer occupied.release() + client, _, writer, cleanup := newClient(t, limiter) + defer cleanup() + serverConn := &shortHandshakeServerConn{mockServerConn: newMockServerConn(nil)} + client.router = &shortHandshakeRouter{response: makeOKPacket(8), sc: serverConn} + + got, err := client.connectToBackend("") + require.ErrorIs(t, err, errProxyConnectionLimit) + require.Nil(t, got) + require.Equal(t, 1, serverConn.closeCount) + require.Zero(t, writer.writeCount, "a backend OK must not escape before tenant admission") + }) + + t.Run("backend auth failure does not consume tenant quota", func(t *testing.T) { + limiter := newConnectionLimiter(2, 1) + client, _, writer, cleanup := newClient(t, limiter) + defer cleanup() + serverConn := &shortHandshakeServerConn{mockServerConn: newMockServerConn(nil)} + client.router = &shortHandshakeRouter{response: makeErrPacket(8), sc: serverConn} + + got, err := client.connectToBackend("") + require.Error(t, err) + require.Nil(t, got) + require.Empty(t, limiter.byTenant) + require.Equal(t, 1, writer.writeCount) + }) + + t.Run("authenticated cache hit obeys tenant admission", func(t *testing.T) { + limiter := newConnectionLimiter(3, 1) + occupied, ok := limiter.acquire() + require.True(t, ok) + require.True(t, occupied.bindTenant("tenant1")) + defer occupied.release() + client, _, writer, cleanup := newClient(t, limiter) + defer cleanup() + serverConn := &shortHandshakeServerConn{ + mockServerConn: newMockServerConn(nil), + connResponse: makeOKPacket(8)[4:], + } + client.router = &testRouter{} + client.connCache = &mockConnCache{popFn: func(cacheKey, uint32, []byte, []byte) ServerConn { + return serverConn + }} + + got, err := client.connectToBackend("") + require.ErrorIs(t, err, errProxyConnectionLimit) + require.Nil(t, got) + require.Equal(t, 1, serverConn.closeCount) + require.Zero(t, writer.writeCount) + }) +} + func Test_connectToBackend_SkipCacheOnMigration(t *testing.T) { rt := runtime.DefaultRuntime() logger := rt.Logger() diff --git a/pkg/proxy/connection_limit.go b/pkg/proxy/connection_limit.go index 603ea2b4389a1..02840d544dc84 100644 --- a/pkg/proxy/connection_limit.go +++ b/pkg/proxy/connection_limit.go @@ -28,7 +28,8 @@ var errProxyConnectionLimit = moerr.NewInvalidInputNoCtx("proxy connection limit // connectionLimiter provides two-stage admission. The global slot is acquired // before allocating per-connection protocol state. The tenant slot is bound -// only after the MySQL login packet reveals the tenant. +// only after a backend or cached authenticator accepts the login credentials; +// the tenant name in an unauthenticated packet is not an admission identity. type connectionLimiter struct { mu sync.Mutex maxTotal int diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index f11f195088335..c657f66eaddc0 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -161,10 +161,6 @@ func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { if err := c.clientInfo.parse(c.mysqlProto.GetUserName()); err != nil { return err } - if c.admission != nil && !c.admission.bindTenant(c.clientInfo.Tenant) { - return errProxyConnectionLimit - } - li := &c.clientInfo.labelInfo c.clientInfo.labelInfo = newLabelInfo(c.clientInfo.Tenant, li.Labels) From fbba8e166739d810dacb7e44124550fcbfde09f0 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sun, 19 Jul 2026 23:41:08 +0800 Subject: [PATCH 08/15] fix(proxy): reserve transient protocol memory --- pkg/proxy/client_conn.go | 71 ++++++- pkg/proxy/client_conn_test.go | 57 ++++-- pkg/proxy/config.go | 31 +-- pkg/proxy/config_test.go | 21 +- pkg/proxy/handler.go | 9 + pkg/proxy/handshake.go | 32 ++- pkg/proxy/handshake_buffer.go | 8 + pkg/proxy/ppv2.go | 44 ++-- pkg/proxy/ppv2_test.go | 44 ++++ pkg/proxy/protocol_memory.go | 322 ++++++++++++++++++++++++++++++ pkg/proxy/protocol_memory_test.go | 121 +++++++++++ pkg/proxy/server.go | 2 +- pkg/proxy/tunnel.go | 6 + pkg/proxy/tunnel_test.go | 9 +- 14 files changed, 710 insertions(+), 67 deletions(-) create mode 100644 pkg/proxy/protocol_memory.go create mode 100644 pkg/proxy/protocol_memory_test.go diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 481a9c533ab2e..84bf37012583b 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -159,6 +159,12 @@ func withClientConnAdmission(lease *connectionLease) clientConnOption { } } +func withClientConnProtocolMemoryLimiter(limiter *protocolMemoryLimiter) clientConnOption { + return func(c *clientConn) { + c.protocolMemoryLimiter = limiter + } +} + // clientConn is the connection between proxy and client. type clientConn struct { ctx context.Context @@ -209,6 +215,9 @@ type clientConn struct { // clientHandshakeTimeout bounds only the unauthenticated login read. Its // deadline is cleared before the connection enters the tunnel data path. clientHandshakeTimeout time.Duration + // clientHandshakePacketLimit also bounds bytes read ahead after the login; + // those bytes move from goetty into the shared allocator at phase handoff. + clientHandshakePacketLimit int // ipNetList is the list of ip net, which is parsed from CIDRs. ipNetList []*net.IPNet // queryClient is used to send query request to CN servers. @@ -228,6 +237,9 @@ type clientConn struct { // admission owns this connection's global and, after login parsing, // per-tenant capacity. The handler releases it on every exit path. admission *connectionLease + // protocolMemoryLimiter admits short-lived overlap that is not part of one + // connection's steady client/backend/login reservation. + protocolMemoryLimiter *protocolMemoryLimiter // quit tracks quit/cleanup status. It is shared by quit event path and // EOF/connection-end fallback path. quit struct { @@ -317,6 +329,7 @@ func newClientConn( if handshakePacketLimit == 0 { handshakePacketLimit = defaultClientHandshakePacketLimit } + c.clientHandshakePacketLimit = int(handshakePacketLimit) ios, err := frontend.NewIOSessionWithOptions( c.RawConn(), pu, @@ -429,6 +442,7 @@ func (c *clientConn) SendErrToClient(err error) { // BuildConnWithServer implements the ClientConn interface. func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { + var transientLease *protocolMemoryLease if prevAddr == "" { // Step 1, proxy write initial handshake to client. if err := c.writeInitialHandshake(); err != nil { @@ -440,7 +454,9 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { } // Step 2, client send handshake response, which is auth request, // to proxy. - if err := c.handleHandshakeResp(); err != nil { + var err error + transientLease, err = c.handleHandshakeResp() + if err != nil { // This connection may come from heartbeat of LB, and receive EOF error // from it. Just return error and do not log it. if errors.Is(err, io.EOF) { @@ -457,9 +473,22 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { // otherwise survives the handshake and would abort a legitimate long // query later in the raw tunnel. if err := c.RawConn().SetReadDeadline(time.Time{}); err != nil { + transientLease.release() + return nil, err + } + } else { + var err error + transientLease, err = c.acquireBackendProtocolMemory(defaultTransferTimeout) + if err != nil { return nil, err } } + leaseOwned := transientLease != nil + defer func() { + if leaseOwned { + transientLease.release() + } + }() // Step 3, proxy connects to a CN server to build connection. conn, err := c.connectToBackend(prevAddr) if err != nil { @@ -470,6 +499,13 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { } return nil, err } + if prevAddr != "" && transientLease != nil { + conn = &protocolMemoryServerConn{ + ServerConn: conn, + lease: transientLease, + } + leaseOwned = false + } // bind the server connection to the client connection. c.sc = conn @@ -516,6 +552,15 @@ func (c *clientConn) sendErr(err error, resp chan<- []byte) { } func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) error { + lease, err := c.acquireBackendProtocolMemory(defaultTransferTimeout) + if err != nil { + if resp != nil { + c.sendErr(err, resp) + } + return err + } + defer lease.release() + sc, r, err := c.connectWithHandshakePack(func(pack *frontend.Packet) (ServerConn, []byte, error) { return c.router.Connect(cn, pack, nil) }) @@ -549,6 +594,30 @@ func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) return nil } +func (c *clientConn) acquireProtocolMemory( + timeout time.Duration, + bytes uint64, +) (*protocolMemoryLease, error) { + if c.protocolMemoryLimiter == nil { + return nil, nil + } + ctx := c.ctx + if ctx == nil { + ctx = context.Background() + } + deadline := time.Now().Add(timeout) + return acquireProtocolMemoryBefore(ctx, c.protocolMemoryLimiter, bytes, deadline) +} + +func (c *clientConn) acquireBackendProtocolMemory( + timeout time.Duration, +) (*protocolMemoryLease, error) { + if c.protocolMemoryLimiter == nil { + return nil, nil + } + return c.acquireProtocolMemory(timeout, c.protocolMemoryLimiter.budget.backendBytes) +} + // KillCurrentBackendConn implements the ClientConn interface. func (c *clientConn) KillCurrentBackendConn(sc ServerConn) error { if sc == nil { diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index f9f1bf6c00510..9fc3682d6e2f5 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -70,6 +70,12 @@ func (c *readDeadlineTrackingConn) readDeadline() time.Time { return c.deadline } +func handleHandshakeRespForTest(c *clientConn) error { + lease, err := c.handleHandshakeResp() + lease.release() + return err +} + func newMockNetConn( localIP string, localPort int, remoteIP string, remotePort int, c net.Conn, ) *mockNetConn { @@ -648,7 +654,7 @@ func TestClientConnHandshakePhases(t *testing.T) { _, _ = remote.Write(response) }() - require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, handleHandshakeRespForTest(client)) <-writeDone assertLogin(t, client) require.NoError(t, client.Close()) @@ -690,7 +696,7 @@ func TestClientConnHandshakePhases(t *testing.T) { clientDone <- err }() - require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, handleHandshakeRespForTest(client)) require.NoError(t, <-clientDone) assertLogin(t, client) require.NoError(t, client.Close()) @@ -740,7 +746,7 @@ func TestClientConnHandshakePhases(t *testing.T) { serverDone := make(chan error, 1) go func() { - serverDone <- client.handleHandshakeResp() + serverDone <- handleHandshakeRespForTest(client) }() select { case err := <-serverDone: @@ -780,7 +786,7 @@ func TestClientConnHandshakePhases(t *testing.T) { _, _ = remote.Write([]byte{0, 0, 0, 0, 0}) }() - require.ErrorContains(t, client.handleHandshakeResp(), "TLS handshake error") + require.ErrorContains(t, handleHandshakeRespForTest(client), "TLS handshake error") <-clientDone require.Nil(t, client.handshakePack) require.NoError(t, client.Close()) @@ -805,7 +811,7 @@ func TestClientConnHandshakePhases(t *testing.T) { _, _ = remote.Write(response) }() - require.ErrorContains(t, client.handleHandshakeResp(), "get collationName") + require.ErrorContains(t, handleHandshakeRespForTest(client), "get collationName") <-writeDone require.NotNil(t, client.handshakePack) require.False(t, allocator.CheckBalance()) @@ -894,7 +900,7 @@ func TestClientConnHandshakeDoesNotTrustClaimedTenant(t *testing.T) { defer close(writeDone) _, _ = remote.Write(makeClientHandshakeResp()) }() - err := client.handleHandshakeResp() + err := handleHandshakeRespForTest(client) <-writeDone return err } @@ -999,7 +1005,7 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { _, _ = remote.Write(response) }() - require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, handleHandshakeRespForTest(client)) <-writeDone require.Equal(t, response[frontend.PacketHeaderLength:], client.handshakePack.Payload) buffered, ok := session.(goetty.BufferedIOSession) @@ -1026,7 +1032,7 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { _, _ = remote.Write(makeClientHandshakeResp()) }() - require.ErrorContains(t, client.handleHandshakeResp(), "injected allocator exhaustion") + require.ErrorContains(t, handleHandshakeRespForTest(client), "injected allocator exhaustion") <-writeDone require.Nil(t, client.handshakePack) require.NoError(t, client.Close()) @@ -1097,6 +1103,11 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { allocator := frontend.NewLeakCheckAllocator() client, session, remote := newClient(t, allocator) defer remote.Close() + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 100, + initialBytes: 100, + }) + client.protocolMemoryLimiter = limiter command := []byte{1, 0, 0, 0, 0x0e} stream := append(append([]byte{}, makeClientHandshakeResp()...), command...) @@ -1106,14 +1117,17 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { _, _ = remote.Write(stream) }() - require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, handleHandshakeRespForTest(client)) <-writeDone require.IsType(t, &handshakeBufferedConn{}, client.RawConn()) + require.Equal(t, int64(100), limiter.used.Load(), + "read-ahead must retain transient ownership after authentication") require.NoError(t, client.RawConn().SetReadDeadline(time.Now().Add(time.Second))) got := make([]byte, len(command)) _, err := io.ReadFull(client.RawConn(), got) require.NoError(t, err) require.Equal(t, command, got) + require.Zero(t, limiter.used.Load(), "draining read-ahead must release transient ownership") buffered := session.(goetty.BufferedIOSession) require.Nil(t, buffered.InBuf().RawBuf(), "phase-owned handshake buffer must be released") @@ -1139,13 +1153,30 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { _, _ = remote.Write(stream) }() - require.ErrorContains(t, client.handleHandshakeResp(), "injected handoff allocator exhaustion") + require.ErrorContains(t, handleHandshakeRespForTest(client), "injected handoff allocator exhaustion") <-writeDone require.NoError(t, client.Close()) require.NoError(t, session.Close()) require.True(t, leakCheck.CheckBalance()) }) + t.Run("oversized read-ahead is rejected before shared allocation", func(t *testing.T) { + allocator := frontend.NewLeakCheckAllocator() + client, session, remote := newClient(t, allocator) + defer remote.Close() + + buffered := session.(goetty.BufferedIOSession) + readAhead := make([]byte, client.clientHandshakePacketLimit+1) + _, err := buffered.InBuf().Write(readAhead) + require.NoError(t, err) + require.ErrorIs(t, client.handoffHandshakeBuffer(nil), frontend.ErrPacketTooLarge) + require.NotNil(t, buffered.InBuf().RawBuf(), "rejected buffer remains phase-owned until close") + + require.NoError(t, client.Close()) + require.NoError(t, session.Close()) + require.True(t, allocator.CheckBalance()) + }) + t.Run("coalesced TLS login hands off trailing packet", func(t *testing.T) { allocator := frontend.NewLeakCheckAllocator() client, session, remote := newClient(t, allocator) @@ -1199,7 +1230,7 @@ func TestClientConnHandshakeMemoryLifecycle(t *testing.T) { _ = tlsClient.Close() }() - require.NoError(t, client.handleHandshakeResp()) + require.NoError(t, handleHandshakeRespForTest(client)) require.NoError(t, <-clientDone) require.IsType(t, &handshakeBufferedConn{}, client.RawConn()) require.NoError(t, client.RawConn().SetReadDeadline(time.Now().Add(time.Second))) @@ -1269,7 +1300,7 @@ func TestClientConnHandshakeTimeout(t *testing.T) { client.conn.UseConn(local) client.mysqlProto.UseConn(local) - err := client.handleHandshakeResp() + err := handleHandshakeRespForTest(client) require.Error(t, err) var netErr net.Error require.ErrorAs(t, err, &netErr) @@ -1304,7 +1335,7 @@ func TestClientConnHandshakeTimeout(t *testing.T) { start := time.Now() errC := make(chan error, 1) go func() { - errC <- client.handleHandshakeResp() + errC <- handleHandshakeRespForTest(client) }() select { case err := <-errC: diff --git a/pkg/proxy/config.go b/pkg/proxy/config.go index c3f93ffedf557..089ee3f4e59cf 100644 --- a/pkg/proxy/config.go +++ b/pkg/proxy/config.go @@ -342,35 +342,8 @@ func (c *Config) Validate() error { 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 _, err := calculateProtocolMemoryBudget(c); err != nil { + return err } if c.Plugin != nil { if c.Plugin.Backend == "" { diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index c21a03d77c329..8e1ef20e6e7c1 100644 --- a/pkg/proxy/config_test.go +++ b/pkg/proxy/config_test.go @@ -48,6 +48,7 @@ func TestFillDefault(t *testing.T) { } func TestValidate(t *testing.T) { + transientBytes := toml.ByteSize(proxyIOSessionBufferSize) tests := []struct { name string cfg Config @@ -100,7 +101,7 @@ func TestValidate(t *testing.T) { }, wantErr: true, }, { - name: "protocol memory exactly covers retained buffers", + name: "protocol memory without transient overlap", cfg: Config{ MaxConnections: 10, MaxConnectionsPerTenant: 10, @@ -109,6 +110,17 @@ func TestValidate(t *testing.T) { ), ClientHandshakePacketLimit: 64, }, + wantErr: true, + }, { + name: "protocol memory exactly covers one transient overlap", + cfg: Config{ + MaxConnections: 10, + MaxConnectionsPerTenant: 10, + ProtocolMemoryLimit: toml.ByteSize( + 10*(2*proxyIOSessionBufferSize+64), + ) + transientBytes, + ClientHandshakePacketLimit: 64, + }, }, { name: "connection calculation overflow", cfg: Config{ @@ -137,6 +149,9 @@ func TestValidate(t *testing.T) { }, { name: "handshake packet limit at protocol maximum", cfg: Config{ + MaxConnections: 1, + MaxConnectionsPerTenant: 1, + ProtocolMemoryLimit: 64 << 20, ClientHandshakePacketLimit: maximumClientHandshakePacketLimit, }, }, { @@ -173,8 +188,8 @@ func TestValidate(t *testing.T) { MaxConnections: 10, MaxConnectionsPerTenant: 10, ProtocolMemoryLimit: toml.ByteSize( - 10 * (2*proxyIOSessionBufferSize + 64), - ), + 10*(2*proxyIOSessionBufferSize+64), + ) + transientBytes, ClientHandshakePacketLimit: 64, ConnCacheEnabled: true, Plugin: &PluginConfig{ diff --git a/pkg/proxy/handler.go b/pkg/proxy/handler.go index d53fc8d7bc86f..1ecc553c1ed03 100644 --- a/pkg/proxy/handler.go +++ b/pkg/proxy/handler.go @@ -68,6 +68,9 @@ type handler struct { sessionAllocator frontend.Allocator // connectionLimiter bounds aggregate and per-tenant live connections. connectionLimiter *connectionLimiter + // protocolMemoryLimiter serializes phase overlap against the headroom left + // after steady per-connection protocol memory has been reserved. + protocolMemoryLimiter *protocolMemoryLimiter } var ErrNoAvailableCNServers = moerr.NewInternalErrorNoCtx("no available CN servers") @@ -82,6 +85,10 @@ func newProxyHandler( haKeeperClient logservice.ProxyHAKeeperClient, test bool, ) (*handler, error) { + protocolMemoryLimiter, err := newProtocolMemoryLimiter(&cfg) + if err != nil { + return nil, err + } protocolMemoryLimit := cfg.ProtocolMemoryLimit if protocolMemoryLimit == 0 { protocolMemoryLimit = defaultProtocolMemoryLimit @@ -191,6 +198,7 @@ func newProxyHandler( maxConnections, maxConnectionsPerTenant, ), + protocolMemoryLimiter: protocolMemoryLimiter, } if h.config.ConnCacheEnabled && h.config.Plugin == nil { var cacheOpts []connCacheOption @@ -255,6 +263,7 @@ func (h *handler) handle(c goetty.IOSession) error { h.connCache, withClientConnAllocator(h.sessionAllocator), withClientConnAdmission(admission), + withClientConnProtocolMemoryLimiter(h.protocolMemoryLimiter), ) if err != nil { h.logger.Error("failed to create client conn", zap.Error(err)) diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index c657f66eaddc0..dc6bc1065b6aa 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -34,7 +34,7 @@ func (c *clientConn) writeInitialHandshake() error { // handleHandshakeResp receives login information from client and saves it // in proxy end. -func (c *clientConn) handleHandshakeResp() error { +func (c *clientConn) handleHandshakeResp() (*protocolMemoryLease, error) { deadline := time.Now().Add(c.clientHandshakeTimeout) rawConn := c.conn.RawConn() deadlineConn := newAbsoluteReadDeadlineConn(rawConn, deadline) @@ -51,18 +51,34 @@ func (c *clientConn) handleHandshakeResp() error { c.mysqlProto.UseConn(rawConn) } }() - err := c.handleHandshakeRespBefore(deadline) - if err == nil { - err = c.handoffHandshakeBuffer() + if err := c.handleHandshakeRespBefore(deadline); err != nil { + return nil, err + } + transientBytes := uint64(0) + if c.protocolMemoryLimiter != nil { + transientBytes = c.protocolMemoryLimiter.budget.initialBytes } - return err + lease, err := acquireProtocolMemoryBefore( + c.ctx, + c.protocolMemoryLimiter, + transientBytes, + deadline, + ) + if err != nil { + return nil, err + } + if err := c.handoffHandshakeBuffer(lease); err != nil { + lease.release() + return nil, err + } + return lease, nil } // handoffHandshakeBuffer closes the phase-owned goetty input buffer without // dropping bytes that were read past the final login packet. This applies to // both plaintext and TLS: for TLS, the same ByteBuf is also referenced by the // transport's BufferedConn and must be drained and reset before it is freed. -func (c *clientConn) handoffHandshakeBuffer() error { +func (c *clientConn) handoffHandshakeBuffer(lease *protocolMemoryLease) error { session, ok := c.conn.(goetty.BufferedIOSession) if !ok { return nil @@ -78,11 +94,15 @@ func (c *clientConn) handoffHandshakeBuffer() error { input.Close() return nil } + if readable > c.clientHandshakePacketLimit { + return frontend.ErrPacketTooLarge + } conn, err := newHandshakeBufferedConn( c.conn.RawConn(), input.PeekN(0, readable), c.sessionAllocator, + lease, ) if err != nil { return err diff --git a/pkg/proxy/handshake_buffer.go b/pkg/proxy/handshake_buffer.go index bb3cafb95868e..87fe1f7295a40 100644 --- a/pkg/proxy/handshake_buffer.go +++ b/pkg/proxy/handshake_buffer.go @@ -36,6 +36,7 @@ type handshakeBufferedConn struct { prefix []byte allocation []byte allocator frontend.Allocator + lease *protocolMemoryLease pending atomic.Bool } @@ -43,6 +44,7 @@ func newHandshakeBufferedConn( conn net.Conn, source []byte, allocator frontend.Allocator, + leases ...*protocolMemoryLease, ) (*handshakeBufferedConn, error) { if conn == nil { return nil, moerr.NewInternalErrorNoCtx("nil connection for handshake buffer handoff") @@ -69,6 +71,10 @@ func newHandshakeBufferedConn( allocation: allocation, allocator: allocator, } + if len(leases) > 0 && leases[0] != nil { + leases[0].retain() + c.lease = leases[0] + } c.pending.Store(true) return c, nil } @@ -103,8 +109,10 @@ func (c *handshakeBufferedConn) releasePrefixLocked() { return } c.allocator.Free(c.allocation) + c.lease.release() c.prefix = nil c.allocation = nil c.allocator = nil + c.lease = nil c.pending.Store(false) } diff --git a/pkg/proxy/ppv2.go b/pkg/proxy/ppv2.go index 56259e0e41c64..09afa4bca90ae 100644 --- a/pkg/proxy/ppv2.go +++ b/pkg/proxy/ppv2.go @@ -49,12 +49,30 @@ const ( ipv6AddrLength = 16 ) -func WithProxyProtocolCodec(c codec.Codec) codec.Codec { - return &proxyProtocolCodec{Codec: c} +// ProxyProtocolCodecOption configures the PROXY protocol decoder. +type ProxyProtocolCodecOption func(*proxyProtocolCodec) + +// WithProxyProtocolMaxBodySize rejects a PROXY v2 body larger than size after +// reading only its fixed header. Non-positive values leave the cap disabled. +func WithProxyProtocolMaxBodySize(size int) ProxyProtocolCodecOption { + return func(c *proxyProtocolCodec) { + c.maxBodySize = size + } +} + +func WithProxyProtocolCodec(c codec.Codec, options ...ProxyProtocolCodecOption) codec.Codec { + pp := &proxyProtocolCodec{Codec: c} + for _, option := range options { + if option != nil { + option(pp) + } + } + return pp } type proxyProtocolCodec struct { codec.Codec + maxBodySize int } // mysqlPacketDecodeError retains the header information needed to produce a @@ -91,7 +109,7 @@ type ProxyAddr struct { // Decode implements the Codec interface. func (c *proxyProtocolCodec) Decode(in *buf.ByteBuf) (interface{}, bool, error) { - proxyAddr, ok, needMore, err := parseProxyHeaderV2(in) + proxyAddr, ok, needMore, err := parseProxyHeaderV2(in, c.maxBodySize) if err != nil { return nil, false, err } @@ -123,7 +141,7 @@ func (c *proxyProtocolCodec) Encode(data interface{}, out *buf.ByteBuf, writer i // parseProxyHeader read potential proxy protocol v2 header from the stream // ref: https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt -func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, bool, error) { +func parseProxyHeaderV2(in *buf.ByteBuf, maxBodySize int) (*ProxyAddr, bool, bool, error) { // A stream read may stop anywhere in the signature, fixed header, or // variable-length body. Keep input buffered while every available signature // byte still matches; only a mismatch proves that this is a MySQL packet and @@ -151,6 +169,11 @@ func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, bool, error) { if string(header.Signature[:]) != ProxyProtocolV2Signature { return nil, false, false, nil } + if maxBodySize > 0 && int(header.Length) > maxBodySize { + // Reject from the fixed header alone. Waiting for an attacker-controlled + // uint16 body would retain it in goetty outside the shared allocator. + return nil, false, false, frontend.ErrPacketTooLarge + } // Do not consume the fixed header until the complete declared body is // buffered. Decoders are retried with the same ByteBuf after the next socket @@ -159,9 +182,6 @@ func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, bool, error) { return nil, false, true, nil } - // valid proxy header, consume the bytes - in.Skip(ProxyHeaderLength) - // According to ppv2: // // When a sender presents a @@ -172,11 +192,11 @@ func parseProxyHeaderV2(in *buf.ByteBuf) (*ProxyAddr, bool, bool, error) { // In practice, the proxy may add extra information (AWS NLB do this) in the proxy header // which makes the body length > address length. So here we consume the entire body and // read address from it. - body := make([]byte, header.Length) - if _, err := io.ReadFull(in, body); err != nil { - return nil, false, false, moerr.NewInternalErrorNoCtxf("cannot read proxy address, %s", err.Error()) - } - bodyBuf := bytes.NewBuffer(body) + // Parse directly from the ByteBuf. The complete-frame check above keeps the + // view valid, and avoiding a second body copy removes the unaccounted peak. + body := in.PeekN(ProxyHeaderLength, int(header.Length)) + bodyBuf := bytes.NewReader(body) + in.Skip(ProxyHeaderLength + int(header.Length)) addr := &ProxyAddr{} switch header.ProtocolFamily { case tcpOverIPv4, udpOverIPv4: diff --git a/pkg/proxy/ppv2_test.go b/pkg/proxy/ppv2_test.go index 60ce1ae73b1cb..1d0c31bd136e1 100644 --- a/pkg/proxy/ppv2_test.go +++ b/pkg/proxy/ppv2_test.go @@ -30,6 +30,50 @@ func TestProxyProtocolOptions(t *testing.T) { } func TestProxyProtocolCodec_Decode(t *testing.T) { + t.Run("oversized body rejected from fixed header", func(t *testing.T) { + data := buf.NewByteBuf(ProxyHeaderLength) + header := make([]byte, ProxyHeaderLength) + copy(header, ProxyProtocolV2Signature) + header[12] = 0x21 + header[13] = unspec + binary.BigEndian.PutUint16(header[14:], 65) + _, err := data.Write(header) + require.NoError(t, err) + + pp := WithProxyProtocolCodec( + frontend.NewSqlCodec(), + WithProxyProtocolMaxBodySize(64), + ) + res, ok, err := pp.Decode(data) + require.ErrorIs(t, err, frontend.ErrPacketTooLarge) + require.False(t, ok) + require.Nil(t, res) + require.Equal(t, ProxyHeaderLength, data.Readable()) + }) + + t.Run("body at configured limit is accepted", func(t *testing.T) { + data := buf.NewByteBuf(ProxyHeaderLength + 64) + header := make([]byte, ProxyHeaderLength) + copy(header, ProxyProtocolV2Signature) + header[12] = 0x21 + header[13] = unspec + binary.BigEndian.PutUint16(header[14:], 64) + _, err := data.Write(header) + require.NoError(t, err) + _, err = data.Write(make([]byte, 64)) + require.NoError(t, err) + + pp := WithProxyProtocolCodec( + frontend.NewSqlCodec(), + WithProxyProtocolMaxBodySize(64), + ) + res, ok, err := pp.Decode(data) + require.NoError(t, err) + require.True(t, ok) + require.IsType(t, &ProxyAddr{}, res) + require.Zero(t, data.Readable()) + }) + t.Run("fragmented header and body", func(t *testing.T) { data := buf.NewByteBuf(100) pp := WithProxyProtocolCodec(frontend.NewSqlCodec( diff --git a/pkg/proxy/protocol_memory.go b/pkg/proxy/protocol_memory.go new file mode 100644 index 0000000000000..3479e29b34138 --- /dev/null +++ b/pkg/proxy/protocol_memory.go @@ -0,0 +1,322 @@ +// Copyright 2026 Matrix Origin +// +// 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 proxy + +import ( + "context" + "errors" + "math" + "sync" + "sync/atomic" + "time" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/frontend" + "golang.org/x/sync/semaphore" +) + +// protocolMemoryBudget separates long-lived memory, which is reserved by +// connection admission, from short-lived phase overlap. The minimum transient +// headroom is deliberately sized for the larger of: +// +// - initial login forwarding: pipelined client prefix + dynamic backend write +// - migration/control: one additional backend session + dynamic login or +// captured control-statement write +// +// The shared allocator remains the byte-level hard limit. Weighted transient +// leases prevent individually valid phase transitions from racing for the same +// unreserved headroom, without charging small normal logins for a 16 MiB write. +type protocolMemoryBudget struct { + steadyBytes uint64 + headroomBytes uint64 + transientBytes uint64 + initialBytes uint64 + backendBytes uint64 + loginDynamicBytes uint64 +} + +func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { + maxConnections := c.MaxConnections + if maxConnections == 0 { + maxConnections = defaultMaxConnections + } + if maxConnections < 0 { + return protocolMemoryBudget{}, moerr.NewInternalErrorNoCtx( + "proxy max-connections must be positive") + } + + handshakeLimit := uint64(c.ClientHandshakePacketLimit) + if handshakeLimit == 0 { + handshakeLimit = uint64(defaultClientHandshakePacketLimit) + } + memoryLimit := uint64(c.ProtocolMemoryLimit) + if memoryLimit == 0 { + memoryLimit = uint64(defaultProtocolMemoryLimit) + } + if memoryLimit > math.MaxInt64 { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + + cachedSessions := uint64(0) + if c.ConnCacheEnabled && c.Plugin == nil { + cachedSessions = defaultMaxNumTotal + } + connections := uint64(maxConnections) + fixedSessions, ok := checkedMulUint64(connections, 2) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + fixedSessions, ok = checkedAddUint64(fixedSessions, cachedSessions) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + fixedBytes, ok := checkedMulUint64(fixedSessions, proxyIOSessionBufferSize) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + retainedLoginBytes, ok := checkedMulUint64(connections, handshakeLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + steadyBytes, ok := checkedAddUint64(fixedBytes, retainedLoginBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + + loginDynamicBytes, ok := dynamicProtocolWriteBytes(handshakeLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + // Only complete statements inside msgBuf's fixed event window are captured + // for migration/control replay. Larger client packets are forwarded without + // becoming events, so they cannot reach ServerConn.ExecStmt here. + controlDynamicBytes, ok := dynamicProtocolWriteBytes(defaultBufLen + defaultExtraBufLen) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + initialBytes, ok := checkedAddUint64(handshakeLimit, loginDynamicBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + backendBytes, ok := checkedAddUint64( + proxyIOSessionBufferSize, + max(loginDynamicBytes, controlDynamicBytes), + ) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + transientBytes := max(initialBytes, backendBytes) + minimumBytes, ok := checkedAddUint64(steadyBytes, transientBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + if memoryLimit < minimumBytes { + return protocolMemoryBudget{}, moerr.NewInternalErrorNoCtx( + "proxy protocol-memory-limit is smaller than steady protocol buffers plus one transient operation") + } + return protocolMemoryBudget{ + steadyBytes: steadyBytes, + headroomBytes: memoryLimit - steadyBytes, + transientBytes: transientBytes, + initialBytes: initialBytes, + backendBytes: backendBytes, + loginDynamicBytes: loginDynamicBytes, + }, nil +} + +// dynamicProtocolWriteBytes mirrors frontend.Conn.AppendPart: WritePacket adds +// a four-byte header, consumes the fixed IOSession block first, then rounds one +// dynamic allocation up to an IOSession block boundary. +func dynamicProtocolWriteBytes(payloadBytes uint64) (uint64, bool) { + packetBytes, ok := checkedAddUint64(payloadBytes, frontend.PacketHeaderLength) + if !ok { + return 0, false + } + if packetBytes <= proxyIOSessionBufferSize { + return 0, true + } + remaining := packetBytes - proxyIOSessionBufferSize + allocation := max(uint64(proxyIOSessionBufferSize), remaining) + return roundUpUint64(allocation, proxyIOSessionBufferSize) +} + +func protocolMemoryConfigOverflow() error { + return moerr.NewInternalErrorNoCtx("proxy protocol memory configuration is too large") +} + +func checkedAddUint64(a, b uint64) (uint64, bool) { + if a > math.MaxUint64-b { + return 0, false + } + return a + b, true +} + +func checkedMulUint64(a, b uint64) (uint64, bool) { + if a != 0 && b > math.MaxUint64/a { + return 0, false + } + return a * b, true +} + +func roundUpUint64(value, unit uint64) (uint64, bool) { + if unit == 0 { + return 0, false + } + remainder := value % unit + if remainder == 0 { + return value, true + } + return checkedAddUint64(value, unit-remainder) +} + +type protocolMemoryLimiter struct { + weighted *semaphore.Weighted + budget protocolMemoryBudget + used atomic.Int64 +} + +func newProtocolMemoryLimiter(c *Config) (*protocolMemoryLimiter, error) { + budget, err := calculateProtocolMemoryBudget(c) + if err != nil { + return nil, err + } + return newProtocolMemoryLimiterWithBudget(budget), nil +} + +func newProtocolMemoryLimiterWithBudget(budget protocolMemoryBudget) *protocolMemoryLimiter { + return &protocolMemoryLimiter{ + weighted: semaphore.NewWeighted(int64(budget.headroomBytes)), + budget: budget, + } +} + +// acquire waits instead of spuriously failing a valid phase transition, but +// every caller supplies a lifecycle deadline so admission cannot hang. +func (l *protocolMemoryLimiter) acquire( + ctx context.Context, + bytes uint64, +) (*protocolMemoryLease, error) { + if l == nil || bytes == 0 { + return nil, nil + } + if bytes > l.budget.headroomBytes || bytes > math.MaxInt64 { + return nil, errProxyConnectionLimit + } + if ctx == nil { + ctx = context.Background() + } + amount := int64(bytes) + if err := l.weighted.Acquire(ctx, amount); err != nil { + return nil, errors.Join(errProxyConnectionLimit, context.Cause(ctx)) + } + l.used.Add(amount) + lease := &protocolMemoryLease{limiter: l, bytes: amount} + lease.refs.Store(1) + return lease, nil +} + +type protocolMemoryLease struct { + limiter *protocolMemoryLimiter + bytes int64 + refs atomic.Int32 +} + +// retain adds another concrete owner (for example, a pipelined prefix whose +// lifetime extends beyond backend authentication). It is only called before +// that owner becomes concurrently visible. +func (l *protocolMemoryLease) retain() { + if l == nil { + return + } + for { + refs := l.refs.Load() + if refs <= 0 { + panic("retaining released proxy protocol memory lease") + } + if l.refs.CompareAndSwap(refs, refs+1) { + return + } + } +} + +func (l *protocolMemoryLease) release() { + if l == nil { + return + } + for { + refs := l.refs.Load() + if refs <= 0 { + panic("proxy protocol memory lease released more than retained") + } + if !l.refs.CompareAndSwap(refs, refs-1) { + continue + } + if refs == 1 { + l.limiter.used.Add(-l.bytes) + l.limiter.weighted.Release(l.bytes) + } + return + } +} + +// protocolMemoryServerConn owns transient overlap while a newly authenticated +// backend coexists with the old backend. Close covers every failure path; +// promote releases the overlap only after the tunnel has closed the old side. +type protocolMemoryServerConn struct { + ServerConn + lease *protocolMemoryLease + once sync.Once +} + +func (c *protocolMemoryServerConn) releaseProtocolMemory() { + if c == nil { + return + } + c.once.Do(func() { + c.lease.release() + c.lease = nil + }) +} + +func (c *protocolMemoryServerConn) promoteProtocolMemory() { + c.releaseProtocolMemory() +} + +func (c *protocolMemoryServerConn) Quit() error { + defer c.releaseProtocolMemory() + return c.ServerConn.Quit() +} + +func (c *protocolMemoryServerConn) Close() error { + defer c.releaseProtocolMemory() + return c.ServerConn.Close() +} + +func acquireProtocolMemoryBefore( + ctx context.Context, + limiter *protocolMemoryLimiter, + bytes uint64, + deadline time.Time, +) (*protocolMemoryLease, error) { + if limiter == nil { + return nil, nil + } + if ctx == nil { + ctx = context.Background() + } + acquireCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + return limiter.acquire(acquireCtx, bytes) +} diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go new file mode 100644 index 0000000000000..b11b8a4b4a1bc --- /dev/null +++ b/pkg/proxy/protocol_memory_test.go @@ -0,0 +1,121 @@ +// Copyright 2026 Matrix Origin +// +// 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 proxy + +import ( + "context" + "testing" + "time" + + "github.com/matrixorigin/matrixone/pkg/util/toml" + "github.com/stretchr/testify/require" +) + +func TestCalculateProtocolMemoryBudget(t *testing.T) { + t.Run("small login needs one extra backend buffer", func(t *testing.T) { + transient := uint64(proxyIOSessionBufferSize) + steady := uint64(10 * (2*proxyIOSessionBufferSize + 64)) + cfg := Config{ + MaxConnections: 10, + ProtocolMemoryLimit: toml.ByteSize(steady), + ClientHandshakePacketLimit: 64, + } + _, err := calculateProtocolMemoryBudget(&cfg) + require.Error(t, err) + + cfg.ProtocolMemoryLimit += toml.ByteSize(transient) + budget, err := calculateProtocolMemoryBudget(&cfg) + require.NoError(t, err) + require.Equal(t, steady, budget.steadyBytes) + require.Equal(t, transient, budget.transientBytes) + require.Equal(t, transient, budget.headroomBytes) + require.Equal(t, uint64(64), budget.initialBytes) + require.Equal(t, uint64(proxyIOSessionBufferSize), budget.backendBytes) + }) + + t.Run("large login includes rounded backend dynamic write", func(t *testing.T) { + const handshakeLimit = 64 << 10 + steady := uint64(2*proxyIOSessionBufferSize + handshakeLimit) + transient := uint64(2 * handshakeLimit) + cfg := Config{ + MaxConnections: 1, + ProtocolMemoryLimit: toml.ByteSize(steady + 2*transient), + ClientHandshakePacketLimit: handshakeLimit, + } + budget, err := calculateProtocolMemoryBudget(&cfg) + require.NoError(t, err) + require.Equal(t, steady, budget.steadyBytes) + require.Equal(t, transient, budget.transientBytes) + require.Equal(t, 2*transient, budget.headroomBytes) + require.Equal(t, uint64(2*handshakeLimit), budget.initialBytes) + require.Equal(t, uint64(proxyIOSessionBufferSize+handshakeLimit), budget.backendBytes) + }) +} + +func TestProtocolMemoryLimiterLifecycle(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 100}) + first, err := limiter.acquire(context.Background(), 60) + require.NoError(t, err) + first.retain() + + timedOut, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + _, err = limiter.acquire(timedOut, 60) + require.ErrorIs(t, err, errProxyConnectionLimit) + require.ErrorIs(t, err, context.DeadlineExceeded) + + // Releasing only the authentication owner must not make capacity reusable + // while a pipelined prefix still owns the same transient peak. + first.release() + stillBlocked, cancelBlocked := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancelBlocked() + _, err = limiter.acquire(stillBlocked, 60) + require.ErrorIs(t, err, context.DeadlineExceeded) + + first.release() + second, err := limiter.acquire(context.Background(), 60) + require.NoError(t, err) + second.release() + require.Zero(t, limiter.used.Load()) +} + +func TestProtocolMemoryServerConnLifecycle(t *testing.T) { + t.Run("failed connection closes lease", func(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 1}) + lease, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + conn := &protocolMemoryServerConn{ + ServerConn: &mockServerConn{}, + lease: lease, + } + require.NoError(t, conn.Close()) + require.NoError(t, conn.Close()) + require.Zero(t, limiter.used.Load()) + }) + + t.Run("promotion is idempotent", func(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 1}) + lease, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + conn := &protocolMemoryServerConn{ + ServerConn: &mockServerConn{}, + lease: lease, + } + conn.promoteProtocolMemory() + conn.promoteProtocolMemory() + require.NoError(t, conn.Close()) + require.Zero(t, limiter.used.Load()) + }) +} diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index a56da5abef54f..7faf74f89aa6c 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -120,7 +120,7 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err goetty.WithAppSessionOptions( goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec( frontend.WithSQLCodecMaxPayloadSize(int(config.ClientHandshakePacketLimit)), - ))), + ), WithProxyProtocolMaxBodySize(int(config.ClientHandshakePacketLimit)))), goetty.WithSessionLogger(s.runtime.Logger().RawLogger()), ), ) diff --git a/pkg/proxy/tunnel.go b/pkg/proxy/tunnel.go index 5e378e42df3ec..38c7d8a4db1a9 100644 --- a/pkg/proxy/tunnel.go +++ b/pkg/proxy/tunnel.go @@ -387,6 +387,12 @@ func (t *tunnel) replaceServerConn(newServerConn *MySQLConn, newSC ServerConn, s t.mu.csp = t.newPipe(pipeClientToServer, t.mu.clientConn, t.mu.serverConn) t.mu.scp = t.newPipe(pipeServerToClient, t.mu.serverConn, t.mu.clientConn) } + // The new backend is now the one steady backend reserved for this client. + // Release transient overlap only after the old backend has been closed and + // every tunnel pointer has been switched. + if leased, ok := newSC.(interface{ promoteProtocolMemory() }); ok { + leased.promoteProtocolMemory() + } } // canStartTransfer checks whether the transfer can be started. diff --git a/pkg/proxy/tunnel_test.go b/pkg/proxy/tunnel_test.go index 8996151ce7957..ca783279efdd2 100644 --- a/pkg/proxy/tunnel_test.go +++ b/pkg/proxy/tunnel_test.go @@ -843,10 +843,15 @@ func TestReplaceServerConn(t *testing.T) { require.NoError(t, scp.pause(ctx)) newServerProxy, newServer := net.Pipe() - newSC := newMockServerConn(newServerProxy) - require.NotNil(t, sc) + backendSC := newMockServerConn(newServerProxy) + require.NotNil(t, backendSC) + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 1}) + lease, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + newSC := &protocolMemoryServerConn{ServerConn: backendSC, lease: lease} newServerC := newMySQLConn("new-server", newSC.RawConn(), 0, nil, nil, false, 0) tu.replaceServerConn(newServerC, newSC, false) + require.Zero(t, limiter.used.Load(), "replacement must become steady only after tunnel switch") _, newMysqlSC := tu.getConns() require.Equal(t, newServerC, newMysqlSC) require.NoError(t, tu.kickoff()) From f4cc9b40c588bad6198c5076ecc26e98b4712608 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sun, 19 Jul 2026 23:48:30 +0800 Subject: [PATCH 09/15] fix(proxy): propagate transfer admission cancellation --- pkg/proxy/client_conn.go | 22 ++++++++++++++-------- pkg/proxy/client_conn_test.go | 12 ++++++------ pkg/proxy/handler.go | 2 +- pkg/proxy/handshake.go | 4 ++-- pkg/proxy/protocol_memory_test.go | 21 +++++++++++++++++++++ pkg/proxy/tunnel.go | 2 +- 6 files changed, 45 insertions(+), 18 deletions(-) diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 84bf37012583b..bdf6af19bf2ac 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -100,7 +100,9 @@ type ClientConn interface { // prevAddr is empty if it is the first time to build connection with // a cn server; otherwise, it is the address of the previous cn node // when it is transferring connection and the handshake phase is ignored. - BuildConnWithServer(prevAddr string) (ServerConn, error) + // ctx bounds phase admission; transfer cancellation must interrupt any wait + // for transient protocol-memory headroom. + BuildConnWithServer(ctx context.Context, prevAddr string) (ServerConn, error) // HandleEvent handles event that comes from tunnel data flow. HandleEvent(ctx context.Context, e IEvent, resp chan<- []byte) error // KillCurrentBackendConn kills the backend connection that is currently @@ -441,7 +443,7 @@ func (c *clientConn) SendErrToClient(err error) { } // BuildConnWithServer implements the ClientConn interface. -func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { +func (c *clientConn) BuildConnWithServer(ctx context.Context, prevAddr string) (ServerConn, error) { var transientLease *protocolMemoryLease if prevAddr == "" { // Step 1, proxy write initial handshake to client. @@ -455,7 +457,7 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { // Step 2, client send handshake response, which is auth request, // to proxy. var err error - transientLease, err = c.handleHandshakeResp() + transientLease, err = c.handleHandshakeResp(ctx) if err != nil { // This connection may come from heartbeat of LB, and receive EOF error // from it. Just return error and do not log it. @@ -478,7 +480,7 @@ func (c *clientConn) BuildConnWithServer(prevAddr string) (ServerConn, error) { } } else { var err error - transientLease, err = c.acquireBackendProtocolMemory(defaultTransferTimeout) + transientLease, err = c.acquireBackendProtocolMemory(ctx, defaultTransferTimeout) if err != nil { return nil, err } @@ -552,7 +554,7 @@ func (c *clientConn) sendErr(err error, resp chan<- []byte) { } func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) error { - lease, err := c.acquireBackendProtocolMemory(defaultTransferTimeout) + lease, err := c.acquireBackendProtocolMemory(c.ctx, defaultTransferTimeout) if err != nil { if resp != nil { c.sendErr(err, resp) @@ -595,27 +597,31 @@ func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) } func (c *clientConn) acquireProtocolMemory( + ctx context.Context, timeout time.Duration, bytes uint64, ) (*protocolMemoryLease, error) { if c.protocolMemoryLimiter == nil { return nil, nil } - ctx := c.ctx if ctx == nil { - ctx = context.Background() + ctx = c.ctx + if ctx == nil { + ctx = context.Background() + } } deadline := time.Now().Add(timeout) return acquireProtocolMemoryBefore(ctx, c.protocolMemoryLimiter, bytes, deadline) } func (c *clientConn) acquireBackendProtocolMemory( + ctx context.Context, timeout time.Duration, ) (*protocolMemoryLease, error) { if c.protocolMemoryLimiter == nil { return nil, nil } - return c.acquireProtocolMemory(timeout, c.protocolMemoryLimiter.budget.backendBytes) + return c.acquireProtocolMemory(ctx, timeout, c.protocolMemoryLimiter.budget.backendBytes) } // KillCurrentBackendConn implements the ClientConn interface. diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 9fc3682d6e2f5..1716ac6fcdda0 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -71,7 +71,7 @@ func (c *readDeadlineTrackingConn) readDeadline() time.Time { } func handleHandshakeRespForTest(c *clientConn) error { - lease, err := c.handleHandshakeResp() + lease, err := c.handleHandshakeResp(context.Background()) lease.release() return err } @@ -161,7 +161,7 @@ func (c *mockClientConn) GetHandshakePack() *frontend.Packet { return nil } func (c *mockClientConn) RawConn() net.Conn { return c.conn } func (c *mockClientConn) GetTenant() Tenant { return c.tenant } func (c *mockClientConn) SendErrToClient(err error) {} -func (c *mockClientConn) BuildConnWithServer(_ string) (ServerConn, error) { +func (c *mockClientConn) BuildConnWithServer(_ context.Context, _ string) (ServerConn, error) { var err error li := &c.clientInfo.labelInfo c.clientInfo.labelInfo = newLabelInfo(c.clientInfo.Tenant, li.Labels) @@ -839,7 +839,7 @@ func TestClientConn_ConnectToBackend(t *testing.T) { return nil, moerr.NewInternalErrorNoCtx("123 456") } - sc, err := cc.BuildConnWithServer("aaa") + sc, err := cc.BuildConnWithServer(context.Background(), "aaa") require.ErrorContains(t, err, "123 456") require.Nil(t, sc) }) @@ -875,7 +875,7 @@ func TestClientConn_ConnectToBackend(t *testing.T) { require.Equal(t, len(resp), n) }() - _, err := cc.BuildConnWithServer("") + _, err := cc.BuildConnWithServer(context.Background(), "") require.Error(t, err) // just test client, no router set require.Equal(t, "tenant1", string(cc.GetTenant())) require.NotNil(t, cc.GetHandshakePack()) @@ -1378,7 +1378,7 @@ func TestClientConnHandshakeTimeout(t *testing.T) { _, _ = remote.Write(makeClientHandshakeResp()) }() - _, err := client.BuildConnWithServer("") + _, err := client.BuildConnWithServer(context.Background(), "") require.Error(t, err) // no router is configured after the handshake <-peerDone require.Same(t, tracked, client.RawConn()) @@ -1569,7 +1569,7 @@ func TestClientConn_SendErrToClient(t *testing.T) { require.True(t, strings.Contains(string(b[4+1+2+1+5:n]), "internal error: msg1")) }() - _, err := cc.BuildConnWithServer("") + _, err := cc.BuildConnWithServer(context.Background(), "") require.Error(t, err) // just test client, no router set require.Equal(t, "tenant1", string(cc.GetTenant())) require.NotNil(t, cc.GetHandshakePack()) diff --git a/pkg/proxy/handler.go b/pkg/proxy/handler.go index 1ecc553c1ed03..1aaa8cc7873cc 100644 --- a/pkg/proxy/handler.go +++ b/pkg/proxy/handler.go @@ -274,7 +274,7 @@ func (h *handler) handle(c goetty.IOSession) error { // client builds connections with a best CN server and returns // the server connection. - sc, err := cc.BuildConnWithServer("") + sc, err := cc.BuildConnWithServer(h.ctx, "") if err != nil { if isConnEndErr(err) { return nil diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index dc6bc1065b6aa..1009f86eec433 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -34,7 +34,7 @@ func (c *clientConn) writeInitialHandshake() error { // handleHandshakeResp receives login information from client and saves it // in proxy end. -func (c *clientConn) handleHandshakeResp() (*protocolMemoryLease, error) { +func (c *clientConn) handleHandshakeResp(ctx context.Context) (*protocolMemoryLease, error) { deadline := time.Now().Add(c.clientHandshakeTimeout) rawConn := c.conn.RawConn() deadlineConn := newAbsoluteReadDeadlineConn(rawConn, deadline) @@ -59,7 +59,7 @@ func (c *clientConn) handleHandshakeResp() (*protocolMemoryLease, error) { transientBytes = c.protocolMemoryLimiter.budget.initialBytes } lease, err := acquireProtocolMemoryBefore( - c.ctx, + ctx, c.protocolMemoryLimiter, transientBytes, deadline, diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go index b11b8a4b4a1bc..d828e1466edd5 100644 --- a/pkg/proxy/protocol_memory_test.go +++ b/pkg/proxy/protocol_memory_test.go @@ -91,6 +91,27 @@ func TestProtocolMemoryLimiterLifecycle(t *testing.T) { require.Zero(t, limiter.used.Load()) } +func TestClientConnProtocolMemoryAdmissionUsesTransferContext(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + backendBytes: 1, + }) + occupied, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + defer occupied.release() + + client := &clientConn{ + ctx: context.Background(), + protocolMemoryLimiter: limiter, + } + transferCtx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = client.BuildConnWithServer(transferCtx, "old-backend") + require.ErrorIs(t, err, errProxyConnectionLimit) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, int64(1), limiter.used.Load()) +} + func TestProtocolMemoryServerConnLifecycle(t *testing.T) { t.Run("failed connection closes lease", func(t *testing.T) { limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 1}) diff --git a/pkg/proxy/tunnel.go b/pkg/proxy/tunnel.go index 38c7d8a4db1a9..83af4ef7c7ca6 100644 --- a/pkg/proxy/tunnel.go +++ b/pkg/proxy/tunnel.go @@ -582,7 +582,7 @@ func (t *tunnel) getNewServerConn(ctx context.Context) (ServerConn, *MySQLConn, } prevAddr := t.mu.serverConn.RemoteAddr().String() t.logger.Info("build connection with new server", zap.String("prev addr", prevAddr)) - newConn, err := t.cc.BuildConnWithServer(prevAddr) + newConn, err := t.cc.BuildConnWithServer(ctx, prevAddr) if err != nil { t.logger.Error("failed to build connection with new server", zap.String("prev addr", prevAddr), From b594a264c70af1aa462c84c3486e5df2c52ab70f Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Mon, 20 Jul 2026 00:55:13 +0800 Subject: [PATCH 10/15] fix(proxy): bound transient connection phases --- pkg/proxy/client_conn.go | 149 ++++++++++++++++---- pkg/proxy/conn_cache.go | 45 +++++- pkg/proxy/conn_cache_test.go | 57 +++++++- pkg/proxy/conn_migration.go | 48 +++++-- pkg/proxy/conn_migration_test.go | 26 ++++ pkg/proxy/plugin.go | 21 +++ pkg/proxy/protocol_memory.go | 8 ++ pkg/proxy/protocol_memory_test.go | 69 +++++++++ pkg/proxy/router.go | 86 +++++++++++- pkg/proxy/server_conn.go | 226 +++++++++++++++++++++++++++--- pkg/proxy/server_conn_test.go | 88 ++++++++++++ 11 files changed, 748 insertions(+), 75 deletions(-) diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index bdf6af19bf2ac..90eb5e34bc885 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -100,8 +100,8 @@ type ClientConn interface { // prevAddr is empty if it is the first time to build connection with // a cn server; otherwise, it is the address of the previous cn node // when it is transferring connection and the handshake phase is ignored. - // ctx bounds phase admission; transfer cancellation must interrupt any wait - // for transient protocol-memory headroom. + // ctx bounds the complete connection phase: transient-memory admission, + // routing, dial, backend authentication, and migration replay. BuildConnWithServer(ctx context.Context, prevAddr string) (ServerConn, error) // HandleEvent handles event that comes from tunnel data flow. HandleEvent(ctx context.Context, e IEvent, resp chan<- []byte) error @@ -120,9 +120,9 @@ type migration struct { type clientConnOption func(*clientConn) -// absoluteReadDeadlineConn prevents a relative timeout set by an upper layer -// from extending an already established lifecycle deadline. It remains usable -// after the handshake; disable makes subsequent deadline updates transparent. +// absoluteReadDeadlineConn prevents a relative timeout set by goetty from +// extending the client-handshake lifecycle deadline. It remains usable after +// TLS upgrade; disable makes subsequent deadline updates transparent. type absoluteReadDeadlineConn struct { net.Conn deadline time.Time @@ -226,7 +226,8 @@ type clientConn struct { queryClient qclient.QueryClient // testHelper is used for testing. testHelper struct { - connectToBackend func() (ServerConn, error) + connectToBackend func() (ServerConn, error) + connectToBackendContext func(context.Context) (ServerConn, error) } migration migration // sc is the server connection bond with the client conn. @@ -444,6 +445,17 @@ func (c *clientConn) SendErrToClient(err error) { // BuildConnWithServer implements the ClientConn interface. func (c *clientConn) BuildConnWithServer(ctx context.Context, prevAddr string) (ServerConn, error) { + if ctx == nil { + ctx = c.ctx + if ctx == nil { + ctx = context.Background() + } + } + if prevAddr != "" { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, defaultTransferTimeout) + defer cancel() + } var transientLease *protocolMemoryLease if prevAddr == "" { // Step 1, proxy write initial handshake to client. @@ -492,15 +504,19 @@ func (c *clientConn) BuildConnWithServer(ctx context.Context, prevAddr string) ( } }() // Step 3, proxy connects to a CN server to build connection. - conn, err := c.connectToBackend(prevAddr) + conn, err := c.connectToBackendContext(ctx, prevAddr) if err != nil { - if isProxyAdmissionError(err) { - c.log.Debug("backend connection rejected by proxy admission", zap.Error(err)) + if isProxyAdmissionError(err) || operationContextCause(ctx) != nil { + c.log.Debug("backend connection phase ended", zap.Error(err)) } else { c.log.Error("failed to connect to backend", zap.Error(err)) } return nil, err } + if cause := operationContextCause(ctx); cause != nil { + _ = conn.Close() + return nil, cause + } if prevAddr != "" && transientLease != nil { conn = &protocolMemoryServerConn{ ServerConn: conn, @@ -518,7 +534,7 @@ func (c *clientConn) BuildConnWithServer(ctx context.Context, prevAddr string) ( func (c *clientConn) HandleEvent(ctx context.Context, e IEvent, resp chan<- []byte) error { switch ev := e.(type) { case *killEvent: - return c.handleKill(ev, resp) + return c.handleKill(ctx, ev, resp) case *setVarEvent: return c.handleSetVar(ev) case *quitEvent: @@ -535,7 +551,7 @@ func (c *clientConn) HandleEvent(ctx context.Context, e IEvent, resp chan<- []by }() return nil case *upgradeEvent: - return c.handleUpgradeEvent(ev, resp) + return c.handleUpgradeEvent(ctx, ev, resp) default: } return nil @@ -553,8 +569,19 @@ func (c *clientConn) sendErr(err error, resp chan<- []byte) { sendResp(packetToBytes(r), resp) } -func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) error { - lease, err := c.acquireBackendProtocolMemory(c.ctx, defaultTransferTimeout) +func (c *clientConn) connAndExec( + ctx context.Context, + cn *CNServer, + stmt string, + resp chan<- []byte, +) error { + if ctx == nil { + ctx = c.ctx + if ctx == nil { + ctx = context.Background() + } + } + lease, err := c.acquireBackendProtocolMemory(ctx, defaultTransferTimeout) if err != nil { if resp != nil { c.sendErr(err, resp) @@ -564,7 +591,7 @@ func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) defer lease.release() sc, r, err := c.connectWithHandshakePack(func(pack *frontend.Packet) (ServerConn, []byte, error) { - return c.router.Connect(cn, pack, nil) + return connectRouterWithContext(ctx, c.router, cn, pack, nil, false) }) if err != nil { c.log.Error("failed to connect to backend server", zap.Error(err)) @@ -584,7 +611,7 @@ func (c *clientConn) connAndExec(cn *CNServer, stmt string, resp chan<- []byte) return moerr.NewInternalErrorNoCtx("access error") } - ok, err := sc.ExecStmt(internalStmt{cmdType: cmdQuery, s: stmt}, resp) + ok, err := execStmtWithContext(ctx, sc, internalStmt{cmdType: cmdQuery, s: stmt}, resp) if err != nil { c.log.Error("failed to send query to server", zap.String("query", stmt), zap.Error(err)) @@ -649,11 +676,21 @@ func (c *clientConn) KillCurrentBackendConn(sc ServerConn) error { } tempCN.connID = cid - return c.connAndExec(tempCN, fmt.Sprintf("kill connection %d", c.ConnID()), nil) + ctx := c.ctx + if ctx == nil { + ctx = context.Background() + } + ctx, cancel := context.WithTimeout(ctx, defaultTransferTimeout) + defer cancel() + return c.connAndExec(ctx, tempCN, fmt.Sprintf("kill connection %d", c.ConnID()), nil) } // handleKill handles the kill event. -func (c *clientConn) handleKill(e *killEvent, resp chan<- []byte) error { +func (c *clientConn) handleKill( + ctx context.Context, + e *killEvent, + resp chan<- []byte, +) error { cn, err := c.router.SelectByConnID(e.connID) if err != nil { // If no server found, means that the query has been terminated. @@ -677,7 +714,12 @@ func (c *clientConn) handleKill(e *killEvent, resp chan<- []byte) error { } cn.connID = cid - return c.connAndExec(cn, e.stmt, resp) + if ctx == nil { + ctx = context.Background() + } + ctx, cancel := context.WithTimeout(ctx, defaultTransferTimeout) + defer cancel() + return c.connAndExec(ctx, cn, e.stmt, resp) } // handleSetVar handles the set variable event. @@ -737,7 +779,11 @@ func (c *clientConn) isConnCached() bool { return c.quit.cached } -func (c *clientConn) handleUpgradeEvent(e *upgradeEvent, resp chan<- []byte) error { +func (c *clientConn) handleUpgradeEvent( + ctx context.Context, + e *upgradeEvent, + resp chan<- []byte, +) error { defer e.notify() if !c.clientInfo.isSuperTenant() { @@ -769,7 +815,7 @@ func (c *clientConn) handleUpgradeEvent(e *upgradeEvent, resp chan<- []byte) err // In the loop, do not pass the resp, because it only receives response once. // It everything is ok, send ok response at last out of the loop. - if err := c.connAndExec(cn, e.stmt, nil); err != nil { + if err := c.connAndExec(ctx, cn, e.stmt, nil); err != nil { c.log.Error("failed to execute upgrade query", zap.Error(err)) c.sendErr(err, resp) return err @@ -822,9 +868,24 @@ func (c *clientConn) connectWithHandshakePack( return connect(c.handshakePack) } -// connectToBackend connect to the real CN server. +// connectToBackend connects to the real CN server using the client lifecycle +// context. Tests and legacy internal callers retain the old helper, while +// BuildConnWithServer passes its narrower handshake/transfer context below. func (c *clientConn) connectToBackend(prevAdd string) (ServerConn, error) { + return c.connectToBackendContext(c.ctx, prevAdd) +} + +func (c *clientConn) connectToBackendContext( + ctx context.Context, + prevAdd string, +) (ServerConn, error) { + if ctx == nil { + ctx = context.Background() + } // Testing path. + if c.testHelper.connectToBackendContext != nil { + return c.testHelper.connectToBackendContext(ctx) + } if c.testHelper.connectToBackend != nil { return c.testHelper.connectToBackend() } @@ -846,7 +907,24 @@ func (c *clientConn) connectToBackend(prevAdd string) (ServerConn, error) { if _, pluginMode := c.router.(*pluginRouter); pluginMode { goto skipConnCache } - sc = c.connCache.Pop(c.clientInfo.hash, c.connID, c.mysqlProto.GetSalt(), c.mysqlProto.GetAuthResponse()) + if contextual, ok := c.connCache.(contextConnCache); ok { + cacheCtx, cancel := context.WithTimeout(ctx, defaultTransferTimeout) + sc = contextual.PopContext( + cacheCtx, + c.clientInfo.hash, + c.connID, + c.mysqlProto.GetSalt(), + c.mysqlProto.GetAuthResponse(), + ) + cancel() + } else { + sc = c.connCache.Pop( + c.clientInfo.hash, + c.connID, + c.mysqlProto.GetSalt(), + c.mysqlProto.GetAuthResponse(), + ) + } if sc != nil { if err := c.bindAuthenticatedTenant(); err != nil { _ = sc.Close() @@ -881,18 +959,24 @@ skipConnCache: var cn *CNServer var r []byte for { + if err := operationContextCause(ctx); err != nil { + return nil, err + } // Select the best CN server from backend. // // NB: The selected CNServer must have label hash in it. if prevAdd == "" { - cn, err = c.router.Route(c.ctx, c.sid, c.clientInfo, filterFn) + cn, err = c.router.Route(ctx, c.sid, c.clientInfo, filterFn) } else if tr, ok := c.router.(transferRouter); ok { - cn, err = tr.RouteForTransfer(c.ctx, c.sid, c.clientInfo, filterFn) + cn, err = tr.RouteForTransfer(ctx, c.sid, c.clientInfo, filterFn) } else { - cn, err = c.router.Route(c.ctx, c.sid, c.clientInfo, filterFn) + cn, err = c.router.Route(ctx, c.sid, c.clientInfo, filterFn) } if err != nil { v2.ProxyConnectRouteFailCounter.Inc() + if cause := operationContextCause(ctx); cause != nil { + return nil, cause + } // Check if all failed CN servers were due to timeout. // If so, return a more specific error message. if len(timeoutCNServers) > 0 && len(timeoutCNServers) == len(badCNServers) { @@ -922,10 +1006,7 @@ skipConnCache: if prevAdd == "" { sc, r, err = c.connectWithHandshakePack( func(pack *frontend.Packet) (ServerConn, []byte, error) { - if rr, ok := c.router.(routeSelectedConnector); ok { - return rr.ConnectRouteSelected(cn, pack, c.tun) - } - return c.router.Connect(cn, pack, c.tun) + return connectRouterWithContext(ctx, c.router, cn, pack, c.tun, true) }, ) } else { @@ -934,11 +1015,14 @@ skipConnCache: // Route-selected new-session connect. sc, r, err = c.connectWithHandshakePack( func(pack *frontend.Packet) (ServerConn, []byte, error) { - return c.router.Connect(cn, pack, c.tun) + return connectRouterWithContext(ctx, c.router, cn, pack, c.tun, false) }, ) } if err != nil { + if cause := operationContextCause(ctx); cause != nil { + return nil, cause + } if isRetryableErr(err) { v2.ProxyConnectRetryCounter.Inc() badCNServers[cn.addr] = struct{}{} @@ -998,11 +1082,14 @@ skipConnCache: } else { // The connection has been transferred to a new server, but migration fails, // but we don't return error, which will cause unknown issue. - if err := c.migrateConn(prevAdd, sc); err != nil { + if err := c.migrateConnContext(ctx, prevAdd, sc); err != nil { closeErr := sc.Close() if closeErr != nil { c.log.Error("failed to close server connection", zap.Error(closeErr)) } + if cause := operationContextCause(ctx); cause != nil { + return nil, cause + } c.log.Error("failed to migrate connection to cn, will retry", zap.Uint32("conn ID", c.connID), zap.String("current uuid", cn.uuid), diff --git a/pkg/proxy/conn_cache.go b/pkg/proxy/conn_cache.go index 504cac286f7b3..574d4ea7239a7 100644 --- a/pkg/proxy/conn_cache.go +++ b/pkg/proxy/conn_cache.go @@ -152,6 +152,14 @@ func (s *serverConnAuth) close() error { return err } +func (s *serverConnAuth) ExecStmtContext( + ctx context.Context, + stmt internalStmt, + resp chan<- []byte, +) (bool, error) { + return execStmtWithContext(ctx, s.ServerConn, stmt, resp) +} + // cacheStore is the storage which stores the cached connections. type cacheStore struct { connections []*serverConnAuth @@ -188,6 +196,10 @@ type ConnCache interface { Close() error } +type contextConnCache interface { + PopContext(context.Context, cacheKey, uint32, []byte, []byte) ServerConn +} + // the main cache struct. type connCache struct { ctx context.Context @@ -407,7 +419,23 @@ func (c *connCache) closeCachedConnection(sc *serverConnAuth) { // Pop implements the ConnCache interface. func (c *connCache) Pop(key cacheKey, connID uint32, salt []byte, authResp []byte) ServerConn { + return c.PopContext(context.Background(), key, connID, salt, authResp) +} + +func (c *connCache) PopContext( + ctx context.Context, + key cacheKey, + connID uint32, + salt []byte, + authResp []byte, +) ServerConn { + if ctx == nil { + ctx = context.Background() + } for { + if operationContextCause(ctx) != nil { + return nil + } // Reserve one available connection by removing it from the store. Keep // it in allConns until it is handed to the caller so Close can terminate // an in-flight SET CONNECTION ID that is blocked in backend I/O. @@ -443,15 +471,22 @@ func (c *connCache) Pop(key cacheKey, connID uint32, salt []byte, authResp []byt // Check if the connection is expired. if time.Since(sc.CreateTime()) < c.connTimeout { - ok, err := sc.ExecStmt(internalStmt{ + ok, err := execStmtWithContext(ctx, sc, internalStmt{ cmdType: cmdQuery, s: fmt.Sprintf(setConnectionIDSQL, connID), }, nil) if err != nil || !ok { - c.logger.Error("failed to set conn id", - zap.Uint32("conn ID", sc.ConnID()), - zap.Error(err), - ) + if operationContextCause(ctx) != nil { + c.logger.Debug("set conn id canceled", + zap.Uint32("conn ID", sc.ConnID()), + zap.Error(err), + ) + } else { + c.logger.Error("failed to set conn id", + zap.Uint32("conn ID", sc.ConnID()), + zap.Error(err), + ) + } c.closeCachedConnection(sc) continue } diff --git a/pkg/proxy/conn_cache_test.go b/pkg/proxy/conn_cache_test.go index 0aa9d063aadf7..54619ac5098e4 100644 --- a/pkg/proxy/conn_cache_test.go +++ b/pkg/proxy/conn_cache_test.go @@ -121,7 +121,7 @@ func simulateScramble(password string, salt []byte) []byte { return scrambled } -type blockingCacheServerConn struct { +type blockingContextServerConn struct { *mockServerConn enteredOnce sync.Once closeOnce sync.Once @@ -129,21 +129,35 @@ type blockingCacheServerConn struct { closed chan struct{} } -func newBlockingCacheServerConn(conn net.Conn) *blockingCacheServerConn { - return &blockingCacheServerConn{ +func newBlockingContextServerConn(conn net.Conn) *blockingContextServerConn { + return &blockingContextServerConn{ mockServerConn: newMockServerConn(conn), entered: make(chan struct{}), closed: make(chan struct{}), } } -func (s *blockingCacheServerConn) ExecStmt(internalStmt, chan<- []byte) (bool, error) { +func (s *blockingContextServerConn) ExecStmt(internalStmt, chan<- []byte) (bool, error) { s.enteredOnce.Do(func() { close(s.entered) }) <-s.closed return false, net.ErrClosed } -func (s *blockingCacheServerConn) Close() error { +func (s *blockingContextServerConn) ExecStmtContext( + ctx context.Context, + _ internalStmt, + _ chan<- []byte, +) (bool, error) { + s.enteredOnce.Do(func() { close(s.entered) }) + select { + case <-ctx.Done(): + return false, context.Cause(ctx) + case <-s.closed: + return false, net.ErrClosed + } +} + +func (s *blockingContextServerConn) Close() error { s.closeOnce.Do(func() { close(s.closed) _ = s.mockServerConn.Close() @@ -343,7 +357,7 @@ func TestConnCacheBlockedPopDoesNotBlockOtherTenantsOrClose(t *testing.T) { clientSide, serverSide := net.Pipe() defer clientSide.Close() - blocked := newBlockingCacheServerConn(serverSide) + blocked := newBlockingContextServerConn(serverSide) require.True(t, cache.Push("tenant-a", blocked)) popDone := make(chan ServerConn, 1) @@ -388,6 +402,37 @@ func TestConnCacheBlockedPopDoesNotBlockOtherTenantsOrClose(t *testing.T) { } } +func TestConnCachePopContextCancelsBackendValidation(t *testing.T) { + cache := newConnCache(context.Background(), "", runtime.DefaultRuntime().Logger(), + withResetSessionFunc(func(ServerConn) ([]byte, error) { return nil, nil }), + withAuthConstructor(nil), + ) + clientSide, serverSide := net.Pipe() + defer clientSide.Close() + blocked := newBlockingContextServerConn(serverSide) + require.True(t, cache.Push("tenant-a", &protocolMemoryServerConn{ServerConn: blocked})) + + ctx, cancel := context.WithCancel(context.Background()) + popDone := make(chan ServerConn, 1) + go func() { + popDone <- cache.(*connCache).PopContext(ctx, "tenant-a", 1, nil, nil) + }() + select { + case <-blocked.entered: + case <-time.After(time.Second): + t.Fatal("PopContext did not enter backend validation") + } + cancel() + select { + case sc := <-popDone: + require.Nil(t, sc) + case <-time.After(time.Second): + t.Fatal("PopContext ignored lifecycle cancellation") + } + require.Zero(t, cache.Count()) + require.NoError(t, cache.Close()) +} + func TestConnCacheCloseDoesNotWaitForPushReset(t *testing.T) { ctx := context.Background() logger := runtime.DefaultRuntime().Logger() diff --git a/pkg/proxy/conn_migration.go b/pkg/proxy/conn_migration.go index dfd0adde4a29e..f83ba4d174246 100644 --- a/pkg/proxy/conn_migration.go +++ b/pkg/proxy/conn_migration.go @@ -27,15 +27,25 @@ import ( ) func (c *clientConn) migrateConnFrom(sqlAddr string) (*query.MigrateConnFromResponse, error) { + return c.migrateConnFromContext(c.ctx, sqlAddr) +} + +func (c *clientConn) migrateConnFromContext( + parent context.Context, + sqlAddr string, +) (*query.MigrateConnFromResponse, error) { + if parent == nil { + parent = context.Background() + } req := c.queryClient.NewRequest(query.CmdMethod_MigrateConnFrom) req.MigrateConnFromRequest = &query.MigrateConnFromRequest{ ConnID: c.connID, } - ctx, cancel := context.WithTimeoutCause(c.ctx, time.Second*3, moerr.CauseMigrateConnFrom) + ctx, cancel := context.WithTimeoutCause(parent, time.Second*3, moerr.CauseMigrateConnFrom) defer cancel() addr := getQueryAddress(c.moCluster, sqlAddr) if addr == "" { - return nil, moerr.NewInternalError(c.ctx, "cannot get query service address") + return nil, moerr.NewInternalError(parent, "cannot get query service address") } resp, err := c.queryClient.SendMessage(ctx, addr, req) if err != nil { @@ -57,11 +67,22 @@ func (c *clientConn) migrateConnFrom(sqlAddr string) (*query.MigrateConnFromResp } func (c *clientConn) migrateConnTo(sc ServerConn, info *query.MigrateConnFromResponse) error { + return c.migrateConnToContext(c.ctx, sc, info) +} + +func (c *clientConn) migrateConnToContext( + parent context.Context, + sc ServerConn, + info *query.MigrateConnFromResponse, +) error { + if parent == nil { + parent = context.Background() + } // Before migrate session info with RPC, we need to execute some // SQLs to initialize the session and account in handler. // Currently, the session variable transferred is not used anywhere else, // and just used here. - if _, err := sc.ExecStmt(internalStmt{ + if _, err := execStmtWithContext(parent, sc, internalStmt{ cmdType: cmdQuery, s: "/* cloud_nonuser */ set transferred=1;", }, nil); err != nil { @@ -70,7 +91,7 @@ func (c *clientConn) migrateConnTo(sc ServerConn, info *query.MigrateConnFromRes // First, we re-run the set variables statements. for _, stmt := range c.migration.setVarStmts { - if _, err := sc.ExecStmt(internalStmt{ + if _, err := execStmtWithContext(parent, sc, internalStmt{ cmdType: cmdQuery, s: stmt, }, nil); err != nil { @@ -82,7 +103,7 @@ func (c *clientConn) migrateConnTo(sc ServerConn, info *query.MigrateConnFromRes // Then, migrate other info with RPC. addr := getQueryAddress(c.moCluster, sc.RawConn().RemoteAddr().String()) if addr == "" { - return moerr.NewInternalError(c.ctx, "cannot get query service address") + return moerr.NewInternalError(parent, "cannot get query service address") } c.log.Info("connection migrate to server", zap.String("server address", addr), zap.String("tenant", string(c.clientInfo.Tenant)), @@ -97,7 +118,7 @@ func (c *clientConn) migrateConnTo(sc ServerConn, info *query.MigrateConnFromRes PrepareStmts: info.PrepareStmts, LastAffectedRows: info.LastAffectedRows, } - ctx, cancel := context.WithTimeoutCause(c.ctx, time.Second*3, moerr.CauseMigrateConnTo) + ctx, cancel := context.WithTimeoutCause(parent, time.Second*3, moerr.CauseMigrateConnTo) defer cancel() resp, err := c.queryClient.SendMessage(ctx, addr, req) if err != nil { @@ -107,13 +128,20 @@ func (c *clientConn) migrateConnTo(sc ServerConn, info *query.MigrateConnFromRes return nil } -func (c *clientConn) migrateConn(prevAddr string, sc ServerConn) error { - resp, err := c.migrateConnFrom(prevAddr) +func (c *clientConn) migrateConnContext( + ctx context.Context, + prevAddr string, + sc ServerConn, +) error { + if ctx == nil { + ctx = context.Background() + } + resp, err := c.migrateConnFromContext(ctx, prevAddr) if err != nil { return err } if resp == nil { - return moerr.NewInternalError(c.ctx, "bad response") + return moerr.NewInternalError(ctx, "bad response") } - return c.migrateConnTo(sc, resp) + return c.migrateConnToContext(ctx, sc, resp) } diff --git a/pkg/proxy/conn_migration_test.go b/pkg/proxy/conn_migration_test.go index 6034027d24c79..2c1f5656463d8 100644 --- a/pkg/proxy/conn_migration_test.go +++ b/pkg/proxy/conn_migration_test.go @@ -146,3 +146,29 @@ func TestQueryServiceMigrateTo(t *testing.T) { assert.NoError(t, err) }) } + +func TestMigrateConnToContextCancelsReplay(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + blocked := newBlockingContextServerConn(local) + defer blocked.Close() + cc := &clientConn{} + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + result <- cc.migrateConnToContext(ctx, blocked, &pb.MigrateConnFromResponse{}) + }() + + select { + case <-blocked.entered: + case <-time.After(time.Second): + t.Fatal("migration replay did not enter backend ExecStmt") + } + cancel() + select { + case err := <-result: + assert.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("migration replay ignored transfer cancellation") + } +} diff --git a/pkg/proxy/plugin.go b/pkg/proxy/plugin.go index 53c3d7485d69b..f7ec93baf497f 100644 --- a/pkg/proxy/plugin.go +++ b/pkg/proxy/plugin.go @@ -54,6 +54,27 @@ func (r *pluginRouter) ConnectRouteSelected( return r.Router.Connect(cn, handshakeResp, t) } +func (r *pluginRouter) ConnectContext( + ctx context.Context, cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, +) (ServerConn, []byte, error) { + if rr, ok := r.Router.(contextConnector); ok { + return rr.ConnectContext(ctx, cn, handshakeResp, t) + } + return r.Router.Connect(cn, handshakeResp, t) +} + +func (r *pluginRouter) ConnectRouteSelectedContext( + ctx context.Context, cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, +) (ServerConn, []byte, error) { + if rr, ok := r.Router.(contextRouteSelectedConnector); ok { + return rr.ConnectRouteSelectedContext(ctx, cn, handshakeResp, t) + } + if rr, ok := r.Router.(routeSelectedConnector); ok { + return rr.ConnectRouteSelected(cn, handshakeResp, t) + } + return r.ConnectContext(ctx, cn, handshakeResp, t) +} + // RouteForTransfer preserves plugin routing semantics for session transfer / // migration, but intentionally bypasses the breaker/probe gate: transfer // traffic must not consume a recovery probe that belongs to new-session diff --git a/pkg/proxy/protocol_memory.go b/pkg/proxy/protocol_memory.go index 3479e29b34138..cf8b725fe3440 100644 --- a/pkg/proxy/protocol_memory.go +++ b/pkg/proxy/protocol_memory.go @@ -304,6 +304,14 @@ func (c *protocolMemoryServerConn) Close() error { return c.ServerConn.Close() } +func (c *protocolMemoryServerConn) ExecStmtContext( + ctx context.Context, + stmt internalStmt, + resp chan<- []byte, +) (bool, error) { + return execStmtWithContext(ctx, c.ServerConn, stmt, resp) +} + func acquireProtocolMemoryBefore( ctx context.Context, limiter *protocolMemoryLimiter, diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go index d828e1466edd5..5fff0f6c1ca9f 100644 --- a/pkg/proxy/protocol_memory_test.go +++ b/pkg/proxy/protocol_memory_test.go @@ -16,9 +16,11 @@ package proxy import ( "context" + "net" "testing" "time" + "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/util/toml" "github.com/stretchr/testify/require" ) @@ -112,6 +114,73 @@ func TestClientConnProtocolMemoryAdmissionUsesTransferContext(t *testing.T) { require.Equal(t, int64(1), limiter.used.Load()) } +func TestClientConnProtocolMemoryTransferCancellationAfterAdmission(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + backendBytes: 1, + }) + entered := make(chan struct{}) + client := &clientConn{ + ctx: context.Background(), + log: runtime.DefaultRuntime().Logger(), + protocolMemoryLimiter: limiter, + } + client.testHelper.connectToBackendContext = func(ctx context.Context) (ServerConn, error) { + close(entered) + <-ctx.Done() + return nil, context.Cause(ctx) + } + + transferCtx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + _, err := client.BuildConnWithServer(transferCtx, "old-backend") + result <- err + }() + <-entered + require.Equal(t, int64(1), limiter.used.Load(), "transfer owns the admitted transient slot") + cancel() + + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("transfer cancellation did not terminate post-admission work") + } + require.Zero(t, limiter.used.Load(), "canceled transfer must release global headroom") + reused, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + reused.release() +} + +func TestClientConnProtocolMemoryCancellationWinsConnectHandoff(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + backendBytes: 1, + }) + local, remote := net.Pipe() + defer remote.Close() + backend := newMockServerConn(local) + client := &clientConn{ + ctx: context.Background(), + log: runtime.DefaultRuntime().Logger(), + protocolMemoryLimiter: limiter, + } + transferCtx, cancel := context.WithCancel(context.Background()) + client.testHelper.connectToBackendContext = func(context.Context) (ServerConn, error) { + cancel() + return backend, nil + } + + conn, err := client.BuildConnWithServer(transferCtx, "old-backend") + require.Nil(t, conn) + require.ErrorIs(t, err, context.Canceled) + require.Zero(t, limiter.used.Load(), "canceled handoff must release global headroom") + _ = remote.SetWriteDeadline(time.Now().Add(time.Second)) + _, err = remote.Write([]byte{1}) + require.Error(t, err, "canceled handoff must close the unowned backend") +} + func TestProtocolMemoryServerConnLifecycle(t *testing.T) { t.Run("failed connection closes lease", func(t *testing.T) { limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 1}) diff --git a/pkg/proxy/router.go b/pkg/proxy/router.go index c82686187416f..3d0b1740fb5ab 100644 --- a/pkg/proxy/router.go +++ b/pkg/proxy/router.go @@ -73,6 +73,55 @@ type routeSelectedConnector interface { ConnectRouteSelected(c *CNServer, handshakeResp *frontend.Packet, t *tunnel) (ServerConn, []byte, error) } +// contextConnector and contextRouteSelectedConnector keep Router source +// compatibility while allowing connection phases that own transient memory to +// propagate their lifecycle context through dial and backend authentication. +type contextConnector interface { + ConnectContext( + ctx context.Context, + c *CNServer, + handshakeResp *frontend.Packet, + t *tunnel, + ) (ServerConn, []byte, error) +} + +type contextRouteSelectedConnector interface { + ConnectRouteSelectedContext( + ctx context.Context, + c *CNServer, + handshakeResp *frontend.Packet, + t *tunnel, + ) (ServerConn, []byte, error) +} + +func connectRouterWithContext( + ctx context.Context, + r Router, + cn *CNServer, + handshakeResp *frontend.Packet, + t *tunnel, + routeSelected bool, +) (ServerConn, []byte, error) { + if ctx == nil { + ctx = context.Background() + } + if cause := operationContextCause(ctx); cause != nil { + return nil, nil, cause + } + if routeSelected { + if rr, ok := r.(contextRouteSelectedConnector); ok { + return rr.ConnectRouteSelectedContext(ctx, cn, handshakeResp, t) + } + if rr, ok := r.(routeSelectedConnector); ok { + return rr.ConnectRouteSelected(cn, handshakeResp, t) + } + } + if rr, ok := r.(contextConnector); ok { + return rr.ConnectContext(ctx, cn, handshakeResp, t) + } + return r.Connect(cn, handshakeResp, t) +} + // transferRouter is implemented by routers that want session transfer / // migration traffic to select a target CN without consuming a breaker-managed // recovery probe. The transfer path is control-plane traffic, not a new @@ -365,10 +414,19 @@ func (r *router) CanReuseCachedCN(cn *CNServer) bool { } func (r *router) connect( - cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, accountHealth bool, + ctx context.Context, + cn *CNServer, + handshakeResp *frontend.Packet, + t *tunnel, + accountHealth bool, ) (ServerConn, []byte, error) { + if ctx == nil { + ctx = context.Background() + } // Creates a server connection. - sc, err := newServerConn(cn, t, r.rebalancer, r.connectTimeout, r.sessionAllocator) + sc, err := newServerConnContext( + ctx, cn, t, r.rebalancer, r.connectTimeout, r.sessionAllocator, + ) if err != nil { // Connection failed, remove the placeholder that was added in selectOne. r.rebalancer.connManager.selectOneFailed(cn.hash, cn.uuid) @@ -405,7 +463,13 @@ func (r *router) connect( WithLabelValues(result). Observe(time.Since(start).Seconds()) }() - resp, err := sc.HandleHandshake(handshakeResp, r.authTimeout) + var resp *frontend.Packet + var err error + if contextual, ok := sc.(contextHandshakeHandler); ok { + resp, err = contextual.HandleHandshakeContext(ctx, handshakeResp, r.authTimeout) + } else { + resp, err = sc.HandleHandshake(handshakeResp, r.authTimeout) + } if err != nil { result = "error" if isTimeoutErr(err) { @@ -445,7 +509,13 @@ func (r *router) connect( func (r *router) Connect( cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, ) (ServerConn, []byte, error) { - return r.connect(cn, handshakeResp, t, false) + return r.connect(context.Background(), cn, handshakeResp, t, false) +} + +func (r *router) ConnectContext( + ctx context.Context, cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, +) (ServerConn, []byte, error) { + return r.connect(ctx, cn, handshakeResp, t, false) } // ConnectRouteSelected is used only for Route-selected new-session connects. @@ -454,7 +524,13 @@ func (r *router) Connect( func (r *router) ConnectRouteSelected( cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, ) (ServerConn, []byte, error) { - return r.connect(cn, handshakeResp, t, true) + return r.connect(context.Background(), cn, handshakeResp, t, true) +} + +func (r *router) ConnectRouteSelectedContext( + ctx context.Context, cn *CNServer, handshakeResp *frontend.Packet, t *tunnel, +) (ServerConn, []byte, error) { + return r.connect(ctx, cn, handshakeResp, t, true) } // Refresh refreshes the router diff --git a/pkg/proxy/server_conn.go b/pkg/proxy/server_conn.go index 52143df044500..f3c943aef14a2 100644 --- a/pkg/proxy/server_conn.go +++ b/pkg/proxy/server_conn.go @@ -19,6 +19,8 @@ import ( "errors" "io" "net" + "net/url" + "strings" "sync" "sync/atomic" "time" @@ -95,6 +97,40 @@ type serverConn struct { var _ ServerConn = (*serverConn)(nil) +type contextHandshakeHandler interface { + HandleHandshakeContext( + ctx context.Context, + handshakeResp *frontend.Packet, + timeout time.Duration, + ) (*frontend.Packet, error) +} + +type contextStmtExecutor interface { + ExecStmtContext( + ctx context.Context, + stmt internalStmt, + resp chan<- []byte, + ) (bool, error) +} + +func execStmtWithContext( + ctx context.Context, + sc ServerConn, + stmt internalStmt, + resp chan<- []byte, +) (bool, error) { + if ctx == nil { + ctx = context.Background() + } + if cause := operationContextCause(ctx); cause != nil { + return false, cause + } + if contextual, ok := sc.(contextStmtExecutor); ok { + return contextual.ExecStmtContext(ctx, stmt, resp) + } + return sc.ExecStmt(stmt, resp) +} + // Proxy only needs a small retained buffer for normal MySQL packets. Larger // packets use frontend.Conn's dynamic path and are released after use. const proxyIOSessionBufferSize = 16 * 1024 @@ -106,12 +142,25 @@ func newServerConn( r *rebalancer, timeout time.Duration, allocators ...frontend.Allocator, +) (ServerConn, error) { + return newServerConnContext( + context.Background(), cn, tun, r, timeout, allocators..., + ) +} + +func newServerConnContext( + ctx context.Context, + cn *CNServer, + tun *tunnel, + r *rebalancer, + timeout time.Duration, + allocators ...frontend.Allocator, ) (ServerConn, error) { var logger *zap.Logger if r != nil && r.logger != nil { logger = r.logger.RawLogger() } - c, err := cn.Connect(logger, timeout) + c, err := cn.ConnectContext(ctx, logger, timeout) if err != nil { return nil, err } @@ -196,8 +245,32 @@ func (s *serverConn) RawConn() net.Conn { func (s *serverConn) HandleHandshake( handshakeResp *frontend.Packet, timeout time.Duration, ) (*frontend.Packet, error) { - ctx, cancel := context.WithTimeoutCause(context.Background(), timeout, moerr.CauseHandleHandshake) + return s.HandleHandshakeContext(context.Background(), handshakeResp, timeout) +} + +func (s *serverConn) HandleHandshakeContext( + parent context.Context, + handshakeResp *frontend.Packet, + timeout time.Duration, +) (*frontend.Packet, error) { + if parent == nil { + parent = context.Background() + } + ctx, cancel := context.WithTimeoutCause(parent, timeout, moerr.CauseHandleHandshake) defer cancel() + if s == nil || s.conn == nil { + return nil, moerr.NewInternalErrorNoCtx("backend connection is unavailable") + } + raw := s.conn.RawConn() + if raw == nil { + return nil, moerr.NewInternalErrorNoCtx("backend connection is unavailable") + } + if cause := operationContextCause(ctx); cause != nil { + _ = raw.Close() + return nil, newTimeoutConnectErr(cause) + } + joinInterrupt := interruptConnectionOnDone(ctx, raw) + defer joinInterrupt() type result struct { resp *frontend.Packet @@ -220,6 +293,13 @@ func (s *serverConn) HandleHandshake( select { case ret := <-resultC: + // Join cancellation before accepting success. If cancellation won the + // race, the transport may already be closed and cannot enter the tunnel. + joinInterrupt() + if cause := operationContextCause(ctx); cause != nil { + _ = raw.Close() + return nil, newTimeoutConnectErr(errors.Join(ret.err, cause)) + } return ret.resp, ret.err case <-ctx.Done(): // A caller may release or reuse handshakeResp as soon as this method @@ -227,14 +307,48 @@ func (s *serverConn) HandleHandshake( // returning; calling s.Close here would free mysqlProto buffers while // the worker may still be using them. net.Conn.Close unblocks both the // goetty read and frontend write paths. - if conn := s.conn.RawConn(); conn != nil { - _ = conn.Close() - } + _ = raw.Close() <-resultC + joinInterrupt() logutil.Errorf("handshake to cn %s timeout %v, conn ID: %d goId:%d", s.cnServer.addr, timeout, s.connID, goid.Get()) // Return a retryable error with timeout flag set. - return nil, newTimeoutConnectErr(moerr.AttachCause(ctx, context.DeadlineExceeded)) + return nil, newTimeoutConnectErr(moerr.AttachCause(ctx, context.Cause(ctx))) + } +} + +func operationContextCause(ctx context.Context) error { + if cause := context.Cause(ctx); cause != nil { + return cause + } + if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) { + return context.DeadlineExceeded + } + return nil +} + +// interruptConnectionOnDone makes cancellation a terminal event for a +// phase-owned backend transport. These connections are sacrificial until the +// phase succeeds, so closing is both the strongest I/O interrupt and avoids a +// wrapper or deadline check in the steady tunnel hot path. The returned join +// function prevents a late cancellation callback from closing a connection +// after ownership has been handed off. +func interruptConnectionOnDone(ctx context.Context, conn net.Conn) func() { + if ctx == nil || conn == nil { + return func() {} + } + done := make(chan struct{}) + stop := context.AfterFunc(ctx, func() { + _ = conn.Close() + close(done) + }) + var once sync.Once + return func() { + once.Do(func() { + if !stop() { + <-done + } + }) } } @@ -269,6 +383,35 @@ func (s *serverConn) ExecStmt(stmt internalStmt, resp chan<- []byte) (bool, erro return execOK, nil } +func (s *serverConn) ExecStmtContext( + ctx context.Context, + stmt internalStmt, + resp chan<- []byte, +) (bool, error) { + if ctx == nil { + ctx = context.Background() + } + if s == nil || s.conn == nil { + return false, moerr.NewInternalErrorNoCtx("backend connection is unavailable") + } + raw := s.conn.RawConn() + if raw == nil { + return false, moerr.NewInternalErrorNoCtx("backend connection is unavailable") + } + if cause := operationContextCause(ctx); cause != nil { + _ = raw.Close() + return false, cause + } + joinInterrupt := interruptConnectionOnDone(ctx, raw) + ok, err := s.ExecStmt(stmt, resp) + joinInterrupt() + if cause := operationContextCause(ctx); cause != nil { + _ = raw.Close() + return false, errors.Join(err, cause) + } + return ok, err +} + // GetCNServer implements the ServerConn interface. func (s *serverConn) GetCNServer() *CNServer { return s.cnServer @@ -365,28 +508,53 @@ type CNServer struct { // Connect connects to backend server and returns IOSession. func (s *CNServer) Connect(logger *zap.Logger, timeout time.Duration) (goetty.IOSession, error) { - c := goetty.NewIOSession( - goetty.WithSessionCodec(frontend.NewSqlCodec()), - goetty.WithSessionLogger(logger), - ) - owned := true - defer func() { - if owned { - _ = c.Close() - } - }() - err := c.Connect(s.addr, timeout) + return s.ConnectContext(context.Background(), logger, timeout) +} + +func (s *CNServer) ConnectContext( + parent context.Context, + logger *zap.Logger, + timeout time.Duration, +) (goetty.IOSession, error) { + if parent == nil { + parent = context.Background() + } + ctx := parent + cancel := func() {} + if timeout > 0 { + ctx, cancel = context.WithTimeout(parent, timeout) + } + defer cancel() + + network, address, err := parseBackendAddress(s.addr) + if err != nil { + return nil, newConnectErr(err) + } + raw, err := (&net.Dialer{}).DialContext(ctx, network, address) if err != nil { logutil.Errorf("failed to connect to cn server, timeout: %v, conn ID: %d, cn: %s, goId: %d, error: %v", timeout, s.connID, s.addr, goid.Get(), err) if ne, ok := err.(net.Error); ok && ne.Timeout() { return nil, newTimeoutConnectErr(err) } - if errors.Is(err, context.DeadlineExceeded) { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { return nil, newTimeoutConnectErr(err) } return nil, newConnectErr(err) } + c := goetty.NewIOSession( + goetty.WithSessionConn(0, raw), + goetty.WithSessionCodec(frontend.NewSqlCodec()), + goetty.WithSessionLogger(logger), + ) + owned := true + defer func() { + if owned { + _ = c.Close() + } + }() + joinInterrupt := interruptConnectionOnDone(ctx, raw) + defer joinInterrupt() if len(s.salt) != 20 { return nil, moerr.NewInternalErrorNoCtx("salt is empty") } @@ -404,8 +572,30 @@ func (s *CNServer) Connect(logger *zap.Logger, timeout time.Duration) (goetty.IO // When build connection with backend server, proxy send its salt, request // labels and other information to the backend server. if err := c.Write(data, goetty.WriteOptions{Flush: true}); err != nil { + joinInterrupt() + if cause := operationContextCause(ctx); cause != nil { + return nil, newTimeoutConnectErr(errors.Join(err, cause)) + } return nil, err } + joinInterrupt() + if cause := operationContextCause(ctx); cause != nil { + return nil, newTimeoutConnectErr(cause) + } owned = false return c, nil } + +func parseBackendAddress(address string) (string, string, error) { + if !strings.Contains(address, "//") { + return "tcp4", address, nil + } + u, err := url.Parse(address) + if err != nil { + return "", "", err + } + if strings.EqualFold(u.Scheme, "unix") { + return u.Scheme, u.Path, nil + } + return u.Scheme, u.Host, nil +} diff --git a/pkg/proxy/server_conn_test.go b/pkg/proxy/server_conn_test.go index 373debe55d63d..9549f05591e6b 100644 --- a/pkg/proxy/server_conn_test.go +++ b/pkg/proxy/server_conn_test.go @@ -632,6 +632,94 @@ func TestCNServerConnectClosesSessionOnInvalidSalt(t *testing.T) { require.ErrorIs(t, readErr, io.EOF) } +func TestServerConnExecStmtContextInterruptsSessionTimeout(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + frontend.InitServerLevelVars("test") + transport := goetty.NewIOSession( + goetty.WithSessionConn(1, local), + goetty.WithSessionCodec(frontend.NewSqlCodec()), + ) + fp := config.FrontendParameters{} + fp.SetDefaultValues() + pu := config.NewParameterUnit(&fp, nil, nil, nil) + allocator := frontend.NewLeakCheckAllocator() + ios, err := frontend.NewIOSessionWithOptions( + local, + pu, + "test", + frontend.WithIOSessionBufferSize(proxyIOSessionBufferSize), + frontend.WithIOSessionAllocator(allocator), + ) + require.NoError(t, err) + sc := &serverConn{ + conn: transport, + mysqlProto: frontend.NewMysqlClientProtocol("test", 1, ios, 0, &fp), + } + defer func() { + require.NoError(t, sc.Close()) + require.True(t, allocator.CheckBalance()) + require.NoError(t, transport.Close()) + }() + + requestRead := make(chan error, 1) + go func() { + header := make([]byte, frontend.PacketHeaderLength) + if _, err := io.ReadFull(remote, header); err != nil { + requestRead <- err + return + } + length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16) + _, err := io.ReadFull(remote, make([]byte, length)) + requestRead <- err + }() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + started := time.Now() + _, err = sc.ExecStmtContext(ctx, internalStmt{cmdType: cmdQuery, s: "set transferred=1"}, nil) + require.Error(t, err) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Less(t, time.Since(started), time.Second, + "context must override frontend's default 24-hour session read timeout") + require.NoError(t, <-requestRead) +} + +func TestInterruptConnectionOnDoneOwnershipHandoff(t *testing.T) { + t.Run("cancellation closes phase-owned transport", func(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + ctx, cancel := context.WithCancel(context.Background()) + join := interruptConnectionOnDone(ctx, local) + cancel() + join() + _, err := remote.Write([]byte{1}) + require.Error(t, err) + }) + + t.Run("successful handoff stops late cancellation", func(t *testing.T) { + local, remote := net.Pipe() + defer local.Close() + defer remote.Close() + ctx, cancel := context.WithCancel(context.Background()) + join := interruptConnectionOnDone(ctx, local) + join() + cancel() + + writeDone := make(chan error, 1) + go func() { + _, err := remote.Write([]byte{1}) + writeDone <- err + }() + require.NoError(t, local.SetReadDeadline(time.Now().Add(time.Second))) + buf := make([]byte, 1) + _, err := io.ReadFull(local, buf) + require.NoError(t, err) + require.Equal(t, byte(1), buf[0]) + require.NoError(t, <-writeDone) + }) +} + func TestServerConn_HandleHandshakeEarlyReadFailureIsNotTimeout(t *testing.T) { defer leaktest.AfterTest(t) temp := os.TempDir() From a96c63a744df5d102b0a82e670d0592640aa3296 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Mon, 20 Jul 2026 13:03:39 +0800 Subject: [PATCH 11/15] fix(proxy): close backend memory and admission gaps --- pkg/frontend/mysql_buffer.go | 5 ++ pkg/proxy/client_conn.go | 43 ++++++++- pkg/proxy/client_conn_test.go | 68 ++++++++++++++ pkg/proxy/config.go | 25 +++++- pkg/proxy/config_test.go | 51 +++++++---- pkg/proxy/handshake.go | 6 ++ pkg/proxy/ppv2_test.go | 22 +++++ pkg/proxy/protocol_memory.go | 143 ++++++++++++++++++++++++++---- pkg/proxy/protocol_memory_test.go | 64 ++++++++++--- pkg/proxy/server.go | 11 ++- pkg/proxy/server_conn.go | 88 +++++++++++------- pkg/proxy/server_conn_test.go | 81 ++++++++++++++--- 12 files changed, 508 insertions(+), 99 deletions(-) diff --git a/pkg/frontend/mysql_buffer.go b/pkg/frontend/mysql_buffer.go index 00783034b762f..f5542f726505a 100644 --- a/pkg/frontend/mysql_buffer.go +++ b/pkg/frontend/mysql_buffer.go @@ -324,6 +324,11 @@ func (c *Conn) RawConn() net.Conn { return c.conn } +// GetSequenceID returns the sequence expected for the next MySQL packet. +func (c *Conn) GetSequenceID() uint8 { + return c.sequenceId +} + func (c *Conn) UseConn(conn net.Conn) { c.conn = conn } diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 90eb5e34bc885..67ccaf9784303 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -574,6 +574,7 @@ func (c *clientConn) connAndExec( cn *CNServer, stmt string, resp chan<- []byte, + lane protocolMemoryLane, ) error { if ctx == nil { ctx = c.ctx @@ -581,7 +582,13 @@ func (c *clientConn) connAndExec( ctx = context.Background() } } - lease, err := c.acquireBackendProtocolMemory(ctx, defaultTransferTimeout) + var lease *protocolMemoryLease + var err error + if lane == protocolMemoryBackground { + lease, err = c.acquireBackgroundBackendProtocolMemory(ctx, defaultTransferTimeout) + } else { + lease, err = c.acquireBackendProtocolMemory(ctx, defaultTransferTimeout) + } if err != nil { if resp != nil { c.sendErr(err, resp) @@ -651,6 +658,28 @@ func (c *clientConn) acquireBackendProtocolMemory( return c.acquireProtocolMemory(ctx, timeout, c.protocolMemoryLimiter.budget.backendBytes) } +func (c *clientConn) acquireBackgroundBackendProtocolMemory( + ctx context.Context, + timeout time.Duration, +) (*protocolMemoryLease, error) { + if c.protocolMemoryLimiter == nil { + return nil, nil + } + if ctx == nil { + ctx = c.ctx + if ctx == nil { + ctx = context.Background() + } + } + deadline := time.Now().Add(timeout) + return acquireBackgroundProtocolMemoryBefore( + ctx, + c.protocolMemoryLimiter, + c.protocolMemoryLimiter.budget.backendBytes, + deadline, + ) +} + // KillCurrentBackendConn implements the ClientConn interface. func (c *clientConn) KillCurrentBackendConn(sc ServerConn) error { if sc == nil { @@ -682,7 +711,13 @@ func (c *clientConn) KillCurrentBackendConn(sc ServerConn) error { } ctx, cancel := context.WithTimeout(ctx, defaultTransferTimeout) defer cancel() - return c.connAndExec(ctx, tempCN, fmt.Sprintf("kill connection %d", c.ConnID()), nil) + return c.connAndExec( + ctx, + tempCN, + fmt.Sprintf("kill connection %d", c.ConnID()), + nil, + protocolMemoryCritical, + ) } // handleKill handles the kill event. @@ -719,7 +754,7 @@ func (c *clientConn) handleKill( } ctx, cancel := context.WithTimeout(ctx, defaultTransferTimeout) defer cancel() - return c.connAndExec(ctx, cn, e.stmt, resp) + return c.connAndExec(ctx, cn, e.stmt, resp, protocolMemoryCritical) } // handleSetVar handles the set variable event. @@ -815,7 +850,7 @@ func (c *clientConn) handleUpgradeEvent( // In the loop, do not pass the resp, because it only receives response once. // It everything is ok, send ok response at last out of the loop. - if err := c.connAndExec(ctx, cn, e.stmt, nil); err != nil { + if err := c.connAndExec(ctx, cn, e.stmt, nil, protocolMemoryBackground); err != nil { c.log.Error("failed to execute upgrade query", zap.Error(err)) c.sendErr(err, resp) return err diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 1716ac6fcdda0..1de3c3f1f2b49 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -308,6 +308,21 @@ func (s *killExecServerConn) CreateTime() time.Time { return time.Now() } func (s *killExecServerConn) Quit() error { return nil } func (s *killExecServerConn) Close() error { return nil } +type stalledContextServerConn struct { + killExecServerConn + entered chan struct{} +} + +func (s *stalledContextServerConn) ExecStmtContext( + ctx context.Context, + _ internalStmt, + _ chan<- []byte, +) (bool, error) { + close(s.entered) + <-ctx.Done() + return false, context.Cause(ctx) +} + func testStartClient(t *testing.T, tp *testProxyHandler, ci clientInfo, cn *CNServer) func() { if cn.salt == nil || len(cn.salt) != 20 { cn.salt = testSlat @@ -372,6 +387,59 @@ func TestClientConn_KillCurrentBackendConn(t *testing.T) { require.Equal(t, "kill connection 42", execSC.stmt.s) } +func TestBackgroundExecDoesNotStarveCriticalProtocolAdmission(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 2, + backgroundBytes: 1, + backendBytes: 1, + }) + stalled := &stalledContextServerConn{entered: make(chan struct{})} + router := &killTestRouter{ + connectFn: func(*CNServer, *frontend.Packet, *tunnel) (ServerConn, []byte, error) { + return stalled, makeOKPacket(8), nil + }, + } + client := &clientConn{ + ctx: context.Background(), + log: runtime.DefaultRuntime().Logger(), + router: router, + protocolMemoryLimiter: limiter, + handshakePack: &frontend.Packet{}, + } + + upgradeCtx, cancelUpgrade := context.WithCancel(context.Background()) + upgradeDone := make(chan error, 1) + go func() { + upgradeDone <- client.connAndExec( + upgradeCtx, + &CNServer{}, + "upgrade account all", + nil, + protocolMemoryBackground, + ) + }() + select { + case <-stalled.entered: + case <-time.After(time.Second): + t.Fatal("background control statement did not start") + } + + criticalCtx, cancelCritical := context.WithTimeout(context.Background(), time.Second) + defer cancelCritical() + critical, err := limiter.acquire(criticalCtx, 1) + require.NoError(t, err, "stalled UPGRADE must not consume login and migration admission") + critical.release() + + cancelUpgrade() + select { + case err := <-upgradeDone: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("canceled background control statement did not release its lease") + } + require.Zero(t, limiter.used.Load()) +} + func TestClientConn_HandleQuitEventMarksExpectedCacheQuit(t *testing.T) { tun := &tunnel{} tun.mu.scp = &pipe{} diff --git a/pkg/proxy/config.go b/pkg/proxy/config.go index 089ee3f4e59cf..4a1d774b931d1 100644 --- a/pkg/proxy/config.go +++ b/pkg/proxy/config.go @@ -65,6 +65,9 @@ var ( // their independent per-connection bound is intentionally small. defaultProtocolMemoryLimit = toml.ByteSize(1 << 30) defaultClientHandshakePacketLimit = toml.ByteSize(64 << 10) + // Proxy only consumes the address block; leave bounded room for common TLVs + // without coupling this unauthenticated framing limit to the MySQL login. + defaultProxyProtocolBodyLimit = toml.ByteSize(4 << 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 @@ -74,6 +77,8 @@ var ( minimumClientHandshakePacketLimit = protocol41SSLRequestPayloadSize + 3 // The packet header itself can encode at most (1<<24)-1 payload bytes. maximumClientHandshakePacketLimit = toml.ByteSize((1 << 24) - 1) + minimumProxyProtocolBodyLimit = toml.ByteSize(2*ipv6AddrLength + 4) + maximumProxyProtocolBodyLimit = toml.ByteSize((1 << 16) - 1) ) type RebalancePolicy int @@ -142,13 +147,16 @@ type Config struct { // 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 bounds retained and phase-overlap MySQL protocol + // memory, including shared session buffers, pre-auth framing and login + // packets 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"` + // ProxyProtocolBodyLimit bounds the unauthenticated PROXY v2 address and + // TLV body independently of the MySQL login packet. + ProxyProtocolBodyLimit toml.ByteSize `toml:"proxy-protocol-body-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 @@ -312,6 +320,9 @@ func (c *Config) FillDefault() { if c.ClientHandshakePacketLimit == 0 { c.ClientHandshakePacketLimit = defaultClientHandshakePacketLimit } + if c.ProxyProtocolBodyLimit == 0 { + c.ProxyProtocolBodyLimit = defaultProxyProtocolBodyLimit + } } // Validate validates the configuration of proxy server. @@ -342,6 +353,14 @@ func (c *Config) Validate() error { return moerr.NewInternalError(noReport, "proxy client-handshake-packet-limit exceeds the MySQL packet payload limit") } + if c.ProxyProtocolBodyLimit > 0 && c.ProxyProtocolBodyLimit < minimumProxyProtocolBodyLimit { + return moerr.NewInternalError(noReport, + "proxy proxy-protocol-body-limit is smaller than a TCP/IPv6 address block") + } + if c.ProxyProtocolBodyLimit > maximumProxyProtocolBodyLimit { + return moerr.NewInternalError(noReport, + "proxy proxy-protocol-body-limit exceeds the PROXY v2 body length") + } if _, err := calculateProtocolMemoryBudget(c); err != nil { return err } diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index 8e1ef20e6e7c1..1eed8e6869166 100644 --- a/pkg/proxy/config_test.go +++ b/pkg/proxy/config_test.go @@ -39,6 +39,7 @@ func TestFillDefault(t *testing.T) { require.Equal(t, defaultMaxConnectionsPerTenant, c.MaxConnectionsPerTenant) require.Equal(t, defaultProtocolMemoryLimit, c.ProtocolMemoryLimit) require.Equal(t, defaultClientHandshakePacketLimit, c.ClientHandshakePacketLimit) + require.Equal(t, defaultProxyProtocolBodyLimit, c.ProxyProtocolBodyLimit) require.NoError(t, c.Validate(), "default limits must form a valid retained-memory budget") c = Config{MaxConnections: 1000} @@ -48,7 +49,10 @@ func TestFillDefault(t *testing.T) { } func TestValidate(t *testing.T) { - transientBytes := toml.ByteSize(proxyIOSessionBufferSize) + steadyBytes := toml.ByteSize(10 * (2*proxyIOSessionBufferSize + 64 + + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + + proxyBackendRetainedResponseLimit)) + transientBytes := toml.ByteSize(2 * (proxyIOSessionBufferSize + 2*proxyBackendPacketLimit)) tests := []struct { name string cfg Config @@ -95,7 +99,7 @@ func TestValidate(t *testing.T) { MaxConnections: 10, MaxConnectionsPerTenant: 10, ProtocolMemoryLimit: toml.ByteSize( - 10*(2*proxyIOSessionBufferSize+64) - 1, + steadyBytes - 1, ), ClientHandshakePacketLimit: 64, }, @@ -103,22 +107,18 @@ func TestValidate(t *testing.T) { }, { name: "protocol memory without transient overlap", cfg: Config{ - MaxConnections: 10, - MaxConnectionsPerTenant: 10, - ProtocolMemoryLimit: toml.ByteSize( - 10 * (2*proxyIOSessionBufferSize + 64), - ), + MaxConnections: 10, + MaxConnectionsPerTenant: 10, + ProtocolMemoryLimit: steadyBytes, ClientHandshakePacketLimit: 64, }, wantErr: true, }, { name: "protocol memory exactly covers one transient overlap", cfg: Config{ - MaxConnections: 10, - MaxConnectionsPerTenant: 10, - ProtocolMemoryLimit: toml.ByteSize( - 10*(2*proxyIOSessionBufferSize+64), - ) + transientBytes, + MaxConnections: 10, + MaxConnectionsPerTenant: 10, + ProtocolMemoryLimit: steadyBytes + transientBytes, ClientHandshakePacketLimit: 64, }, }, { @@ -151,7 +151,7 @@ func TestValidate(t *testing.T) { cfg: Config{ MaxConnections: 1, MaxConnectionsPerTenant: 1, - ProtocolMemoryLimit: 64 << 20, + ProtocolMemoryLimit: 80 << 20, ClientHandshakePacketLimit: maximumClientHandshakePacketLimit, }, }, { @@ -160,6 +160,23 @@ func TestValidate(t *testing.T) { ClientHandshakePacketLimit: maximumClientHandshakePacketLimit + 1, }, wantErr: true, + }, { + name: "proxy protocol body below IPv6 minimum", + cfg: Config{ + ProxyProtocolBodyLimit: minimumProxyProtocolBodyLimit - 1, + }, + wantErr: true, + }, { + name: "proxy protocol body at IPv6 minimum", + cfg: Config{ + ProxyProtocolBodyLimit: minimumProxyProtocolBodyLimit, + }, + }, { + name: "proxy protocol body above wire maximum", + cfg: Config{ + ProxyProtocolBodyLimit: maximumProxyProtocolBodyLimit + 1, + }, + wantErr: true, }, { name: "plugin enabled but no backend", cfg: Config{ @@ -185,11 +202,9 @@ func TestValidate(t *testing.T) { }, { name: "plugin mode does not reserve disabled connection cache", cfg: Config{ - MaxConnections: 10, - MaxConnectionsPerTenant: 10, - ProtocolMemoryLimit: toml.ByteSize( - 10*(2*proxyIOSessionBufferSize+64), - ) + transientBytes, + MaxConnections: 10, + MaxConnectionsPerTenant: 10, + ProtocolMemoryLimit: steadyBytes + transientBytes, ClientHandshakePacketLimit: 64, ConnCacheEnabled: true, Plugin: &PluginConfig{ diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index 1009f86eec433..2084359de3255 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -244,6 +244,9 @@ func (s *serverConn) readInitialHandshake() error { if err != nil { return err } + if len(r.Payload) > proxyBackendAuthResponseLimit { + return frontend.ErrPacketTooLarge + } if err := s.parseConnID(r); err != nil { return err } @@ -261,5 +264,8 @@ func (s *serverConn) writeHandshakeResp(handshakeResp *frontend.Packet) (*fronte if err != nil { return nil, err } + if len(data.Payload) > proxyBackendAuthResponseLimit { + return nil, frontend.ErrPacketTooLarge + } return data, nil } diff --git a/pkg/proxy/ppv2_test.go b/pkg/proxy/ppv2_test.go index 1d0c31bd136e1..c3d7daee3de7b 100644 --- a/pkg/proxy/ppv2_test.go +++ b/pkg/proxy/ppv2_test.go @@ -29,6 +29,28 @@ func TestProxyProtocolOptions(t *testing.T) { require.NotNil(t, ret) } +func TestProxyServerCodecAcceptsIPv6AtIndependentMinimum(t *testing.T) { + data := buf.NewByteBuf(ProxyHeaderLength + int(minimumProxyProtocolBodyLimit)) + header := make([]byte, ProxyHeaderLength) + copy(header, ProxyProtocolV2Signature) + header[12] = 0x21 + header[13] = tcpOverIPv6 + binary.BigEndian.PutUint16(header[14:], uint16(minimumProxyProtocolBodyLimit)) + _, err := data.Write(header) + require.NoError(t, err) + _, err = data.Write(make([]byte, minimumProxyProtocolBodyLimit)) + require.NoError(t, err) + + configured := Config{ + ClientHandshakePacketLimit: minimumClientHandshakePacketLimit, + ProxyProtocolBodyLimit: minimumProxyProtocolBodyLimit, + } + res, ok, err := newProxySessionCodec(configured).Decode(data) + require.NoError(t, err) + require.True(t, ok) + require.IsType(t, &ProxyAddr{}, res) +} + func TestProxyProtocolCodec_Decode(t *testing.T) { t.Run("oversized body rejected from fixed header", func(t *testing.T) { data := buf.NewByteBuf(ProxyHeaderLength) diff --git a/pkg/proxy/protocol_memory.go b/pkg/proxy/protocol_memory.go index cf8b725fe3440..ef29cff36b09f 100644 --- a/pkg/proxy/protocol_memory.go +++ b/pkg/proxy/protocol_memory.go @@ -37,11 +37,13 @@ import ( // // The shared allocator remains the byte-level hard limit. Weighted transient // leases prevent individually valid phase transitions from racing for the same -// unreserved headroom, without charging small normal logins for a 16 MiB write. +// unreserved headroom. One backend operation is reserved for the background +// lane so a long-running UPGRADE cannot consume login and migration capacity. type protocolMemoryBudget struct { steadyBytes uint64 headroomBytes uint64 transientBytes uint64 + backgroundBytes uint64 initialBytes uint64 backendBytes uint64 loginDynamicBytes uint64 @@ -61,6 +63,10 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if handshakeLimit == 0 { handshakeLimit = uint64(defaultClientHandshakePacketLimit) } + proxyBodyLimit := uint64(c.ProxyProtocolBodyLimit) + if proxyBodyLimit == 0 { + proxyBodyLimit = uint64(defaultProxyProtocolBodyLimit) + } memoryLimit := uint64(c.ProtocolMemoryLimit) if memoryLimit == 0 { memoryLimit = uint64(defaultProtocolMemoryLimit) @@ -86,7 +92,31 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - retainedLoginBytes, ok := checkedMulUint64(connections, handshakeLimit) + // Before authentication, one Goetty read may contain both a complete PROXY + // frame and the following MySQL login. After authentication, the retained + // login copy replaces that input. Reserve the larger pre-auth state per + // admitted connection so the independent wire limits still compose into a + // global memory bound. + proxyFrameBytes, ok := checkedAddUint64(ProxyHeaderLength, proxyBodyLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + preAuthBytes, ok := checkedAddUint64(proxyFrameBytes, handshakeLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + retainedLoginBytes, ok := checkedMulUint64(connections, preAuthBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + backendSessions, ok := checkedAddUint64(connections, cachedSessions) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + retainedBackendResponseBytes, ok := checkedMulUint64( + backendSessions, + proxyBackendRetainedResponseLimit, + ) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } @@ -94,6 +124,10 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + steadyBytes, ok = checkedAddUint64(steadyBytes, retainedBackendResponseBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } loginDynamicBytes, ok := dynamicProtocolWriteBytes(handshakeLimit) if !ok { @@ -110,6 +144,10 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + initialBytes, ok = checkedAddUint64(initialBytes, proxyBackendPacketLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } backendBytes, ok := checkedAddUint64( proxyIOSessionBufferSize, max(loginDynamicBytes, controlDynamicBytes), @@ -117,7 +155,26 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - transientBytes := max(initialBytes, backendBytes) + // frontend.Conn returns a payload allocation and packetToBytes builds the + // header-bearing response consumed by the event path. Account both while + // they overlap; successful authentication is covered by the smaller initial + // response cap and this is deliberately conservative for control results. + backendReadBytes, ok := checkedMulUint64(2, proxyBackendPacketLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + backendBytes, ok = checkedAddUint64(backendBytes, backendReadBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + criticalBytes := max(initialBytes, backendBytes) + // UPGRADE may legitimately run for the client lifetime. Reserve it a + // separate single-operation lane instead of imposing an arbitrary SQL + // timeout or allowing it to starve login, migration and kill operations. + transientBytes, ok := checkedAddUint64(criticalBytes, backendBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } minimumBytes, ok := checkedAddUint64(steadyBytes, transientBytes) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() @@ -130,6 +187,7 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { steadyBytes: steadyBytes, headroomBytes: memoryLimit - steadyBytes, transientBytes: transientBytes, + backgroundBytes: backendBytes, initialBytes: initialBytes, backendBytes: backendBytes, loginDynamicBytes: loginDynamicBytes, @@ -182,11 +240,21 @@ func roundUpUint64(value, unit uint64) (uint64, bool) { } type protocolMemoryLimiter struct { - weighted *semaphore.Weighted - budget protocolMemoryBudget - used atomic.Int64 + critical *semaphore.Weighted + criticalBytes uint64 + background *semaphore.Weighted + backgroundCapacity uint64 + budget protocolMemoryBudget + used atomic.Int64 } +type protocolMemoryLane uint8 + +const ( + protocolMemoryCritical protocolMemoryLane = iota + protocolMemoryBackground +) + func newProtocolMemoryLimiter(c *Config) (*protocolMemoryLimiter, error) { budget, err := calculateProtocolMemoryBudget(c) if err != nil { @@ -196,9 +264,14 @@ func newProtocolMemoryLimiter(c *Config) (*protocolMemoryLimiter, error) { } func newProtocolMemoryLimiterWithBudget(budget protocolMemoryBudget) *protocolMemoryLimiter { + backgroundBytes := min(budget.backgroundBytes, budget.headroomBytes) + criticalBytes := budget.headroomBytes - backgroundBytes return &protocolMemoryLimiter{ - weighted: semaphore.NewWeighted(int64(budget.headroomBytes)), - budget: budget, + critical: semaphore.NewWeighted(int64(criticalBytes)), + criticalBytes: criticalBytes, + background: semaphore.NewWeighted(int64(backgroundBytes)), + backgroundCapacity: backgroundBytes, + budget: budget, } } @@ -207,30 +280,53 @@ func newProtocolMemoryLimiterWithBudget(budget protocolMemoryBudget) *protocolMe func (l *protocolMemoryLimiter) acquire( ctx context.Context, bytes uint64, +) (*protocolMemoryLease, error) { + if l == nil { + return nil, nil + } + return l.acquireFrom(ctx, l.critical, l.criticalBytes, bytes) +} + +func (l *protocolMemoryLimiter) acquireBackground( + ctx context.Context, + bytes uint64, +) (*protocolMemoryLease, error) { + if l == nil { + return nil, nil + } + return l.acquireFrom(ctx, l.background, l.backgroundCapacity, bytes) +} + +func (l *protocolMemoryLimiter) acquireFrom( + ctx context.Context, + weighted *semaphore.Weighted, + capacity uint64, + bytes uint64, ) (*protocolMemoryLease, error) { if l == nil || bytes == 0 { return nil, nil } - if bytes > l.budget.headroomBytes || bytes > math.MaxInt64 { + if weighted == nil || bytes > capacity || bytes > math.MaxInt64 { return nil, errProxyConnectionLimit } if ctx == nil { ctx = context.Background() } amount := int64(bytes) - if err := l.weighted.Acquire(ctx, amount); err != nil { + if err := weighted.Acquire(ctx, amount); err != nil { return nil, errors.Join(errProxyConnectionLimit, context.Cause(ctx)) } l.used.Add(amount) - lease := &protocolMemoryLease{limiter: l, bytes: amount} + lease := &protocolMemoryLease{limiter: l, weighted: weighted, bytes: amount} lease.refs.Store(1) return lease, nil } type protocolMemoryLease struct { - limiter *protocolMemoryLimiter - bytes int64 - refs atomic.Int32 + limiter *protocolMemoryLimiter + weighted *semaphore.Weighted + bytes int64 + refs atomic.Int32 } // retain adds another concrete owner (for example, a pipelined prefix whose @@ -265,7 +361,7 @@ func (l *protocolMemoryLease) release() { } if refs == 1 { l.limiter.used.Add(-l.bytes) - l.limiter.weighted.Release(l.bytes) + l.weighted.Release(l.bytes) } return } @@ -328,3 +424,20 @@ func acquireProtocolMemoryBefore( defer cancel() return limiter.acquire(acquireCtx, bytes) } + +func acquireBackgroundProtocolMemoryBefore( + ctx context.Context, + limiter *protocolMemoryLimiter, + bytes uint64, + deadline time.Time, +) (*protocolMemoryLease, error) { + if limiter == nil { + return nil, nil + } + if ctx == nil { + ctx = context.Background() + } + acquireCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + return limiter.acquireBackground(acquireCtx, bytes) +} diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go index 5fff0f6c1ca9f..d69c1bd5cdebd 100644 --- a/pkg/proxy/protocol_memory_test.go +++ b/pkg/proxy/protocol_memory_test.go @@ -26,46 +26,82 @@ import ( ) func TestCalculateProtocolMemoryBudget(t *testing.T) { - t.Run("small login needs one extra backend buffer", func(t *testing.T) { - transient := uint64(proxyIOSessionBufferSize) - steady := uint64(10 * (2*proxyIOSessionBufferSize + 64)) + t.Run("small login reserves critical and background backend capacity", func(t *testing.T) { + const connections = 10 + initial := uint64(64 + proxyBackendPacketLimit) + backend := uint64(proxyIOSessionBufferSize + 2*proxyBackendPacketLimit) + transient := backend + backend + steady := uint64(connections * (2*proxyIOSessionBufferSize + + 64 + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + + proxyBackendRetainedResponseLimit)) cfg := Config{ - MaxConnections: 10, - ProtocolMemoryLimit: toml.ByteSize(steady), + MaxConnections: connections, + ProtocolMemoryLimit: toml.ByteSize(steady + transient - 1), ClientHandshakePacketLimit: 64, } _, err := calculateProtocolMemoryBudget(&cfg) require.Error(t, err) - cfg.ProtocolMemoryLimit += toml.ByteSize(transient) + cfg.ProtocolMemoryLimit++ budget, err := calculateProtocolMemoryBudget(&cfg) require.NoError(t, err) require.Equal(t, steady, budget.steadyBytes) require.Equal(t, transient, budget.transientBytes) require.Equal(t, transient, budget.headroomBytes) - require.Equal(t, uint64(64), budget.initialBytes) - require.Equal(t, uint64(proxyIOSessionBufferSize), budget.backendBytes) + require.Equal(t, backend, budget.backgroundBytes) + require.Equal(t, initial, budget.initialBytes) + require.Equal(t, backend, budget.backendBytes) }) t.Run("large login includes rounded backend dynamic write", func(t *testing.T) { const handshakeLimit = 64 << 10 - steady := uint64(2*proxyIOSessionBufferSize + handshakeLimit) - transient := uint64(2 * handshakeLimit) + steady := uint64(2*proxyIOSessionBufferSize + handshakeLimit + + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + + proxyBackendRetainedResponseLimit) + initial := uint64(3 * handshakeLimit) + backend := uint64(proxyIOSessionBufferSize + 3*handshakeLimit) + transient := backend + backend cfg := Config{ MaxConnections: 1, - ProtocolMemoryLimit: toml.ByteSize(steady + 2*transient), + ProtocolMemoryLimit: toml.ByteSize(steady + transient), ClientHandshakePacketLimit: handshakeLimit, } budget, err := calculateProtocolMemoryBudget(&cfg) require.NoError(t, err) require.Equal(t, steady, budget.steadyBytes) require.Equal(t, transient, budget.transientBytes) - require.Equal(t, 2*transient, budget.headroomBytes) - require.Equal(t, uint64(2*handshakeLimit), budget.initialBytes) - require.Equal(t, uint64(proxyIOSessionBufferSize+handshakeLimit), budget.backendBytes) + require.Equal(t, transient, budget.headroomBytes) + require.Equal(t, backend, budget.backgroundBytes) + require.Equal(t, initial, budget.initialBytes) + require.Equal(t, backend, budget.backendBytes) }) } +func TestProtocolMemoryLimiterSeparatesBackgroundAdmission(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 2, + backgroundBytes: 1, + backendBytes: 1, + }) + upgrade, err := limiter.acquireBackground(context.Background(), 1) + require.NoError(t, err) + + critical, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err, "a stalled background operation must not starve login or migration") + critical.release() + + waitCtx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + _, err = limiter.acquireBackground(waitCtx, 1) + require.ErrorIs(t, err, context.DeadlineExceeded) + + upgrade.release() + reused, err := limiter.acquireBackground(context.Background(), 1) + require.NoError(t, err) + reused.release() + require.Zero(t, limiter.used.Load()) +} + func TestProtocolMemoryLimiterLifecycle(t *testing.T) { limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 100}) first, err := limiter.acquire(context.Background(), 60) diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index 7faf74f89aa6c..ae06bae7465e6 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -19,6 +19,7 @@ import ( "time" "github.com/fagongzi/goetty/v2" + "github.com/fagongzi/goetty/v2/codec" "go.uber.org/zap" "github.com/matrixorigin/matrixone/pkg/common/moerr" @@ -118,9 +119,7 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err goetty.WithAppLogger(s.runtime.Logger().RawLogger()), goetty.WithAppHandleSessionFunc(s.handler.handle), goetty.WithAppSessionOptions( - goetty.WithSessionCodec(WithProxyProtocolCodec(frontend.NewSqlCodec( - frontend.WithSQLCodecMaxPayloadSize(int(config.ClientHandshakePacketLimit)), - ), WithProxyProtocolMaxBodySize(int(config.ClientHandshakePacketLimit)))), + goetty.WithSessionCodec(newProxySessionCodec(config)), goetty.WithSessionLogger(s.runtime.Logger().RawLogger()), ), ) @@ -131,6 +130,12 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err return s, nil } +func newProxySessionCodec(config Config) codec.Codec { + return WithProxyProtocolCodec(frontend.NewSqlCodec( + frontend.WithSQLCodecMaxPayloadSize(int(config.ClientHandshakePacketLimit)), + ), WithProxyProtocolMaxBodySize(int(config.ProxyProtocolBodyLimit))) +} + func runBootstrapTask(ctx context.Context, st *stopper.Stopper, h *handler) error { return st.RunNamedTask("proxy bootstrap", func(taskCtx context.Context) { bootstrapCtx, cancel := context.WithCancelCause(ctx) diff --git a/pkg/proxy/server_conn.go b/pkg/proxy/server_conn.go index f3c943aef14a2..e23b8e7bbe89e 100644 --- a/pkg/proxy/server_conn.go +++ b/pkg/proxy/server_conn.go @@ -25,7 +25,6 @@ import ( "sync/atomic" "time" - "github.com/fagongzi/goetty/v2" "github.com/petermattis/goid" "go.uber.org/zap" @@ -77,7 +76,7 @@ type serverConn struct { // cnServer is the backend CN server. cnServer *CNServer // conn is the raw TCP connection between proxy and server. - conn goetty.IOSession + conn net.Conn // connID is the proxy connection ID, which is not useful in most case. connID uint32 // mysqlProto is used to build handshake info. @@ -131,9 +130,21 @@ func execStmtWithContext( return sc.ExecStmt(stmt, resp) } -// Proxy only needs a small retained buffer for normal MySQL packets. Larger -// packets use frontend.Conn's dynamic path and are released after use. -const proxyIOSessionBufferSize = 16 * 1024 +const ( + // Proxy only needs a small retained buffer for normal MySQL packets. + // Larger packets use frontend.Conn's dynamic path and are released after use. + proxyIOSessionBufferSize = 16 * 1024 + // A backend connection is used only for authentication and bounded Proxy + // control statements before the raw tunnel owns it. It must not accept the + // general 16 MiB MySQL payload maximum: doing so lets one CN response create + // an unaccounted, connection-lifetime buffer in the Proxy. + proxyBackendPacketLimit = 64 * 1024 + // A successful authentication response is retained for connection-cache + // reuse. Normal OK packets are tiny; cap this distinct lifetime class so a + // CN cannot turn the transient packet allowance into steady memory. + proxyBackendAuthResponseLimit = 1024 + proxyBackendRetainedResponseLimit = proxyBackendAuthResponseLimit + frontend.PacketHeaderLength +) // newServerConn creates a connection to CN server. func newServerConn( @@ -190,10 +201,11 @@ func newServerConnContext( allocator = frontend.NewSessionAllocator(pu) } ios, err := frontend.NewIOSessionWithOptions( - c.RawConn(), + c, pu, cn.uuid, frontend.WithIOSessionBufferSize(proxyIOSessionBufferSize), + frontend.WithIOSessionAllowedPacketSize(proxyBackendPacketLimit), frontend.WithIOSessionAllocator(allocator), ) if err != nil { @@ -227,16 +239,16 @@ func (s *serverConn) ConnID() uint32 { // RawConn implements the ServerConn interface. func (s *serverConn) RawConn() net.Conn { - if s != nil { - if s.cnServer != nil { + if s != nil && s.conn != nil { + if s.cnServer != nil && s.rebalancer != nil { return &wrappedConn{ - Conn: s.conn.RawConn(), + Conn: s.conn, closeFn: func() { s.rebalancer.connManager.disconnect(s.cnServer, s.tun) }, } } - return s.conn.RawConn() + return s.conn } return nil } @@ -261,7 +273,7 @@ func (s *serverConn) HandleHandshakeContext( if s == nil || s.conn == nil { return nil, moerr.NewInternalErrorNoCtx("backend connection is unavailable") } - raw := s.conn.RawConn() + raw := s.conn if raw == nil { return nil, moerr.NewInternalErrorNoCtx("backend connection is unavailable") } @@ -394,7 +406,7 @@ func (s *serverConn) ExecStmtContext( if s == nil || s.conn == nil { return false, moerr.NewInternalErrorNoCtx("backend connection is unavailable") } - raw := s.conn.RawConn() + raw := s.conn if raw == nil { return false, moerr.NewInternalErrorNoCtx("backend connection is unavailable") } @@ -454,6 +466,8 @@ func (s *serverConn) Close() error { _ = tcpConn.Close() } s.mysqlProto.Close() + } else if s.conn != nil { + _ = s.conn.Close() } }) if s.rebalancer != nil { @@ -465,15 +479,20 @@ func (s *serverConn) Close() error { // readPacket reads packet from CN server, usually used in handshake phase. func (s *serverConn) readPacket() (*frontend.Packet, error) { - msg, err := s.conn.Read(goetty.ReadOptions{}) + if s == nil || s.mysqlProto == nil || s.mysqlProto.GetTcpConnection() == nil { + return nil, moerr.NewInternalErrorNoCtx("backend protocol is unavailable") + } + payload, err := s.mysqlProto.GetTcpConnection().Read() if err != nil { return nil, err } - packet, ok := msg.(*frontend.Packet) - if !ok { - return nil, moerr.NewInternalErrorNoCtx("message is not a Packet") - } - return packet, nil + // frontend.Conn advances the expected sequence after each successful read. + sequenceID := s.mysqlProto.GetTcpConnection().GetSequenceID() - 1 + return &frontend.Packet{ + Length: int32(len(payload)), + SequenceID: int8(sequenceID), + Payload: payload, + }, nil } // nextServerConnID increases baseConnID by 1 and returns the result. @@ -506,16 +525,16 @@ type CNServer struct { clientAddr string } -// Connect connects to backend server and returns IOSession. -func (s *CNServer) Connect(logger *zap.Logger, timeout time.Duration) (goetty.IOSession, error) { +// Connect connects to a backend server and writes the Proxy ExtraInfo preface. +func (s *CNServer) Connect(logger *zap.Logger, timeout time.Duration) (net.Conn, error) { return s.ConnectContext(context.Background(), logger, timeout) } func (s *CNServer) ConnectContext( parent context.Context, - logger *zap.Logger, + _ *zap.Logger, timeout time.Duration, -) (goetty.IOSession, error) { +) (net.Conn, error) { if parent == nil { parent = context.Background() } @@ -542,15 +561,10 @@ func (s *CNServer) ConnectContext( } return nil, newConnectErr(err) } - c := goetty.NewIOSession( - goetty.WithSessionConn(0, raw), - goetty.WithSessionCodec(frontend.NewSqlCodec()), - goetty.WithSessionLogger(logger), - ) owned := true defer func() { if owned { - _ = c.Close() + _ = raw.Close() } }() joinInterrupt := interruptConnectionOnDone(ctx, raw) @@ -571,7 +585,7 @@ func (s *CNServer) ConnectContext( } // When build connection with backend server, proxy send its salt, request // labels and other information to the backend server. - if err := c.Write(data, goetty.WriteOptions{Flush: true}); err != nil { + if err := writeAll(raw, data); err != nil { joinInterrupt() if cause := operationContextCause(ctx); cause != nil { return nil, newTimeoutConnectErr(errors.Join(err, cause)) @@ -583,7 +597,21 @@ func (s *CNServer) ConnectContext( return nil, newTimeoutConnectErr(cause) } owned = false - return c, nil + return raw, nil +} + +func writeAll(conn net.Conn, data []byte) error { + for len(data) > 0 { + n, err := conn.Write(data) + if err != nil { + return err + } + if n == 0 { + return io.ErrShortWrite + } + data = data[n:] + } + return nil } func parseBackendAddress(address string) (string, string, error) { diff --git a/pkg/proxy/server_conn_test.go b/pkg/proxy/server_conn_test.go index 9549f05591e6b..11c29b80c18ad 100644 --- a/pkg/proxy/server_conn_test.go +++ b/pkg/proxy/server_conn_test.go @@ -636,10 +636,6 @@ func TestServerConnExecStmtContextInterruptsSessionTimeout(t *testing.T) { local, remote := net.Pipe() defer remote.Close() frontend.InitServerLevelVars("test") - transport := goetty.NewIOSession( - goetty.WithSessionConn(1, local), - goetty.WithSessionCodec(frontend.NewSqlCodec()), - ) fp := config.FrontendParameters{} fp.SetDefaultValues() pu := config.NewParameterUnit(&fp, nil, nil, nil) @@ -653,13 +649,12 @@ func TestServerConnExecStmtContextInterruptsSessionTimeout(t *testing.T) { ) require.NoError(t, err) sc := &serverConn{ - conn: transport, + conn: local, mysqlProto: frontend.NewMysqlClientProtocol("test", 1, ios, 0, &fp), } defer func() { require.NoError(t, sc.Close()) require.True(t, allocator.CheckBalance()) - require.NoError(t, transport.Close()) }() requestRead := make(chan error, 1) @@ -685,6 +680,53 @@ func TestServerConnExecStmtContextInterruptsSessionTimeout(t *testing.T) { require.NoError(t, <-requestRead) } +func TestServerConnRejectsOversizedBackendPacket(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + frontend.InitServerLevelVars("test") + fp := config.FrontendParameters{} + fp.SetDefaultValues() + pu := config.NewParameterUnit(&fp, nil, nil, nil) + allocator := frontend.NewLeakCheckAllocator() + ios, err := frontend.NewIOSessionWithOptions( + local, + pu, + "test", + frontend.WithIOSessionBufferSize(proxyIOSessionBufferSize), + frontend.WithIOSessionAllowedPacketSize(proxyBackendPacketLimit), + frontend.WithIOSessionAllocator(allocator), + ) + require.NoError(t, err) + sc := &serverConn{ + conn: local, + mysqlProto: frontend.NewMysqlClientProtocol("test", 1, ios, 0, &fp), + } + defer func() { + require.NoError(t, sc.Close()) + require.True(t, allocator.CheckBalance()) + }() + + peerDone := make(chan error, 1) + go func() { + header := []byte{ + byte((proxyBackendPacketLimit + 1) & 0xff), + byte(((proxyBackendPacketLimit + 1) >> 8) & 0xff), + byte(((proxyBackendPacketLimit + 1) >> 16) & 0xff), + 0, + } + packetPrefix := append(header, make([]byte, proxyIOSessionBufferSize-len(header))...) + if _, err := remote.Write(packetPrefix); err != nil { + peerDone <- err + return + } + peerDone <- nil + }() + + _, err = sc.readPacket() + require.Error(t, err) + require.NoError(t, <-peerDone) +} + func TestInterruptConnectionOnDoneOwnershipHandoff(t *testing.T) { t.Run("cancellation closes phase-owned transport", func(t *testing.T) { local, remote := net.Pipe() @@ -797,18 +839,33 @@ func TestServerConn_HandleHandshakeTimeoutStopsWorker(t *testing.T) { local, remote := net.Pipe() defer remote.Close() require.NoError(t, remote.SetWriteDeadline(time.Now().Add(time.Second))) - session := goetty.NewIOSession( - goetty.WithSessionConn(1, local), - goetty.WithSessionCodec(frontend.NewSqlCodec()), + fp := config.FrontendParameters{} + fp.SetDefaultValues() + pu := config.NewParameterUnit(&fp, nil, nil, nil) + allocator := frontend.NewLeakCheckAllocator() + ios, err := frontend.NewIOSessionWithOptions( + local, + pu, + "test", + frontend.WithIOSessionBufferSize(proxyIOSessionBufferSize), + frontend.WithIOSessionAllowedPacketSize(proxyBackendPacketLimit), + frontend.WithIOSessionAllocator(allocator), ) - defer session.Close() + require.NoError(t, err) sc := &serverConn{ cnServer: &CNServer{addr: "pipe"}, - conn: session, + conn: local, connID: 1, + mysqlProto: frontend.NewMysqlClientProtocol( + "test", 1, ios, 0, &fp, + ), } + defer func() { + require.NoError(t, sc.Close()) + require.True(t, allocator.CheckBalance()) + }() - _, err := sc.HandleHandshake( + _, err = sc.HandleHandshake( &frontend.Packet{Payload: []byte("auth")}, 20*time.Millisecond, ) From 57a8e88d11c1714f0c8b4fc077f007978cb62dca Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Mon, 20 Jul 2026 21:03:31 +0800 Subject: [PATCH 12/15] fix(proxy): close handshake and tunnel memory gaps --- pkg/proxy/client_conn_test.go | 232 ++++++++++++++++++++++++++++++ pkg/proxy/config.go | 12 +- pkg/proxy/config_test.go | 11 +- pkg/proxy/handler.go | 6 +- pkg/proxy/handshake.go | 76 +++++++--- pkg/proxy/mysql_conn_buf.go | 9 ++ pkg/proxy/protocol_memory.go | 71 +++++++-- pkg/proxy/protocol_memory_test.go | 12 +- pkg/proxy/server_conn.go | 10 +- pkg/proxy/tunnel.go | 12 +- pkg/proxy/tunnel_test.go | 24 ++-- 11 files changed, 406 insertions(+), 69 deletions(-) diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 1de3c3f1f2b49..6385267a04332 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -1454,6 +1454,238 @@ func TestClientConnHandshakeTimeout(t *testing.T) { }) } +func TestClientConnHandshakeContextLifecycle(t *testing.T) { + t.Run("admission precedes the first login read", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + initialBytes: 1, + }) + occupied, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + defer occupied.release() + client.protocolMemoryLimiter = limiter + + buffered := client.conn.(goetty.BufferedIOSession) + login := makeClientHandshakeResp() + _, err = buffered.InBuf().Write(login) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + lease, err := client.handleHandshakeResp(ctx) + lease.release() + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, client.handshakePack, + "no decoded or retained login may exist before admission") + require.Equal(t, len(login), buffered.InBuf().Readable(), + "login input must remain untouched until transient memory is admitted") + require.Equal(t, int64(1), limiter.used.Load()) + }) + + t.Run("near-limit concurrent logins materialize only after admission", func(t *testing.T) { + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + initialBytes: 1, + }) + clients := make([]*clientConn, 2) + cancels := make([]context.CancelFunc, 2) + for i := range clients { + cc, cleanup := createNewClientConn(t) + defer cleanup() + clients[i] = cc.(*clientConn) + clients[i].protocolMemoryLimiter = limiter + input := clients[i].conn.(goetty.BufferedIOSession).InBuf() + _, err := input.Write(makeClientHandshakeResp()) + require.NoError(t, err) + } + + type handshakeResult struct { + index int + lease *protocolMemoryLease + err error + } + start := make(chan struct{}) + results := make(chan handshakeResult, len(clients)) + for i, client := range clients { + ctx, cancel := context.WithCancel(context.Background()) + cancels[i] = cancel + go func(index int, client *clientConn) { + <-start + lease, err := client.handleHandshakeResp(ctx) + results <- handshakeResult{index: index, lease: lease, err: err} + }(i, client) + } + defer func() { + for _, cancel := range cancels { + cancel() + } + }() + close(start) + + var first handshakeResult + select { + case first = <-results: + require.NoError(t, first.err) + case <-time.After(time.Second): + t.Fatal("first admitted login did not complete") + } + loser := 1 - first.index + require.NotNil(t, clients[first.index].handshakePack) + require.Nil(t, clients[loser].handshakePack, + "the waiter must not decode a login before transient admission") + require.Positive(t, + clients[loser].conn.(goetty.BufferedIOSession).InBuf().Readable(), + "the waiter's input backing must remain untouched") + require.Equal(t, int64(1), limiter.used.Load()) + + first.lease.release() + select { + case second := <-results: + require.NoError(t, second.err) + require.Equal(t, loser, second.index) + second.lease.release() + case <-time.After(time.Second): + t.Fatal("waiting login did not proceed after admission release") + } + require.Zero(t, limiter.used.Load()) + }) + + t.Run("cancel interrupts a silent client and releases admission", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + initialBytes: 1, + }) + client.protocolMemoryLimiter = limiter + + local, remote := net.Pipe() + defer remote.Close() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + lease, err := client.handleHandshakeResp(ctx) + lease.release() + result <- err + }() + require.Eventually(t, func() bool { + return limiter.used.Load() == 1 + }, time.Second, time.Millisecond, "handshake did not acquire transient admission") + cancel() + + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("caller cancellation did not interrupt the client read") + } + require.Zero(t, limiter.used.Load()) + }) + + t.Run("cancel interrupts TLS negotiation and releases admission", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + initialBytes: 1, + }) + client.protocolMemoryLimiter = limiter + cert, err := certGen(t.TempDir()) + require.NoError(t, err) + client.tlsConfig, err = frontend.ConstructTLSConfig( + context.Background(), cert.caFile, cert.certFile, cert.keyFile) + require.NoError(t, err) + client.tlsConnectTimeout = time.Second + + local, remote := net.Pipe() + defer remote.Close() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + lease, err := client.handleHandshakeResp(ctx) + lease.release() + result <- err + }() + writeDone := make(chan error, 1) + go func() { + _, err := remote.Write(makeClientSSLRequest( + 1, + frontend.CLIENT_PROTOCOL_41| + frontend.CLIENT_SECURE_CONNECTION| + frontend.CLIENT_SSL, + )) + if err == nil { + // net.Pipe completes this write only after tls.Conn has started + // reading its record header, making the cancellation boundary exact. + _, err = remote.Write([]byte{0x16}) + } + writeDone <- err + }() + require.NoError(t, <-writeDone) + cancel() + + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("caller cancellation did not interrupt TLS negotiation") + } + require.Zero(t, limiter.used.Load()) + }) + + t.Run("cancel after handoff cannot close the tunnel transport", func(t *testing.T) { + cc, cleanup := createNewClientConn(t) + defer cleanup() + client := cc.(*clientConn) + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{ + headroomBytes: 1, + initialBytes: 1, + }) + client.protocolMemoryLimiter = limiter + + local, remote := net.Pipe() + defer remote.Close() + client.conn.UseConn(local) + client.mysqlProto.UseConn(local) + + ctx, cancel := context.WithCancel(context.Background()) + writeDone := make(chan error, 1) + go func() { + _, err := remote.Write(makeClientHandshakeResp()) + writeDone <- err + }() + lease, err := client.handleHandshakeResp(ctx) + require.NoError(t, err) + require.NoError(t, <-writeDone) + cancel() + + payload := []byte("tunnel-still-open") + writeDone = make(chan error, 1) + go func() { + _, err := remote.Write(payload) + writeDone <- err + }() + require.NoError(t, client.RawConn().SetReadDeadline(time.Now().Add(time.Second))) + got := make([]byte, len(payload)) + _, err = io.ReadFull(client.RawConn(), got) + require.NoError(t, err) + require.Equal(t, payload, got) + require.NoError(t, <-writeDone) + lease.release() + require.Zero(t, limiter.used.Load()) + }) +} + func TestOversizedHandshakeErrorPreservesPacketSequence(t *testing.T) { for _, sequenceID := range []uint8{1, 2} { t.Run(fmt.Sprintf("sequence-%d", sequenceID), func(t *testing.T) { diff --git a/pkg/proxy/config.go b/pkg/proxy/config.go index 4a1d774b931d1..a237c1f22019a 100644 --- a/pkg/proxy/config.go +++ b/pkg/proxy/config.go @@ -60,10 +60,10 @@ var ( // 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) + // Protocol memory covers both the shared allocator used by client/backend + // protocol sessions and the Go-heap buffers retained by established tunnels. + // The default preserves the 10K connection budget after accounting for both. + defaultProtocolMemoryLimit = toml.ByteSize(2 << 30) defaultClientHandshakePacketLimit = toml.ByteSize(64 << 10) // Proxy only consumes the address block; leave bounded room for common TLVs // without coupling this unauthenticated framing limit to the MySQL login. @@ -148,8 +148,8 @@ type Config struct { // It must not exceed MaxConnections. MaxConnectionsPerTenant int `toml:"max-connections-per-tenant" user_setting:"advanced"` // ProtocolMemoryLimit bounds retained and phase-overlap MySQL protocol - // memory, including shared session buffers, pre-auth framing and login - // packets retained for connection migration. + // memory, including shared session buffers, tunnel message/write buffers, + // pre-auth framing and login packets 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. diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index 1eed8e6869166..e601e8627544f 100644 --- a/pkg/proxy/config_test.go +++ b/pkg/proxy/config_test.go @@ -18,6 +18,7 @@ import ( "testing" "time" + "github.com/matrixorigin/matrixone/pkg/frontend" "github.com/matrixorigin/matrixone/pkg/util/toml" "github.com/stretchr/testify/require" ) @@ -49,10 +50,12 @@ func TestFillDefault(t *testing.T) { } func TestValidate(t *testing.T) { - steadyBytes := toml.ByteSize(10 * (2*proxyIOSessionBufferSize + 64 + + steadyBytes := toml.ByteSize(10 * (2*proxyIOSessionBufferSize + + frontend.PacketHeaderLength + 64 + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + - proxyBackendRetainedResponseLimit)) - transientBytes := toml.ByteSize(2 * (proxyIOSessionBufferSize + 2*proxyBackendPacketLimit)) + proxyBackendRetainedResponseLimit + proxyTunnelBufferSize)) + transientBytes := toml.ByteSize(2 * (proxyIOSessionBufferSize + + 2*proxyBackendPacketLimit)) tests := []struct { name string cfg Config @@ -151,7 +154,7 @@ func TestValidate(t *testing.T) { cfg: Config{ MaxConnections: 1, MaxConnectionsPerTenant: 1, - ProtocolMemoryLimit: 80 << 20, + ProtocolMemoryLimit: 96 << 20, ClientHandshakePacketLimit: maximumClientHandshakePacketLimit, }, }, { diff --git a/pkg/proxy/handler.go b/pkg/proxy/handler.go index 1aaa8cc7873cc..e31dea0d4cd2e 100644 --- a/pkg/proxy/handler.go +++ b/pkg/proxy/handler.go @@ -89,12 +89,8 @@ func newProxyHandler( if err != nil { return nil, err } - protocolMemoryLimit := cfg.ProtocolMemoryLimit - if protocolMemoryLimit == 0 { - protocolMemoryLimit = defaultProtocolMemoryLimit - } frontendParameters := &moconfig.FrontendParameters{ - GuestMmuLimitation: int64(protocolMemoryLimit), + GuestMmuLimitation: int64(protocolMemoryLimiter.budget.managedBytes), } sessionAllocator := frontend.NewSessionAllocator( moconfig.NewParameterUnit(frontendParameters, nil, nil, nil), diff --git a/pkg/proxy/handshake.go b/pkg/proxy/handshake.go index 2084359de3255..62ea89d51f4cb 100644 --- a/pkg/proxy/handshake.go +++ b/pkg/proxy/handshake.go @@ -35,7 +35,39 @@ func (c *clientConn) writeInitialHandshake() error { // handleHandshakeResp receives login information from client and saves it // in proxy end. func (c *clientConn) handleHandshakeResp(ctx context.Context) (*protocolMemoryLease, error) { + if ctx == nil { + ctx = c.ctx + if ctx == nil { + ctx = context.Background() + } + } deadline := time.Now().Add(c.clientHandshakeTimeout) + transientBytes := uint64(0) + if c.protocolMemoryLimiter != nil { + transientBytes = c.protocolMemoryLimiter.budget.initialBytes + } + // Admission must precede the first read. SQLCodec and clientConn each copy + // the login payload before the Goetty input backing can be released, so + // acquiring after parsing would leave the actual peak outside the limiter. + lease, err := acquireProtocolMemoryBefore( + ctx, + c.protocolMemoryLimiter, + transientBytes, + deadline, + ) + if err != nil { + return nil, err + } + leaseOwned := true + defer func() { + if leaseOwned { + lease.release() + } + }() + if cause := operationContextCause(ctx); cause != nil { + return nil, cause + } + rawConn := c.conn.RawConn() deadlineConn := newAbsoluteReadDeadlineConn(rawConn, deadline) c.conn.UseConn(deadlineConn) @@ -51,26 +83,23 @@ func (c *clientConn) handleHandshakeResp(ctx context.Context) (*protocolMemoryLe c.mysqlProto.UseConn(rawConn) } }() - if err := c.handleHandshakeRespBefore(deadline); err != nil { - return nil, err - } - transientBytes := uint64(0) - if c.protocolMemoryLimiter != nil { - transientBytes = c.protocolMemoryLimiter.budget.initialBytes - } - lease, err := acquireProtocolMemoryBefore( - ctx, - c.protocolMemoryLimiter, - transientBytes, - deadline, - ) - if err != nil { + // The client transport is phase-owned until this function succeeds. Closing + // it is the only reliable way to interrupt both Goetty reads and TLS I/O; + // joining the callback before success prevents a late cancel from closing a + // connection after ownership has moved to the tunnel. + joinInterrupt := interruptConnectionOnDone(ctx, rawConn) + defer joinInterrupt() + if err := c.handleHandshakeRespBefore(ctx, deadline); err != nil { return nil, err } if err := c.handoffHandshakeBuffer(lease); err != nil { - lease.release() return nil, err } + joinInterrupt() + if cause := operationContextCause(ctx); cause != nil { + return nil, cause + } + leaseOwned = false return lease, nil } @@ -119,7 +148,7 @@ func (c *clientConn) handoffHandshakeBuffer(lease *protocolMemoryLease) error { return nil } -func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { +func (c *clientConn) handleHandshakeRespBefore(ctx context.Context, deadline time.Time) error { if c.handshakePack != nil { return moerr.NewInvalidInputNoCtx("client handshake has already been processed") } @@ -129,6 +158,9 @@ func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { // packet may be an SSLRequest; the terminal packet must be a login. pack, err := c.readPacketBefore(deadline) if err != nil { + if cause := operationContextCause(ctx); cause != nil { + return cause + } return err } c.mysqlProto.AddSequenceId(1) @@ -150,7 +182,7 @@ func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { } // Parse the phase packet and determine whether it requests TLS. - ssl, err := c.mysqlProto.HandleHandshake(c.ctx, ownedPack.Payload) + ssl, err := c.mysqlProto.HandleHandshake(ctx, ownedPack.Payload) if err != nil { // HandleHandshake may retain a slice of the payload before a later // validation fails. Keep clientConn as the cleanup owner until Close. @@ -165,7 +197,7 @@ func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { if tlsEstablished { return moerr.NewInvalidInputNoCtx("duplicate TLS request during client handshake") } - if err = c.upgradeToTLS(deadline); err != nil { + if err = c.upgradeToTLS(ctx, deadline); err != nil { return err } tlsEstablished = true @@ -193,7 +225,7 @@ func (c *clientConn) handleHandshakeRespBefore(deadline time.Time) error { } // upgradeToTLS upgrades the connection to TLS connection. -func (c *clientConn) upgradeToTLS(handshakeDeadline time.Time) error { +func (c *clientConn) upgradeToTLS(ctx context.Context, handshakeDeadline time.Time) error { if c.tlsConfig == nil { return moerr.NewInternalErrorNoCtx("TLS config is invalid") } @@ -205,9 +237,13 @@ func (c *clientConn) upgradeToTLS(handshakeDeadline time.Time) error { // TLS handshake packet from client might have been read into the buffer, use a wrapped conn to // avoid losing handshake packets. tlsConn := tls.Server(c.conn.(goetty.BufferedIOSession).BufferedConn(), c.tlsConfig) - ctx, cancel := context.WithTimeoutCause(context.Background(), timeout, moerr.CauseUpgradeToTLS) + parentCtx := ctx + ctx, cancel := context.WithTimeoutCause(parentCtx, timeout, moerr.CauseUpgradeToTLS) defer cancel() if err := tlsConn.HandshakeContext(ctx); err != nil { + if cause := operationContextCause(parentCtx); cause != nil { + return cause + } err = moerr.AttachCause(ctx, err) return moerr.NewInternalErrorf(ctx, "TLS handshake error: %v", err) } diff --git a/pkg/proxy/mysql_conn_buf.go b/pkg/proxy/mysql_conn_buf.go index 22c5b59c2d833..9e5b806b8f29a 100644 --- a/pkg/proxy/mysql_conn_buf.go +++ b/pkg/proxy/mysql_conn_buf.go @@ -36,6 +36,15 @@ const ( // in the server-to-client direction. 64KB reduces write syscalls by ~100x // for typical MySQL result set rows (~200-500 bytes each). writeBufLen = 65536 + // proxyTunnelBufferSize is the Go-heap memory retained by an established + // tunnel: one message buffer in each direction and one batched writer for + // server-to-client traffic. + proxyTunnelBufferSize = 2*(defaultBufLen+defaultExtraBufLen) + writeBufLen + // proxyMigrationTunnelBufferSize is the additional Go-heap overlap while a + // replacement backend is installed. The new backend message buffer is + // allocated before the old one becomes unreachable; the writer is moved + // because its client-side destination is unchanged. + proxyMigrationTunnelBufferSize = defaultBufLen + defaultExtraBufLen // MySQL header length is 4 bytes, with 3 bytes data length // and 1 byte sequence number. mysqlHeadLen = 4 diff --git a/pkg/proxy/protocol_memory.go b/pkg/proxy/protocol_memory.go index ef29cff36b09f..223c7c48d02b9 100644 --- a/pkg/proxy/protocol_memory.go +++ b/pkg/proxy/protocol_memory.go @@ -28,8 +28,9 @@ import ( ) // protocolMemoryBudget separates long-lived memory, which is reserved by -// connection admission, from short-lived phase overlap. The minimum transient -// headroom is deliberately sized for the larger of: +// connection admission, from short-lived phase overlap. Steady memory includes +// both shared-allocator sessions and Go-heap tunnel buffers. The minimum +// transient headroom is deliberately sized for the larger of: // // - initial login forwarding: pipelined client prefix + dynamic backend write // - migration/control: one additional backend session + dynamic login or @@ -41,6 +42,7 @@ import ( // lane so a long-running UPGRADE cannot consume login and migration capacity. type protocolMemoryBudget struct { steadyBytes uint64 + managedBytes uint64 headroomBytes uint64 transientBytes uint64 backgroundBytes uint64 @@ -93,15 +95,19 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } // Before authentication, one Goetty read may contain both a complete PROXY - // frame and the following MySQL login. After authentication, the retained - // login copy replaces that input. Reserve the larger pre-auth state per - // admitted connection so the independent wire limits still compose into a - // global memory bound. + // frame and the following header-bearing MySQL login. After authentication, + // the retained login payload replaces that input. Reserve the larger + // pre-auth state per admitted connection so the independent wire limits + // still compose into a global memory bound. proxyFrameBytes, ok := checkedAddUint64(ProxyHeaderLength, proxyBodyLimit) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - preAuthBytes, ok := checkedAddUint64(proxyFrameBytes, handshakeLimit) + loginPacketBytes, ok := checkedAddUint64(frontend.PacketHeaderLength, handshakeLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + preAuthBytes, ok := checkedAddUint64(proxyFrameBytes, loginPacketBytes) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } @@ -128,6 +134,17 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + // MySQLConn and bufio buffers live on the Go heap rather than in the shared + // session allocator, but connection admission makes their count equally + // deterministic. Keep them in the same end-to-end protocol memory budget. + tunnelBytes, ok := checkedMulUint64(connections, proxyTunnelBufferSize) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + steadyBytes, ok = checkedAddUint64(steadyBytes, tunnelBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } loginDynamicBytes, ok := dynamicProtocolWriteBytes(handshakeLimit) if !ok { @@ -140,15 +157,26 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - initialBytes, ok := checkedAddUint64(handshakeLimit, loginDynamicBytes) + // SQLCodec copies the login payload out of Goetty's input and clientConn then + // creates the lifetime-retained copy used for migration. That parsing peak + // owns two extra payloads beyond the input backing reserved as steady state. + initialReadBytes, ok := checkedMulUint64(2, handshakeLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + // During backend authentication, a pipelined client prefix can remain live + // alongside the dynamic login write and backend response. These two phases + // do not overlap, so reserve their maximum rather than summing lifetimes. + initialBackendBytes, ok := checkedAddUint64(handshakeLimit, loginDynamicBytes) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - initialBytes, ok = checkedAddUint64(initialBytes, proxyBackendPacketLimit) + initialBackendBytes, ok = checkedAddUint64(initialBackendBytes, proxyBackendPacketLimit) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - backendBytes, ok := checkedAddUint64( + initialBytes := max(initialReadBytes, initialBackendBytes) + backendWorkBytes, ok := checkedAddUint64( proxyIOSessionBufferSize, max(loginDynamicBytes, controlDynamicBytes), ) @@ -163,10 +191,25 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - backendBytes, ok = checkedAddUint64(backendBytes, backendReadBytes) + backendWorkBytes, ok = checkedAddUint64(backendWorkBytes, backendReadBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + backendHandoffBytes, ok := checkedAddUint64( + proxyIOSessionBufferSize, + proxyMigrationTunnelBufferSize, + ) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + backendHandoffBytes, ok = checkedAddUint64( + backendHandoffBytes, + proxyBackendRetainedResponseLimit, + ) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + backendBytes := max(backendWorkBytes, backendHandoffBytes) criticalBytes := max(initialBytes, backendBytes) // UPGRADE may legitimately run for the client lifetime. Reserve it a // separate single-operation lane instead of imposing an arbitrary SQL @@ -184,7 +227,11 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { "proxy protocol-memory-limit is smaller than steady protocol buffers plus one transient operation") } return protocolMemoryBudget{ - steadyBytes: steadyBytes, + steadyBytes: steadyBytes, + // The shared allocator does not own tunnel buffers. Subtract their + // steady reservation from its emergency ceiling so raising the overall + // budget does not also grant the allocator the same bytes a second time. + managedBytes: memoryLimit - tunnelBytes, headroomBytes: memoryLimit - steadyBytes, transientBytes: transientBytes, backgroundBytes: backendBytes, diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go index d69c1bd5cdebd..e600207902529 100644 --- a/pkg/proxy/protocol_memory_test.go +++ b/pkg/proxy/protocol_memory_test.go @@ -21,6 +21,7 @@ import ( "time" "github.com/matrixorigin/matrixone/pkg/common/runtime" + "github.com/matrixorigin/matrixone/pkg/frontend" "github.com/matrixorigin/matrixone/pkg/util/toml" "github.com/stretchr/testify/require" ) @@ -32,8 +33,9 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { backend := uint64(proxyIOSessionBufferSize + 2*proxyBackendPacketLimit) transient := backend + backend steady := uint64(connections * (2*proxyIOSessionBufferSize + - 64 + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + - proxyBackendRetainedResponseLimit)) + frontend.PacketHeaderLength + 64 + ProxyHeaderLength + + int(defaultProxyProtocolBodyLimit) + proxyBackendRetainedResponseLimit + + proxyTunnelBufferSize)) cfg := Config{ MaxConnections: connections, ProtocolMemoryLimit: toml.ByteSize(steady + transient - 1), @@ -46,6 +48,8 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { budget, err := calculateProtocolMemoryBudget(&cfg) require.NoError(t, err) require.Equal(t, steady, budget.steadyBytes) + require.Equal(t, uint64(cfg.ProtocolMemoryLimit)-uint64(connections*proxyTunnelBufferSize), + budget.managedBytes) require.Equal(t, transient, budget.transientBytes) require.Equal(t, transient, budget.headroomBytes) require.Equal(t, backend, budget.backgroundBytes) @@ -55,9 +59,9 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { t.Run("large login includes rounded backend dynamic write", func(t *testing.T) { const handshakeLimit = 64 << 10 - steady := uint64(2*proxyIOSessionBufferSize + handshakeLimit + + steady := uint64(2*proxyIOSessionBufferSize + frontend.PacketHeaderLength + handshakeLimit + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + - proxyBackendRetainedResponseLimit) + proxyBackendRetainedResponseLimit + proxyTunnelBufferSize) initial := uint64(3 * handshakeLimit) backend := uint64(proxyIOSessionBufferSize + 3*handshakeLimit) transient := backend + backend diff --git a/pkg/proxy/server_conn.go b/pkg/proxy/server_conn.go index e23b8e7bbe89e..6e91554c1f738 100644 --- a/pkg/proxy/server_conn.go +++ b/pkg/proxy/server_conn.go @@ -340,11 +340,11 @@ func operationContextCause(ctx context.Context) error { } // interruptConnectionOnDone makes cancellation a terminal event for a -// phase-owned backend transport. These connections are sacrificial until the -// phase succeeds, so closing is both the strongest I/O interrupt and avoids a -// wrapper or deadline check in the steady tunnel hot path. The returned join -// function prevents a late cancellation callback from closing a connection -// after ownership has been handed off. +// phase-owned transport. Client and backend handshake connections are +// sacrificial until their phase succeeds, so closing is both the strongest I/O +// interrupt and avoids a wrapper or deadline check in the steady tunnel hot +// path. The returned join function prevents a late cancellation callback from +// closing a connection after ownership has been handed off. func interruptConnectionOnDone(ctx context.Context, conn net.Conn) func() { if ctx == nil || conn == nil { return func() {} diff --git a/pkg/proxy/tunnel.go b/pkg/proxy/tunnel.go index 83af4ef7c7ca6..f01ef05664564 100644 --- a/pkg/proxy/tunnel.go +++ b/pkg/proxy/tunnel.go @@ -375,14 +375,16 @@ func (t *tunnel) replaceServerConn(newServerConn *MySQLConn, newSC ServerConn, s // set the new ones. t.mu.serverConn = newServerConn t.mu.sc = newSC + // The writer targets the unchanged client connection, so it is safe to move + // for both transfer modes. Reusing it avoids a 64 KiB allocation and leaves + // no detached writer for the GC after every non-sync migration. + if savedBufDst != nil { + t.mu.serverConn.msgBuf.bufDst = savedBufDst + } if sync { t.mu.csp.dst = t.mu.serverConn t.mu.scp.src = t.mu.serverConn - // Transfer the write buffer to the new server conn's msgBuf. - if savedBufDst != nil { - t.mu.serverConn.msgBuf.bufDst = savedBufDst - } } else { t.mu.csp = t.newPipe(pipeClientToServer, t.mu.clientConn, t.mu.serverConn) t.mu.scp = t.newPipe(pipeServerToClient, t.mu.serverConn, t.mu.clientConn) @@ -697,7 +699,7 @@ func (t *tunnel) newPipe(name string, src, dst *MySQLConn) *pipe { // Enable write batching for the server-to-client direction. // Result sets flow s2c and generate many small write syscalls; // bufDst accumulates them and flushes when the read buffer drains. - if name == pipeServerToClient { + if name == pipeServerToClient && src.msgBuf.bufDst == nil { src.msgBuf.bufDst = bufio.NewWriterSize(dst.Conn, writeBufLen) } return p diff --git a/pkg/proxy/tunnel_test.go b/pkg/proxy/tunnel_test.go index ca783279efdd2..96ddbd35fc0e3 100644 --- a/pkg/proxy/tunnel_test.go +++ b/pkg/proxy/tunnel_test.go @@ -328,24 +328,32 @@ func TestTunnelReplaceConn(t *testing.T) { c, s := tu.getConns() require.Equal(t, clientProxy, c.Conn) require.Equal(t, serverProxy, s.Conn) + require.NotNil(t, s.msgBuf.bufDst) + require.Equal(t, proxyTunnelBufferSize, + cap(c.msgBuf.buf)+cap(s.msgBuf.buf)+s.msgBuf.bufDst.Size(), + "the budget constant must match concrete tunnel buffer capacities") + oldWriter := s.msgBuf.bufDst csp, scp := tu.getPipes() require.NoError(t, csp.pause(ctx)) require.NoError(t, scp.pause(ctx)) newServerProxy, newServer := net.Pipe() + newServerConn := newMySQLConn( + "server", + newServerProxy, + 0, + nil, + nil, + false, 0, + ) tu.replaceServerConn( - newMySQLConn( - "server", - newServerProxy, - 0, - nil, - nil, - false, 0, - ), + newServerConn, nil, false, ) + require.Same(t, oldWriter, newServerConn.msgBuf.bufDst, + "non-sync migration must reuse the writer whose client destination is unchanged") require.NoError(t, tu.kickoff()) go func() { From 29837786c139126ff191349c03097f68e82c371d Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Mon, 20 Jul 2026 22:48:09 +0800 Subject: [PATCH 13/15] fix(proxy): account actual handshake buffer capacity --- pkg/proxy/client_conn_test.go | 57 +++++++-- pkg/proxy/config_test.go | 3 +- pkg/proxy/protocol_memory.go | 113 ++++++++++++++--- pkg/proxy/protocol_memory_test.go | 193 +++++++++++++++++++++++++++++- pkg/proxy/server.go | 67 +++++++++++ 5 files changed, 403 insertions(+), 30 deletions(-) diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index 6385267a04332..b5ab22232c7ce 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -1490,16 +1490,44 @@ func TestClientConnHandshakeContextLifecycle(t *testing.T) { initialBytes: 1, }) clients := make([]*clientConn, 2) + sessions := make([]goetty.IOSession, 2) + remotes := make([]net.Conn, 2) cancels := make([]context.CancelFunc, 2) for i := range clients { - cc, cleanup := createNewClientConn(t) - defer cleanup() - clients[i] = cc.(*clientConn) - clients[i].protocolMemoryLimiter = limiter - input := clients[i].conn.(goetty.BufferedIOSession).InBuf() - _, err := input.Write(makeClientHandshakeResp()) + local, remote := net.Pipe() + allocator := frontend.NewLeakCheckAllocator() + session := goetty.NewIOSession( + goetty.WithSessionConn(uint64(i+1), local), + goetty.WithSessionCodec(newProxySessionCodec(Config{ + ClientHandshakePacketLimit: defaultClientHandshakePacketLimit, + ProxyProtocolBodyLimit: defaultProxyProtocolBodyLimit, + })), + goetty.WithSessionAllocator(newProxyApplicationAllocator(allocator)), + goetty.WithSessionRWBUfferSize( + proxyApplicationReadBufferSize, + proxyApplicationWriteBufferSize, + ), + ) + cc, err := newClientConn( + context.Background(), + &Config{ClientHandshakePacketLimit: defaultClientHandshakePacketLimit}, + runtime.DefaultRuntime().Logger(), newCounterSet(), session, + nil, nil, nil, nil, nil, nil, nil, + withClientConnAllocator(allocator), + withClientConnProtocolMemoryLimiter(limiter), + ) require.NoError(t, err) + clients[i] = cc.(*clientConn) + sessions[i] = session + remotes[i] = remote } + defer func() { + for i := range clients { + _ = clients[i].Close() + _ = sessions[i].Close() + _ = remotes[i].Close() + } + }() type handshakeResult struct { index int @@ -1508,9 +1536,16 @@ func TestClientConnHandshakeContextLifecycle(t *testing.T) { } start := make(chan struct{}) results := make(chan handshakeResult, len(clients)) + writeResults := make([]chan error, len(clients)) for i, client := range clients { ctx, cancel := context.WithCancel(context.Background()) cancels[i] = cancel + writeResults[i] = make(chan error, 1) + go func(index int) { + <-start + _, err := remotes[index].Write(makeClientHandshakeResp()) + writeResults[index] <- err + }(i) go func(index int, client *clientConn) { <-start lease, err := client.handleHandshakeResp(ctx) @@ -1532,12 +1567,15 @@ func TestClientConnHandshakeContextLifecycle(t *testing.T) { t.Fatal("first admitted login did not complete") } loser := 1 - first.index + require.NoError(t, <-writeResults[first.index]) require.NotNil(t, clients[first.index].handshakePack) require.Nil(t, clients[loser].handshakePack, "the waiter must not decode a login before transient admission") - require.Positive(t, - clients[loser].conn.(goetty.BufferedIOSession).InBuf().Readable(), - "the waiter's input backing must remain untouched") + loserInput := clients[loser].conn.(goetty.BufferedIOSession).InBuf() + require.Zero(t, loserInput.Readable(), + "the waiter must not read from its socket before transient admission") + require.Len(t, loserInput.RawBuf(), proxyApplicationReadBufferSize, + "the waiter must retain only its bounded bootstrap input") require.Equal(t, int64(1), limiter.used.Load()) first.lease.release() @@ -1545,6 +1583,7 @@ func TestClientConnHandshakeContextLifecycle(t *testing.T) { case second := <-results: require.NoError(t, second.err) require.Equal(t, loser, second.index) + require.NoError(t, <-writeResults[loser]) second.lease.release() case <-time.After(time.Second): t.Fatal("waiting login did not proceed after admission release") diff --git a/pkg/proxy/config_test.go b/pkg/proxy/config_test.go index e601e8627544f..1db90de91ea8d 100644 --- a/pkg/proxy/config_test.go +++ b/pkg/proxy/config_test.go @@ -53,7 +53,8 @@ func TestValidate(t *testing.T) { steadyBytes := toml.ByteSize(10 * (2*proxyIOSessionBufferSize + frontend.PacketHeaderLength + 64 + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + - proxyBackendRetainedResponseLimit + proxyTunnelBufferSize)) + proxyBackendRetainedResponseLimit + + proxyApplicationSessionPersistentBytes + proxyTunnelBufferSize)) transientBytes := toml.ByteSize(2 * (proxyIOSessionBufferSize + 2*proxyBackendPacketLimit)) tests := []struct { diff --git a/pkg/proxy/protocol_memory.go b/pkg/proxy/protocol_memory.go index 223c7c48d02b9..2c96e4df8c91f 100644 --- a/pkg/proxy/protocol_memory.go +++ b/pkg/proxy/protocol_memory.go @@ -32,14 +32,17 @@ import ( // both shared-allocator sessions and Go-heap tunnel buffers. The minimum // transient headroom is deliberately sized for the larger of: // -// - initial login forwarding: pipelined client prefix + dynamic backend write +// - initial login parsing/forwarding: the outer Goetty read capacity, decoded +// and retained login copies, pipelined client prefix, and dynamic backend +// write // - migration/control: one additional backend session + dynamic login or // captured control-statement write // -// The shared allocator remains the byte-level hard limit. Weighted transient -// leases prevent individually valid phase transitions from racing for the same -// unreserved headroom. One backend operation is reserved for the background -// lane so a long-running UPGRADE cannot consume login and migration capacity. +// The shared allocator caps the buffers it owns. Weighted transient leases also +// cover the Go-heap copies and capacity slack outside that allocator, preventing +// individually valid phase transitions from racing for the same unreserved +// headroom. One backend operation is reserved for the background lane so a +// long-running UPGRADE cannot consume login and migration capacity. type protocolMemoryBudget struct { steadyBytes uint64 managedBytes uint64 @@ -111,7 +114,11 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - retainedLoginBytes, ok := checkedMulUint64(connections, preAuthBytes) + // The outer input exists as soon as Goetty accepts a socket, before the + // handshake phase lease is acquired. Keep its explicit initial capacity in + // steady admission even when unusually small wire limits are configured. + steadyHandshakeBytes := max(preAuthBytes, proxyApplicationReadBufferSize) + steadyHandshakeTotalBytes, ok := checkedMulUint64(connections, steadyHandshakeBytes) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } @@ -126,7 +133,7 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - steadyBytes, ok := checkedAddUint64(fixedBytes, retainedLoginBytes) + steadyBytes, ok := checkedAddUint64(fixedBytes, steadyHandshakeTotalBytes) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } @@ -134,6 +141,21 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + // The application-facing Goetty session keeps only its small bootstrap and + // copy buffers on the Go heap. Grown input is routed through the shared + // allocator after phase admission and is accounted below; the fixed output + // and I/O-copy buffers remain until session close and belong to steady state. + applicationSessionBytes, ok := checkedMulUint64( + connections, + proxyApplicationSessionPersistentBytes, + ) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + steadyBytes, ok = checkedAddUint64(steadyBytes, applicationSessionBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } // MySQLConn and bufio buffers live on the Go heap rather than in the shared // session allocator, but connection admission makes their count equally // deterministic. Keep them in the same end-to-end protocol memory budget. @@ -141,6 +163,10 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + externalSteadyBytes, ok := checkedAddUint64(tunnelBytes, applicationSessionBytes) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } steadyBytes, ok = checkedAddUint64(steadyBytes, tunnelBytes) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() @@ -157,13 +183,52 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } - // SQLCodec copies the login payload out of Goetty's input and clientConn then - // creates the lifetime-retained copy used for migration. That parsing peak - // owns two extra payloads beyond the input backing reserved as steady state. - initialReadBytes, ok := checkedMulUint64(2, handshakeLimit) + // One terminal read can include bytes after the login. They remain in the + // outer input until handoff, so include the largest read-ahead allowed by one + // Goetty socket read when deriving the backing capacity. The capacity bound + // is intentionally based on the allocator's growth contract rather than the + // logical packet length; see proxyApplicationReadCapacityUpperBound. + maxApplicationReadBytes, ok := checkedAddUint64( + preAuthBytes, + proxyApplicationReadChunkSize-1, + ) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + applicationReadCapacity, ok := proxyApplicationReadCapacityUpperBound( + maxApplicationReadBytes, + ) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } + // Goetty allocates a grown input before freeing its previous backing. Cover + // that allocator peak independently from the later parse peak. + applicationReadGrowthBytes, ok := checkedMulUint64(2, applicationReadCapacity) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + // SQLCodec copies the decoded login out of Goetty's input and clientConn + // creates the lifetime-retained allocator copy used for migration. At the + // parsing peak both copies coexist with the still-live outer input backing. + applicationReadParseBytes, ok := checkedMulUint64(2, handshakeLimit) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + applicationReadParseBytes, ok = checkedAddUint64( + applicationReadParseBytes, + applicationReadCapacity, + ) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + initialReadBytes := max(applicationReadGrowthBytes, applicationReadParseBytes) + // steadyHandshakeBytes is already reserved for every admitted connection. + // Subtract that logical input/retained-login state so the transient lease + // charges only the simultaneously live capacity and copies beyond steady. + if initialReadBytes < steadyHandshakeBytes { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + initialReadBytes -= steadyHandshakeBytes // During backend authentication, a pipelined client prefix can remain live // alongside the dynamic login write and backend response. These two phases // do not overlap, so reserve their maximum rather than summing lifetimes. @@ -228,10 +293,10 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { } return protocolMemoryBudget{ steadyBytes: steadyBytes, - // The shared allocator does not own tunnel buffers. Subtract their - // steady reservation from its emergency ceiling so raising the overall - // budget does not also grant the allocator the same bytes a second time. - managedBytes: memoryLimit - tunnelBytes, + // The shared allocator does not own tunnel or application-session inline + // buffers. Subtract their steady reservation from its emergency ceiling + // so raising the overall budget does not grant those bytes twice. + managedBytes: memoryLimit - externalSteadyBytes, headroomBytes: memoryLimit - steadyBytes, transientBytes: transientBytes, backgroundBytes: backendBytes, @@ -241,6 +306,24 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { }, nil } +// proxyApplicationReadCapacityUpperBound bounds Goetty's application-session +// input ByteBuf without allocating it. Before each socket read Goetty requests +// proxyApplicationReadChunkSize writable bytes. On growth, the current pinned +// dependency advances capacity in max(current/2, 256) steps and stops at the +// first value above the requested size. Therefore the new capacity is strictly +// below twice (maximum buffered bytes + one read chunk). The real-ByteBuf +// contract test intentionally fails if a future dependency changes that rule. +func proxyApplicationReadCapacityUpperBound(maxBufferedBytes uint64) (uint64, bool) { + withReadChunk, ok := checkedAddUint64( + maxBufferedBytes, + proxyApplicationReadChunkSize, + ) + if !ok { + return 0, false + } + return checkedMulUint64(2, withReadChunk) +} + // dynamicProtocolWriteBytes mirrors frontend.Conn.AppendPart: WritePacket adds // a four-byte header, consumes the fixed IOSession block first, then rounds one // dynamic allocation up to an IOSession block boundary. diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go index e600207902529..dc4ff3158e6d3 100644 --- a/pkg/proxy/protocol_memory_test.go +++ b/pkg/proxy/protocol_memory_test.go @@ -16,10 +16,15 @@ package proxy import ( "context" + "fmt" + "io" "net" + "reflect" "testing" "time" + "github.com/fagongzi/goetty/v2" + goettybuf "github.com/fagongzi/goetty/v2/buf" "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/frontend" "github.com/matrixorigin/matrixone/pkg/util/toml" @@ -35,7 +40,7 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { steady := uint64(connections * (2*proxyIOSessionBufferSize + frontend.PacketHeaderLength + 64 + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + proxyBackendRetainedResponseLimit + - proxyTunnelBufferSize)) + proxyApplicationSessionPersistentBytes + proxyTunnelBufferSize)) cfg := Config{ MaxConnections: connections, ProtocolMemoryLimit: toml.ByteSize(steady + transient - 1), @@ -48,7 +53,7 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { budget, err := calculateProtocolMemoryBudget(&cfg) require.NoError(t, err) require.Equal(t, steady, budget.steadyBytes) - require.Equal(t, uint64(cfg.ProtocolMemoryLimit)-uint64(connections*proxyTunnelBufferSize), + require.Equal(t, uint64(cfg.ProtocolMemoryLimit)-uint64(connections)*(proxyTunnelBufferSize+proxyApplicationSessionPersistentBytes), budget.managedBytes) require.Equal(t, transient, budget.transientBytes) require.Equal(t, transient, budget.headroomBytes) @@ -61,10 +66,17 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { const handshakeLimit = 64 << 10 steady := uint64(2*proxyIOSessionBufferSize + frontend.PacketHeaderLength + handshakeLimit + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit) + - proxyBackendRetainedResponseLimit + proxyTunnelBufferSize) - initial := uint64(3 * handshakeLimit) + proxyBackendRetainedResponseLimit + proxyApplicationSessionPersistentBytes + + proxyTunnelBufferSize) + preAuth := uint64(frontend.PacketHeaderLength + handshakeLimit + + ProxyHeaderLength + int(defaultProxyProtocolBodyLimit)) + applicationReadCapacity := uint64(2) * (preAuth + + proxyApplicationReadChunkSize - 1 + proxyApplicationReadChunkSize) + readGrowth := 2 * applicationReadCapacity + readParse := uint64(2*handshakeLimit) + applicationReadCapacity + initial := max(readGrowth, readParse) - preAuth backend := uint64(proxyIOSessionBufferSize + 3*handshakeLimit) - transient := backend + backend + transient := max(initial, backend) + backend cfg := Config{ MaxConnections: 1, ProtocolMemoryLimit: toml.ByteSize(steady + transient), @@ -79,6 +91,177 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { require.Equal(t, initial, budget.initialBytes) require.Equal(t, backend, budget.backendBytes) }) + + t.Run("small wire limits still reserve the accepted session input", func(t *testing.T) { + handshakeLimit := int(minimumClientHandshakePacketLimit) + proxyBodyLimit := int(minimumProxyProtocolBodyLimit) + steadyHandshake := max( + frontend.PacketHeaderLength+handshakeLimit+ProxyHeaderLength+proxyBodyLimit, + proxyApplicationReadBufferSize, + ) + expectedSteady := uint64(2*proxyIOSessionBufferSize + steadyHandshake + + proxyBackendRetainedResponseLimit + + proxyApplicationSessionPersistentBytes + proxyTunnelBufferSize) + cfg := Config{ + MaxConnections: 1, + ProtocolMemoryLimit: 16 << 20, + ClientHandshakePacketLimit: toml.ByteSize(handshakeLimit), + ProxyProtocolBodyLimit: toml.ByteSize(proxyBodyLimit), + } + budget, err := calculateProtocolMemoryBudget(&cfg) + require.NoError(t, err) + require.Equal(t, expectedSteady, budget.steadyBytes) + }) +} + +func TestProxyApplicationReadCapacityUpperBound(t *testing.T) { + type testCase struct { + buffered uint64 + fragment int + } + cases := []testCase{ + {buffered: 1, fragment: 1}, + {buffered: 4 << 10, fragment: 1}, + {buffered: 72 << 10, fragment: 17}, + {buffered: 72 << 10, fragment: int(proxyApplicationReadChunkSize)}, + {buffered: 520 << 10, fragment: int(proxyApplicationReadChunkSize)}, + } + for _, test := range cases { + t.Run(fmt.Sprintf("bytes-%d-fragment-%d", test.buffered, test.fragment), func(t *testing.T) { + managed := frontend.NewLeakCheckAllocator() + input := goettybuf.NewByteBuf( + int(proxyApplicationReadBufferSize), + goettybuf.WithMemAllocator(newProxyApplicationAllocator(managed)), + ) + require.True(t, managed.CheckBalance(), + "bootstrap input must not consume bounded capacity before admission") + reader := &fragmentReader{ + remaining: int(test.buffered), + fragment: test.fragment, + } + for reader.remaining > 0 { + _, err := input.ReadFrom(reader) + require.NoError(t, err) + } + + capacity, ok := proxyApplicationReadCapacityUpperBound(test.buffered) + require.True(t, ok) + require.LessOrEqual(t, uint64(len(input.RawBuf())), capacity, + "the budget must dominate the pinned Goetty ByteBuf growth contract") + require.False(t, managed.CheckBalance(), + "grown input must be owned by the shared allocator") + input.Close() + require.True(t, managed.CheckBalance(), + "handshake handoff must release grown input immediately") + input.Close() + require.True(t, managed.CheckBalance(), "repeated Goetty close must be harmless") + }) + } +} + +func TestProxyApplicationSessionBufferContract(t *testing.T) { + local, remote := net.Pipe() + defer remote.Close() + managed := frontend.NewLeakCheckAllocator() + session := goetty.NewIOSession( + goetty.WithSessionConn(1, local), + goetty.WithSessionAllocator(newProxyApplicationAllocator(managed)), + goetty.WithSessionRWBUfferSize( + proxyApplicationReadBufferSize, + proxyApplicationWriteBufferSize, + ), + ) + require.Len(t, session.(goetty.BufferedIOSession).InBuf().RawBuf(), + proxyApplicationReadBufferSize) + require.Len(t, session.OutBuf().RawBuf(), proxyApplicationWriteBufferSize) + + // Goetty does not expose options for its per-session I/O-copy buffers. Pin + // their concrete dependency contract here so an upgrade cannot silently + // invalidate the steady-memory formula. + implementation := reflect.ValueOf(session).Elem() + require.Equal(t, proxyApplicationReadCopyBufferSize, + implementation.FieldByName("readCopyBuf").Len()) + require.Equal(t, proxyApplicationWriteCopyBufferSize, + implementation.FieldByName("writeCopyBuf").Len()) + require.True(t, managed.CheckBalance(), + "all application-session bootstrap buffers must remain inline") + require.NoError(t, session.Close()) + require.True(t, managed.CheckBalance()) +} + +func TestProxyApplicationAllocatorRejectsInvalidGrowth(t *testing.T) { + allocator := newProxyApplicationAllocator(nil) + bootstrap, err := allocator.Alloc(proxyApplicationReadBufferSize) + require.NoError(t, err) + require.Len(t, bootstrap, proxyApplicationReadBufferSize) + allocator.Free(bootstrap) + + _, err = allocator.Alloc(proxyApplicationReadBufferSize + 1) + require.Error(t, err) + _, err = allocator.Alloc(-1) + require.Error(t, err) + require.NotPanics(t, func() { + allocator.Free(make([]byte, proxyApplicationReadBufferSize+1)) + }) + + input := goettybuf.NewByteBuf( + proxyApplicationReadBufferSize, + goettybuf.WithMemAllocator(allocator), + ) + defer input.Close() + _, err = input.ReadFrom(&fragmentReader{remaining: 1, fragment: 1}) + require.Error(t, err, "post-admission growth must fail normally instead of panicking") + require.Len(t, input.RawBuf(), proxyApplicationReadBufferSize) +} + +func TestProtocolMemoryBudgetCoversObservedOuterReadCapacity(t *testing.T) { + const handshakeLimit = 512 << 10 + preAuth := uint64(ProxyHeaderLength) + uint64(defaultProxyProtocolBodyLimit) + + frontend.PacketHeaderLength + handshakeLimit + maxBuffered := preAuth + proxyApplicationReadChunkSize - 1 + input := goettybuf.NewByteBuf(int(proxyApplicationReadBufferSize)) + defer input.Close() + reader := &fragmentReader{ + remaining: int(maxBuffered), + fragment: int(proxyApplicationReadChunkSize), + } + for reader.remaining > 0 { + _, err := input.ReadFrom(reader) + require.NoError(t, err) + } + require.Greater(t, uint64(len(input.RawBuf())), preAuth, + "the regression requires real capacity to exceed logical wire bytes") + + cfg := Config{ + MaxConnections: 1, + ProtocolMemoryLimit: 64 << 20, + ClientHandshakePacketLimit: handshakeLimit, + } + budget, err := calculateProtocolMemoryBudget(&cfg) + require.NoError(t, err) + observedCapacity := uint64(len(input.RawBuf())) + observedReadOverlap := max( + 2*observedCapacity, + observedCapacity+2*handshakeLimit, + ) - max(preAuth, proxyApplicationReadBufferSize) + require.GreaterOrEqual(t, budget.initialBytes, observedReadOverlap) +} + +type fragmentReader struct { + remaining int + fragment int +} + +func (r *fragmentReader) Read(dst []byte) (int, error) { + if r.remaining == 0 { + return 0, io.EOF + } + n := min(len(dst), r.fragment, r.remaining) + for i := range n { + dst[i] = byte(i) + } + r.remaining -= n + return n, nil } func TestProtocolMemoryLimiterSeparatesBackgroundAdmission(t *testing.T) { diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index ae06bae7465e6..616991c3e0d33 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -34,6 +34,21 @@ import ( var statsFamilyName = "proxy counter" +const ( + // These values pin the application-facing Goetty session contract used by + // calculateProtocolMemoryBudget. Goetty allocates the two copy buffers when + // the session is created; the read input is phase-owned and released after + // authentication, while the other three buffers live until session close. + proxyApplicationReadBufferSize = 256 + proxyApplicationWriteBufferSize = 256 + proxyApplicationReadCopyBufferSize = 1024 + proxyApplicationWriteCopyBufferSize = 1024 + proxyApplicationReadChunkSize = 4 << 10 + proxyApplicationSessionPersistentBytes = proxyApplicationWriteBufferSize + + proxyApplicationReadCopyBufferSize + + proxyApplicationWriteCopyBufferSize +) + type Server struct { runtime runtime.Runtime stopper *stopper.Stopper @@ -121,6 +136,15 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err goetty.WithAppSessionOptions( goetty.WithSessionCodec(newProxySessionCodec(config)), goetty.WithSessionLogger(s.runtime.Logger().RawLogger()), + goetty.WithSessionAllocator(newProxyApplicationAllocator( + s.handler.sessionAllocator, + )), + // Keep the dependency defaults explicit because their concrete + // allocations are part of the configured protocol-memory budget. + goetty.WithSessionRWBUfferSize( + int(proxyApplicationReadBufferSize), + int(proxyApplicationWriteBufferSize), + ), ), ) if err != nil { @@ -136,6 +160,49 @@ func newProxySessionCodec(config Config) codec.Codec { ), WithProxyProtocolMaxBodySize(int(config.ProxyProtocolBodyLimit))) } +// proxyApplicationAllocator keeps Goetty's small bootstrap buffers infallible +// before connection admission, then routes every grown protocol buffer through +// the Proxy's shared bounded allocator. Goetty constructs an application +// session before the handler can acquire admission and panics if that initial +// allocation returns an error, so putting the bootstrap allocation itself +// behind the bounded allocator would turn overload into a process crash. +// +// The first socket read always requests a 4 KiB writable region, which moves the +// input above the inline threshold after the handshake lease has been acquired. +// From then on growth and handoff Close have deterministic allocate/free edges. +type proxyApplicationAllocator struct { + managed frontend.Allocator +} + +func newProxyApplicationAllocator(managed frontend.Allocator) *proxyApplicationAllocator { + return &proxyApplicationAllocator{managed: managed} +} + +func (a *proxyApplicationAllocator) Alloc(capacity int) ([]byte, error) { + if capacity < 0 { + return nil, moerr.NewInternalErrorNoCtx("negative proxy application buffer capacity") + } + if capacity <= proxyApplicationReadBufferSize { + return make([]byte, capacity), nil + } + if a == nil || a.managed == nil { + return nil, moerr.NewInternalErrorNoCtx("proxy application allocator is unavailable") + } + return a.managed.Alloc(capacity) +} + +func (a *proxyApplicationAllocator) Free(data []byte) { + if len(data) <= proxyApplicationReadBufferSize { + return + } + if a == nil || a.managed == nil { + // A large buffer cannot be produced by Alloc without a managed owner. + // Keep cleanup non-panicking if a future caller violates that invariant. + return + } + a.managed.Free(data) +} + func runBootstrapTask(ctx context.Context, st *stopper.Stopper, h *handler) error { return st.RunNamedTask("proxy bootstrap", func(taskCtx context.Context) { bootstrapCtx, cancel := context.WithCancelCause(ctx) From 14fefd36dc13bbc22a5f5e3955c36b8f6e93b121 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Tue, 21 Jul 2026 10:32:21 +0800 Subject: [PATCH 14/15] fix(proxy): reject late proxy protocol headers --- pkg/proxy/client_conn.go | 8 ++++++ pkg/proxy/client_conn_test.go | 46 +++++++++++++++++++++++++++++++++++ pkg/proxy/server_conn_test.go | 1 + 3 files changed, 55 insertions(+) diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go index 67ccaf9784303..29414a1149219 100644 --- a/pkg/proxy/client_conn.go +++ b/pkg/proxy/client_conn.go @@ -201,6 +201,10 @@ type clientConn struct { // proxyHeaderReceived prevents repeated PROXY headers from resetting the // handshake read path or growing its recursion depth. proxyHeaderReceived bool + // mysqlPacketReceived closes the PROXY framing phase permanently. A PROXY + // header is valid only at the start of the transport; in particular, it must + // not be accepted from the decrypted stream after an SSLRequest. + mysqlPacketReceived bool // haKeeperClient is the client of HAKeeper. haKeeperClient logservice.ClusterHAKeeperClient // moCluster is the CN server cache, which used to filter CN servers @@ -1194,6 +1198,9 @@ func (c *clientConn) readPacketBefore(deadline time.Time) (*frontend.Packet, err return nil, err } if proxyAddr, ok := msg.(*ProxyAddr); ok { + if c.mysqlPacketReceived { + return nil, moerr.NewInvalidInputNoCtx("PROXY protocol header after MySQL packet") + } if c.proxyHeaderReceived { return nil, moerr.NewInvalidInputNoCtx("duplicate PROXY protocol header") } @@ -1208,6 +1215,7 @@ func (c *clientConn) readPacketBefore(deadline time.Time) (*frontend.Packet, err if !ok { return nil, moerr.NewInternalError(c.ctx, "message is not a Packet") } + c.mysqlPacketReceived = true return packet, nil } } diff --git a/pkg/proxy/client_conn_test.go b/pkg/proxy/client_conn_test.go index b5ab22232c7ce..c7966379b3bd4 100644 --- a/pkg/proxy/client_conn_test.go +++ b/pkg/proxy/client_conn_test.go @@ -771,6 +771,52 @@ func TestClientConnHandshakePhases(t *testing.T) { require.True(t, allocator.CheckBalance()) }) + t.Run("reject PROXY header after TLS request", func(t *testing.T) { + client, session, remote, allocator := newClient(t) + t.Cleanup(func() { + _ = client.Close() + _ = session.Close() + _ = remote.Close() + }) + cert, err := certGen(t.TempDir()) + require.NoError(t, err) + client.tlsConfig, err = frontend.ConstructTLSConfig( + context.Background(), cert.caFile, cert.certFile, cert.keyFile) + require.NoError(t, err) + client.tlsConnectTimeout = time.Second + + capabilities := uint32(frontend.CLIENT_PROTOCOL_41 | + frontend.CLIENT_SECURE_CONNECTION | + frontend.CLIENT_SSL) + proxyHeader := make([]byte, ProxyHeaderLength) + copy(proxyHeader, ProxyProtocolV2Signature) + proxyHeader[12] = 0x20 + proxyHeader[13] = unspec + + clientDone := make(chan error, 1) + go func() { + defer remote.Close() + if _, err := remote.Write(makeClientSSLRequest(1, capabilities)); err != nil { + clientDone <- err + return + } + tlsClient := tls.Client(remote, &tls.Config{InsecureSkipVerify: true}) //nolint:gosec + if err := tlsClient.Handshake(); err != nil { + clientDone <- err + return + } + _, err := tlsClient.Write(proxyHeader) + clientDone <- err + }() + + require.ErrorContains(t, handleHandshakeRespForTest(client), + "PROXY protocol header after MySQL packet") + require.NoError(t, <-clientDone) + require.False(t, client.proxyHeaderReceived) + require.NoError(t, client.Close()) + require.True(t, allocator.CheckBalance()) + }) + t.Run("reject duplicate TLS request", func(t *testing.T) { client, session, remote, allocator := newClient(t) cert, err := certGen(t.TempDir()) diff --git a/pkg/proxy/server_conn_test.go b/pkg/proxy/server_conn_test.go index 11c29b80c18ad..91e5421ae2e8f 100644 --- a/pkg/proxy/server_conn_test.go +++ b/pkg/proxy/server_conn_test.go @@ -838,6 +838,7 @@ func TestServerConn_HandleHandshakeEarlyReadFailureIsNotTimeout(t *testing.T) { func TestServerConn_HandleHandshakeTimeoutStopsWorker(t *testing.T) { local, remote := net.Pipe() defer remote.Close() + frontend.InitServerLevelVars("test") require.NoError(t, remote.SetWriteDeadline(time.Now().Add(time.Second))) fp := config.FrontendParameters{} fp.SetDefaultValues() From deb52ea29c8a8988615c227109d8b58cf7a31c08 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Tue, 21 Jul 2026 13:05:54 +0800 Subject: [PATCH 15/15] fix(proxy): close admission and tunnel lifecycle gaps --- pkg/proxy/connection_limit.go | 81 ++++++++++++++++ pkg/proxy/connection_limit_test.go | 148 +++++++++++++++++++++++++++++ pkg/proxy/handler.go | 24 ++++- pkg/proxy/protocol_memory.go | 10 +- pkg/proxy/protocol_memory_test.go | 22 +++++ pkg/proxy/server.go | 25 +++-- pkg/proxy/server_listener.go | 48 ++++++++++ pkg/proxy/server_listener_linux.go | 31 ++++++ pkg/proxy/server_listener_other.go | 23 +++++ pkg/proxy/tunnel.go | 43 +++++++-- pkg/proxy/tunnel_test.go | 91 ++++++++++++++++++ 11 files changed, 527 insertions(+), 19 deletions(-) create mode 100644 pkg/proxy/server_listener.go create mode 100644 pkg/proxy/server_listener_linux.go create mode 100644 pkg/proxy/server_listener_other.go diff --git a/pkg/proxy/connection_limit.go b/pkg/proxy/connection_limit.go index 02840d544dc84..71ae5947d65e7 100644 --- a/pkg/proxy/connection_limit.go +++ b/pkg/proxy/connection_limit.go @@ -45,6 +45,87 @@ type connectionLease struct { released bool } +// connectionAdmissionListener keeps sockets that have not acquired a global +// connection slot out of Goetty. Accept is the last boundary before Goetty +// materializes an IOSession, inserts it into its session map, and starts a +// handler goroutine, so admission must happen here rather than in the handler. +type connectionAdmissionListener struct { + net.Listener + limiter *connectionLimiter + reject func(net.Conn) +} + +func newConnectionAdmissionListener( + listener net.Listener, + limiter *connectionLimiter, + reject func(net.Conn), +) net.Listener { + return &connectionAdmissionListener{ + Listener: listener, + limiter: limiter, + reject: reject, + } +} + +func (l *connectionAdmissionListener) Accept() (net.Conn, error) { + for { + conn, err := l.Listener.Accept() + if err != nil { + return nil, err + } + lease, ok := l.limiter.acquire() + if ok { + return &connectionAdmissionConn{Conn: conn, lease: lease}, nil + } + if l.reject != nil { + l.reject(conn) + } else { + writeConnectionLimitError(conn) + } + _ = conn.Close() + } +} + +// connectionAdmissionConn owns the listener-acquired lease until the handler +// atomically takes it. If Goetty fails before starting the handler, IOSession +// cleanup closes this connection and releases the lease. Once taken, the +// handler's existing defer is the sole release owner. +type connectionAdmissionConn struct { + net.Conn + mu sync.Mutex + lease *connectionLease +} + +func (c *connectionAdmissionConn) takeAdmission() *connectionLease { + if c == nil { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + lease := c.lease + c.lease = nil + return lease +} + +func (c *connectionAdmissionConn) Close() error { + lease := c.takeAdmission() + if lease != nil { + lease.release() + } + if c == nil || c.Conn == nil { + return nil + } + return c.Conn.Close() +} + +func takeConnectionAdmission(conn net.Conn) (*connectionLease, bool) { + admitted, ok := conn.(*connectionAdmissionConn) + if !ok { + return nil, false + } + return admitted.takeAdmission(), true +} + func newConnectionLimiter(maxTotal, maxPerTenant int) *connectionLimiter { return &connectionLimiter{ maxTotal: maxTotal, diff --git a/pkg/proxy/connection_limit_test.go b/pkg/proxy/connection_limit_test.go index c0eae4dce3f66..0b40f6bbc0335 100644 --- a/pkg/proxy/connection_limit_test.go +++ b/pkg/proxy/connection_limit_test.go @@ -16,6 +16,7 @@ package proxy import ( "encoding/binary" + "errors" "fmt" "io" "net" @@ -34,6 +35,19 @@ type deadlineErrorConn struct { net.Conn } +type scriptedAdmissionListener struct { + accept func() (net.Conn, error) +} + +func (l *scriptedAdmissionListener) Accept() (net.Conn, error) { return l.accept() } +func (l *scriptedAdmissionListener) Close() error { return nil } +func (l *scriptedAdmissionListener) Addr() net.Addr { return testAddr("listener") } + +type testAddr string + +func (a testAddr) Network() string { return "test" } +func (a testAddr) String() string { return string(a) } + func (c *deadlineErrorConn) SetWriteDeadline(time.Time) error { return net.ErrClosed } @@ -183,6 +197,140 @@ func TestConnectionLimiter(t *testing.T) { }) } +func TestConnectionAdmissionListenerRejectsBeforeSessionMaterialization(t *testing.T) { + limiter := newConnectionLimiter(1, 1) + occupied, ok := limiter.acquire() + require.True(t, ok) + defer occupied.release() + + proxySide, peerSide := net.Pipe() + defer peerSide.Close() + sentinel := errors.New("listener stopped") + step := 0 + raw := &scriptedAdmissionListener{accept: func() (net.Conn, error) { + step++ + if step == 1 { + return proxySide, nil + } + return nil, sentinel + }} + rejected := 0 + listener := newConnectionAdmissionListener(raw, limiter, func(net.Conn) { + rejected++ + }) + + conn, err := listener.Accept() + require.Nil(t, conn) + require.ErrorIs(t, err, sentinel) + require.Equal(t, 1, rejected) + require.Equal(t, 1, limiter.total, + "a rejected raw socket must not acquire another connection slot") + + buf := make([]byte, 1) + _, err = peerSide.Read(buf) + require.Error(t, err, "the rejected raw socket must be closed inside Accept") +} + +func TestConnectionAdmissionListenerBoundsBlockedRejections(t *testing.T) { + limiter := newConnectionLimiter(1, 1) + occupied, ok := limiter.acquire() + require.True(t, ok) + defer occupied.release() + + proxySide, peerSide := net.Pipe() + defer peerSide.Close() + sentinel := errors.New("listener stopped") + var accepts atomic.Int64 + raw := &scriptedAdmissionListener{accept: func() (net.Conn, error) { + if accepts.Add(1) == 1 { + return proxySide, nil + } + return nil, sentinel + }} + + entered := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + defer unblock() + listener := newConnectionAdmissionListener(raw, limiter, func(net.Conn) { + close(entered) + <-release + }) + + result := make(chan error, 1) + go func() { + conn, err := listener.Accept() + if conn != nil { + _ = conn.Close() + } + result <- err + }() + + select { + case <-entered: + case <-time.After(time.Second): + t.Fatal("rejection callback did not start") + } + require.Equal(t, int64(1), accepts.Load(), + "a blocked rejection must stop the accept loop before another session can materialize") + require.Equal(t, 1, limiter.total) + select { + case err := <-result: + require.Failf(t, "Accept returned while rejection was blocked", "error: %v", err) + default: + } + + unblock() + select { + case err := <-result: + require.ErrorIs(t, err, sentinel) + case <-time.After(time.Second): + t.Fatal("Accept did not resume after rejection was unblocked") + } + require.Equal(t, int64(2), accepts.Load()) +} + +func TestConnectionAdmissionOwnershipTransfer(t *testing.T) { + t.Run("session close before handler releases listener-owned lease", func(t *testing.T) { + limiter := newConnectionLimiter(1, 1) + proxySide, peerSide := net.Pipe() + defer peerSide.Close() + raw := &scriptedAdmissionListener{accept: func() (net.Conn, error) { + return proxySide, nil + }} + listener := newConnectionAdmissionListener(raw, limiter, nil) + conn, err := listener.Accept() + require.NoError(t, err) + require.Equal(t, 1, limiter.total) + require.NoError(t, conn.Close()) + require.Equal(t, 0, limiter.total) + require.NoError(t, conn.Close()) + require.Equal(t, 0, limiter.total) + }) + + t.Run("handler becomes sole lease owner after take", func(t *testing.T) { + limiter := newConnectionLimiter(1, 1) + proxySide, peerSide := net.Pipe() + defer peerSide.Close() + raw := &scriptedAdmissionListener{accept: func() (net.Conn, error) { + return proxySide, nil + }} + listener := newConnectionAdmissionListener(raw, limiter, nil) + conn, err := listener.Accept() + require.NoError(t, err) + lease, preadmitted := takeConnectionAdmission(conn) + require.True(t, preadmitted) + require.NotNil(t, lease) + require.NoError(t, conn.Close()) + require.Equal(t, 1, limiter.total, + "transport close must not release a handler-owned lease") + lease.release() + lease.release() + require.Equal(t, 0, limiter.total) + }) +} + func TestRewriteProxyError(t *testing.T) { t.Run("connection limit", func(t *testing.T) { err := fmt.Errorf("wrapped: %w", errProxyConnectionLimit) diff --git a/pkg/proxy/handler.go b/pkg/proxy/handler.go index e31dea0d4cd2e..72ed1570eda99 100644 --- a/pkg/proxy/handler.go +++ b/pkg/proxy/handler.go @@ -218,10 +218,20 @@ func (h *handler) handle(c goetty.IOSession) error { v2.ProxyConnectClosedCounter.Inc() }() - admission, ok := h.connectionLimiter.acquire() - if !ok { + admission, preadmitted := takeConnectionAdmission(c.RawConn()) + if !preadmitted { + var ok bool + admission, ok = h.connectionLimiter.acquire() + if !ok { + v2.ProxyConnectRejectCounter.Inc() + writeConnectionLimitError(c.RawConn()) + return nil + } + } else if admission == nil { + // Goetty shutdown may close the raw connection immediately before the + // handler begins. In that ordering the connection wrapper remains the + // lease owner and has already released it. v2.ProxyConnectRejectCounter.Inc() - writeConnectionLimitError(c.RawConn()) return nil } defer admission.release() @@ -364,6 +374,14 @@ func (h *handler) handle(c goetty.IOSession) error { } } +func (h *handler) rejectBeforeSession(conn net.Conn) { + v2.ProxyConnectAcceptedCounter.Inc() + v2.ProxyConnectRejectCounter.Inc() + v2.ProxyConnectClosedCounter.Inc() + h.counterSet.connAccepted.Add(1) + writeConnectionLimitError(conn) +} + func (h *handler) handleTunnelErr(err error, cc ClientConn, t *tunnel, sessionID uint64, goId int64) error { skipCacheQuit := false if getErrorCode(err) == codeClientDisconnect { diff --git a/pkg/proxy/protocol_memory.go b/pkg/proxy/protocol_memory.go index 2c96e4df8c91f..fca2d467b118e 100644 --- a/pkg/proxy/protocol_memory.go +++ b/pkg/proxy/protocol_memory.go @@ -159,7 +159,15 @@ func calculateProtocolMemoryBudget(c *Config) (protocolMemoryBudget, error) { // MySQLConn and bufio buffers live on the Go heap rather than in the shared // session allocator, but connection admission makes their count equally // deterministic. Keep them in the same end-to-end protocol memory budget. - tunnelBytes, ok := checkedMulUint64(connections, proxyTunnelBufferSize) + // Cached server connections retain their originating tunnel through + // serverConn.tun so connManager can untrack them on terminal Close. That + // tunnel still owns both message buffers and the client writer even though + // it no longer consumes a live connection slot. + tunnelOwners, ok := checkedAddUint64(connections, cachedSessions) + if !ok { + return protocolMemoryBudget{}, protocolMemoryConfigOverflow() + } + tunnelBytes, ok := checkedMulUint64(tunnelOwners, proxyTunnelBufferSize) if !ok { return protocolMemoryBudget{}, protocolMemoryConfigOverflow() } diff --git a/pkg/proxy/protocol_memory_test.go b/pkg/proxy/protocol_memory_test.go index dc4ff3158e6d3..327ce33ec8f82 100644 --- a/pkg/proxy/protocol_memory_test.go +++ b/pkg/proxy/protocol_memory_test.go @@ -112,6 +112,28 @@ func TestCalculateProtocolMemoryBudget(t *testing.T) { require.NoError(t, err) require.Equal(t, expectedSteady, budget.steadyBytes) }) + + t.Run("cached connections retain originating tunnel buffers", func(t *testing.T) { + withoutCache := Config{ + MaxConnections: 1, + ProtocolMemoryLimit: 512 << 20, + } + withCache := withoutCache + withCache.ConnCacheEnabled = true + + activeBudget, err := calculateProtocolMemoryBudget(&withoutCache) + require.NoError(t, err) + cachedBudget, err := calculateProtocolMemoryBudget(&withCache) + require.NoError(t, err) + + expectedCachedSteady := uint64(defaultMaxNumTotal) * + (proxyIOSessionBufferSize + proxyBackendRetainedResponseLimit + proxyTunnelBufferSize) + require.Equal(t, expectedCachedSteady, + cachedBudget.steadyBytes-activeBudget.steadyBytes) + require.Equal(t, uint64(defaultMaxNumTotal)*proxyTunnelBufferSize, + activeBudget.managedBytes-cachedBudget.managedBytes, + "external cached tunnel memory must be deducted from the shared allocator ceiling") + }) } func TestProxyApplicationReadCapacityUpperBound(t *testing.T) { diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index 616991c3e0d33..334d1107d78b3 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -16,6 +16,7 @@ package proxy import ( "context" + "net" "time" "github.com/fagongzi/goetty/v2" @@ -130,7 +131,16 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err } s.handler = h - app, err := goetty.NewApplication(config.ListenAddress, nil, + listener, err := newProxyListener(config.ListenAddress) + if err != nil { + return nil, err + } + listener = newConnectionAdmissionListener( + listener, + s.handler.connectionLimiter, + s.handler.rejectBeforeSession, + ) + app, err := goetty.NewApplicationWithListeners([]net.Listener{listener}, nil, goetty.WithAppLogger(s.runtime.Logger().RawLogger()), goetty.WithAppHandleSessionFunc(s.handler.handle), goetty.WithAppSessionOptions( @@ -148,6 +158,7 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err ), ) if err != nil { + _ = listener.Close() return nil, err } s.app = app @@ -160,12 +171,12 @@ func newProxySessionCodec(config Config) codec.Codec { ), WithProxyProtocolMaxBodySize(int(config.ProxyProtocolBodyLimit))) } -// proxyApplicationAllocator keeps Goetty's small bootstrap buffers infallible -// before connection admission, then routes every grown protocol buffer through -// the Proxy's shared bounded allocator. Goetty constructs an application -// session before the handler can acquire admission and panics if that initial -// allocation returns an error, so putting the bootstrap allocation itself -// behind the bounded allocator would turn overload into a process crash. +// proxyApplicationAllocator keeps Goetty's small bootstrap buffers infallible, +// then routes every grown protocol buffer through the Proxy's shared bounded +// allocator. Goetty constructs an application session after listener admission +// but before the handler takes the lease, and panics if that initial allocation +// returns an error. The listener bounds these bootstrap allocations; putting +// them behind a fallible allocator would turn overload into a process crash. // // The first socket read always requests a 4 KiB writable region, which moves the // input above the inline threshold after the handshake lease has been acquired. diff --git a/pkg/proxy/server_listener.go b/pkg/proxy/server_listener.go new file mode 100644 index 0000000000000..80ac5109d8b79 --- /dev/null +++ b/pkg/proxy/server_listener.go @@ -0,0 +1,48 @@ +// Copyright 2026 Matrix Origin +// +// 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 proxy + +import ( + "context" + "net" + "net/url" + "strings" +) + +func newProxyListener(address string) (net.Listener, error) { + network, listenAddress, err := parseProxyListenAddress(address) + if err != nil { + return nil, err + } + listenConfig := &net.ListenConfig{Control: proxyListenControl} + return listenConfig.Listen(context.Background(), network, listenAddress) +} + +// parseProxyListenAddress mirrors the pinned Goetty application address +// contract so moving admission to the listener boundary does not change the +// accepted tcp:// or unix:// configuration forms. +func parseProxyListenAddress(address string) (string, string, error) { + if !strings.Contains(address, "//") { + return "tcp4", address, nil + } + u, err := url.Parse(address) + if err != nil { + return "", "", err + } + if strings.EqualFold(u.Scheme, "unix") { + return u.Scheme, u.Path, nil + } + return u.Scheme, u.Host, nil +} diff --git a/pkg/proxy/server_listener_linux.go b/pkg/proxy/server_listener_linux.go new file mode 100644 index 0000000000000..6526a7ee29bbd --- /dev/null +++ b/pkg/proxy/server_listener_linux.go @@ -0,0 +1,31 @@ +// Copyright 2026 Matrix Origin +// +// 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. + +//go:build linux + +package proxy + +import ( + "syscall" + + "golang.org/x/sys/unix" +) + +// Preserve the socket options used by the pinned Goetty NewApplication path. +func proxyListenControl(_ string, _ string, conn syscall.RawConn) error { + return conn.Control(func(fd uintptr) { + _ = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1) + _ = syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, unix.TCP_FASTOPEN, 1) + }) +} diff --git a/pkg/proxy/server_listener_other.go b/pkg/proxy/server_listener_other.go new file mode 100644 index 0000000000000..5b2f5eead819d --- /dev/null +++ b/pkg/proxy/server_listener_other.go @@ -0,0 +1,23 @@ +// Copyright 2026 Matrix Origin +// +// 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. + +//go:build !linux + +package proxy + +import "syscall" + +func proxyListenControl(_ string, _ string, _ syscall.RawConn) error { + return nil +} diff --git a/pkg/proxy/tunnel.go b/pkg/proxy/tunnel.go index f01ef05664564..0bbb0fa8f4496 100644 --- a/pkg/proxy/tunnel.go +++ b/pkg/proxy/tunnel.go @@ -135,6 +135,10 @@ type tunnel struct { mu struct { sync.Mutex + // closed is the terminal generation state. It shares this lock with + // backend publication so a replacement can never become reachable after + // Close has selected the resources it owns. + closed bool // started indicates that the tunnel has started. started bool // inTransfer means a transfer of server connection is in progress. @@ -353,17 +357,31 @@ func (t *tunnel) kickoff() error { } // replaceServerConn replaces the CN server. -func (t *tunnel) replaceServerConn(newServerConn *MySQLConn, newSC ServerConn, sync bool) { +func (t *tunnel) replaceServerConn(newServerConn *MySQLConn, newSC ServerConn, sync bool) error { t.mu.Lock() + if t.mu.closed { + t.mu.Unlock() + // newSC owns the raw backend transport, connManager registration and + // transient protocol-memory lease. The unpublished MySQL wrapper owns + // only Go-heap buffers and becomes unreachable on return. + if newSC != nil { + _ = newSC.Close() + } else if newServerConn != nil { + _ = newServerConn.Close() + } + return errPipeClosed + } defer t.mu.Unlock() oldServerConn := t.mu.serverConn - // Flush and preserve bufDst before closing old connection. - // bufDst wraps the client conn (unchanged), so it stays valid. + // Preserve bufDst before closing the old connection. It targets the client + // connection, which is unchanged, and may already contain ordered response + // bytes from the old backend. Moving the writer preserves those bytes and + // avoids a blocking network flush while t.mu is held; otherwise Close could + // queue behind the data path it is supposed to terminate. var savedBufDst *bufio.Writer if oldServerConn != nil && oldServerConn.msgBuf != nil && oldServerConn.msgBuf.bufDst != nil { - _ = oldServerConn.flushBufDst() savedBufDst = oldServerConn.msgBuf.bufDst oldServerConn.msgBuf.bufDst = nil // detach before close } @@ -395,6 +413,7 @@ func (t *tunnel) replaceServerConn(newServerConn *MySQLConn, newSC ServerConn, s if leased, ok := newSC.(interface{ promoteProtocolMemory() }); ok { leased.promoteProtocolMemory() } + return nil } // canStartTransfer checks whether the transfer can be started. @@ -403,7 +422,7 @@ func (t *tunnel) canStartTransfer(sync bool) bool { defer t.mu.Unlock() // The tunnel has not started. - if !t.mu.started { + if t.mu.closed || !t.mu.started { return false } @@ -501,7 +520,9 @@ func (t *tunnel) doReplaceConnection(ctx context.Context, sync bool) error { t.logger.Error("failed to get a new connection", zap.Error(err)) return err } - t.replaceServerConn(newConn, newSC, sync) + if err := t.replaceServerConn(newConn, newSC, sync); err != nil { + return err + } t.counterSet.connMigrationSuccess.Add(1) t.logger.Info("transfer to a new CN server", zap.String("addr", newConn.RemoteAddr().String())) @@ -614,6 +635,14 @@ func (t *tunnel) setTransferType(typ transferType) { // Close closes the tunnel. func (t *tunnel) Close() error { t.closeOnce.Do(func() { + // Select the terminal generation and its cleanup resources before any + // cancellation can race a replacement into publishing new state. + t.mu.Lock() + t.mu.closed = true + cc, sc := t.mu.clientConn, t.mu.serverConn + serverC := t.mu.sc + t.mu.Unlock() + if t.ctxCancel != nil { t.ctxCancel() } @@ -621,7 +650,6 @@ func (t *tunnel) Close() error { close(t.reqC) // close(t.respC) - cc, sc := t.getConns() // cc.Close() just only close the raw net connection, and it // is closed in goetty module, so do NOT need to close it here: // cc, sc := t.getConns() @@ -630,7 +658,6 @@ func (t *tunnel) Close() error { } if !t.connCacheEnabled { // close the server connection - serverC := t.getServerConn() if serverC != nil { _ = serverC.Close() } else if sc != nil { diff --git a/pkg/proxy/tunnel_test.go b/pkg/proxy/tunnel_test.go index 96ddbd35fc0e3..39738f2ec367b 100644 --- a/pkg/proxy/tunnel_test.go +++ b/pkg/proxy/tunnel_test.go @@ -15,6 +15,7 @@ package proxy import ( + "bufio" "bytes" "context" "encoding/binary" @@ -24,6 +25,7 @@ import ( "math/rand" "net" "sync" + "sync/atomic" "testing" "testing/iotest" "time" @@ -875,6 +877,95 @@ func TestReplaceServerConn(t *testing.T) { require.Equal(t, "select 1", string(buf[5:n])) } +type closeCountingServerConn struct { + ServerConn + closes atomic.Int64 +} + +func (c *closeCountingServerConn) Close() error { + c.closes.Add(1) + return c.ServerConn.Close() +} + +func TestReplaceServerConnRejectsAfterTunnelClose(t *testing.T) { + defer leaktest.AfterTest(t)() + + rt := runtime.DefaultRuntime() + tu := newTunnel(context.Background(), rt.Logger(), nil) + require.NoError(t, tu.Close()) + + proxySide, peerSide := net.Pipe() + defer peerSide.Close() + backend := &closeCountingServerConn{ServerConn: newMockServerConn(proxySide)} + limiter := newProtocolMemoryLimiterWithBudget(protocolMemoryBudget{headroomBytes: 1}) + lease, err := limiter.acquire(context.Background(), 1) + require.NoError(t, err) + newSC := &protocolMemoryServerConn{ServerConn: backend, lease: lease} + newServerConn := newMySQLConn( + connServerName, + newSC.RawConn(), + 0, + nil, + nil, + false, + 0, + ) + + err = tu.replaceServerConn(newServerConn, newSC, false) + require.ErrorIs(t, err, errPipeClosed) + require.Equal(t, int64(1), backend.closes.Load(), + "the unpublished backend must have exactly one terminal cleanup owner") + require.Zero(t, limiter.used.Load(), + "rejected publication must release transient protocol memory") + require.Nil(t, tu.getServerConn()) + _, installed := tu.getConns() + require.Nil(t, installed, "a closed generation must never publish the replacement wrapper") +} + +func TestReplaceServerConnTransfersBufferedWriterWithoutBlockingFlush(t *testing.T) { + defer leaktest.AfterTest(t)() + + clientProxy, clientPeer := net.Pipe() + defer clientPeer.Close() + oldProxy, oldPeer := net.Pipe() + defer oldPeer.Close() + newProxy, newPeer := net.Pipe() + defer newPeer.Close() + + rt := runtime.DefaultRuntime() + tu := newTunnel(context.Background(), rt.Logger(), nil) + oldServerConn := newMySQLConn(connServerName, oldProxy, 0, nil, nil, false, 0) + writer := bufio.NewWriterSize(clientProxy, 64) + _, err := writer.Write([]byte("pending")) + require.NoError(t, err) + oldServerConn.msgBuf.bufDst = writer + tu.mu.clientConn = newMySQLConn(connClientName, clientProxy, 0, nil, nil, false, 0) + tu.mu.serverConn = oldServerConn + tu.mu.sc = newMockServerConn(oldProxy) + + newSC := newMockServerConn(newProxy) + newServerConn := newMySQLConn(connServerName, newProxy, 0, nil, nil, false, 0) + result := make(chan error, 1) + go func() { + result <- tu.replaceServerConn(newServerConn, newSC, false) + }() + + select { + case err := <-result: + require.NoError(t, err) + case <-time.After(250 * time.Millisecond): + // Releasing the socket makes cleanup deterministic even if a future + // regression reintroduces a blocking Flush under t.mu. + _ = clientPeer.Close() + <-result + t.Fatal("replacement blocked flushing the client data path") + } + require.Same(t, writer, newServerConn.msgBuf.bufDst) + require.Equal(t, len("pending"), writer.Buffered(), + "ordered buffered bytes must move to the new backend wrapper intact") + require.NoError(t, tu.Close()) +} + func TestCheckTxnStatus(t *testing.T) { t.Run("mustOK false", func(t *testing.T) { inTxn, ok := checkTxnStatus(nil, false)