*: reduce memory allocation for plenty of idle connections (#474)

Author: djshow832

pkg/manager/router/router_score.go

@@ -325,12 +325,14 @@ func (router *ScoreBasedRouter) OnBackendChanged(backends map[string]*BackendHea
 }
 
 func (router *ScoreBasedRouter) rebalanceLoop(ctx context.Context) {
+	ticker := time.NewTicker(rebalanceInterval)
 	for {
-		router.rebalance(rebalanceConnsPerLoop)
 		select {
 		case <-ctx.Done():
+			ticker.Stop()
 			return
-		case <-time.After(rebalanceInterval):
+		case <-ticker.C:
+			router.rebalance(rebalanceConnsPerLoop)
 		}
 	}
 }

pkg/proxy/backend/authenticator.go

@@ -8,7 +8,6 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/binary"
-	"fmt"
 	"net"
 	"strings"
 
@@ -36,20 +35,24 @@ const SupportedServerCapabilities = pnet.ClientLongPassword | pnet.ClientFoundRo
 
 // Authenticator handshakes with the client and the backend.
 type Authenticator struct {
+	salt              [20]byte
 	dbname            string // default database name
 	user              string
 	attrs             map[string]string
-	salt              []byte
 	capability        pnet.Capability
 	zstdLevel         int
 	collation         uint8
 	proxyProtocol     bool
 	requireBackendTLS bool
 }
 
-func (auth *Authenticator) String() string {
-	return fmt.Sprintf("user:%s, dbname:%s, capability:%d, collation:%d",
-		auth.user, auth.dbname, auth.capability, auth.collation)
+func NewAuthenticator(config *BCConfig) *Authenticator {
+	auth := &Authenticator{
+		proxyProtocol:     config.ProxyProtocol,
+		requireBackendTLS: config.RequireBackendTLS,
+	}
+	GenerateSalt(&auth.salt)
+	return auth
 }
 
 func (auth *Authenticator) writeProxyProtocol(clientIO, backendIO *pnet.PacketIO) error {

pkg/proxy/backend/backend_conn_mgr.go

@@ -132,9 +132,12 @@ type BackendConnManager struct {
 	backendIO        atomic.Pointer[pnet.PacketIO]
 	backendTLS       *tls.Config
 	handshakeHandler HandshakeHandler
-	ctxmap           sync.Map
-	connectionID     uint64
-	quitSource       ErrorSource
+	ctxmap           struct {
+		sync.Mutex
+		m map[any]any
+	}
+	connectionID uint64
+	quitSource   ErrorSource
 }
 
 // NewBackendConnManager creates a BackendConnManager.
@@ -146,16 +149,13 @@ func NewBackendConnManager(logger *zap.Logger, handshakeHandler HandshakeHandler
 		connectionID:     connectionID,
 		cmdProcessor:     NewCmdProcessor(logger.Named("cp")),
 		handshakeHandler: handshakeHandler,
-		authenticator: &Authenticator{
-			proxyProtocol:     config.ProxyProtocol,
-			requireBackendTLS: config.RequireBackendTLS,
-			salt:              GenerateSalt(20),
-		},
+		authenticator:    NewAuthenticator(config),
 		// There are 2 types of signals, which may be sent concurrently.
 		signalReceived: make(chan signalType, signalTypeNums),
 		redirectResCh:  make(chan *redirectResult, 1),
 		quitSource:     SrcNone,
 	}
+	mgr.ctxmap.m = make(map[any]any)
 	mgr.SetValue(ConnContextKeyConnID, connectionID)
 	return mgr
 }
@@ -618,14 +618,15 @@ func (mgr *BackendConnManager) QuitSource() ErrorSource {
 }
 
 func (mgr *BackendConnManager) SetValue(key, val any) {
-	mgr.ctxmap.Store(key, val)
+	mgr.ctxmap.Lock()
+	mgr.ctxmap.m[key] = val
+	mgr.ctxmap.Unlock()
 }
 
 func (mgr *BackendConnManager) Value(key any) any {
-	v, ok := mgr.ctxmap.Load(key)
-	if !ok {
-		return nil
-	}
+	mgr.ctxmap.Lock()
+	v := mgr.ctxmap.m[key]
+	mgr.ctxmap.Unlock()
 	return v
 }
 

pkg/proxy/backend/backend_conn_mgr_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -1217,3 +1218,22 @@ func TestCloseWhileGracefulClose(t *testing.T) {
 
 	ts.runTests(runners)
 }
+
+func BenchmarkSyncMap(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		var m sync.Map
+		m.Store("1", "1")
+		m.Load("1")
+	}
+}
+
+func BenchmarkLockedMap(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		m := make(map[string]string)
+		var lock sync.Mutex
+		lock.Lock()
+		m["1"] = "1"
+		_ = m["1"]
+		lock.Unlock()
+	}
+}

pkg/proxy/backend/mock_backend_test.go

@@ -13,10 +13,10 @@ import (
 )
 
 type backendConfig struct {
+	salt          [20]byte
 	tlsConfig     *tls.Config
 	authPlugin    string
 	sessionStates string
-	salt          []byte
 	columns       int
 	loops         int
 	params        int

pkg/proxy/backend/testsuite_test.go

@@ -34,7 +34,7 @@ const (
 var (
 	mockUsername      = "test_user"
 	mockDBName        = "test_db"
-	mockSalt          = []byte("01234567890123456789")
+	mockSalt          = [20]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
 	mockAuthData      = []byte("123456")
 	mockToken         = strings.Repeat("t", 512)
 	mockCmdStr        = "str"

pkg/proxy/backend/util.go

@@ -10,13 +10,11 @@ func Uint32N(a uint64) uint64
 
 // Buf generates a random string using ASCII characters but avoid separator character.
 // Ref https://github.com/mysql/mysql-server/blob/5.7/mysys_ssl/crypt_genhash_impl.cc#L435.
-func GenerateSalt(size int) []byte {
-	buf := make([]byte, size)
+func GenerateSalt(buf *[20]byte) {
 	for i := range buf {
 		buf[i] = byte(Uint32N(127))
 		for buf[i] == 0 || buf[i] == byte('$') {
 			buf[i] = byte(Uint32N(127))
 		}
 	}
-	return buf
 }

pkg/proxy/net/packetio.go

@@ -195,15 +195,14 @@ type PacketIO struct {
 	rawConn       net.Conn
 	readWriter    packetReadWriter
 	limitReader   io.LimitedReader // reuse memory to reduce allocation
-	header        []byte           // reuse memory to reduce allocation
 	logger        *zap.Logger
 	remoteAddr    net.Addr
 	wrap          error
+	header        [4]byte // reuse memory to reduce allocation
 }
 
 func NewPacketIO(conn net.Conn, lg *zap.Logger, bufferSize int, opts ...PacketIOption) *PacketIO {
 	p := &PacketIO{
-		header:     make([]byte, 4),
 		rawConn:    conn,
 		logger:     lg,
 		readWriter: newBasicReadWriter(conn, bufferSize),
@@ -243,7 +242,7 @@ func (p *PacketIO) GetSequence() uint8 {
 }
 
 func (p *PacketIO) readOnePacket() ([]byte, bool, error) {
-	if err := ReadFull(p.readWriter, p.header); err != nil {
+	if err := ReadFull(p.readWriter, p.header[:]); err != nil {
 		return nil, false, errors.Wrap(ErrReadConn, err)
 	}
 	p.readWriter.SetSequence(p.header[3] + 1)
@@ -292,7 +291,7 @@ func (p *PacketIO) writeOnePacket(data []byte) (int, bool, error) {
 	p.header[3] = sequence
 	p.readWriter.SetSequence(sequence + 1)
 
-	if _, err := p.readWriter.Write(p.header); err != nil {
+	if _, err := p.readWriter.Write(p.header[:]); err != nil {
 		return 0, more, errors.Wrap(ErrWriteConn, err)
 	}
 

pkg/proxy/net/packetio_mysql.go

@@ -17,14 +17,8 @@ var (
 
 // WriteInitialHandshake writes an initial handshake as a server.
 // It's used for tenant-aware routing and testing.
-func (p *PacketIO) WriteInitialHandshake(capability Capability, salt []byte, authPlugin string, serverVersion string, connID uint64) error {
+func (p *PacketIO) WriteInitialHandshake(capability Capability, salt [20]byte, authPlugin string, serverVersion string, connID uint64) error {
 	saltLen := len(salt)
-	if saltLen < 8 {
-		return ErrSaltNotLongEnough
-	} else if saltLen > 20 {
-		saltLen = 20
-	}
-
 	data := make([]byte, 0, 128)
 
 	// min version 10
@@ -61,14 +55,14 @@ func (p *PacketIO) WriteInitialHandshake(capability Capability, salt []byte, aut
 }
 
 // WriteSwitchRequest writes a switch request to the client. It's only for testing.
-func (p *PacketIO) WriteSwitchRequest(authPlugin string, salt []byte) error {
+func (p *PacketIO) WriteSwitchRequest(authPlugin string, salt [20]byte) error {
 	length := 1 + len(authPlugin) + 1 + len(salt) + 1
 	data := make([]byte, 0, length)
 	// check https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_connection_phase_packets_protocol_auth_switch_request.html
 	data = append(data, byte(AuthSwitchHeader))
 	data = append(data, authPlugin...)
 	data = append(data, 0x00)
-	data = append(data, salt...)
+	data = append(data, salt[:]...)
 	data = append(data, 0x00)
 	return p.WritePacket(data, true)
 }

pkg/proxy/net/packetio_test.go

@@ -76,7 +76,7 @@ func TestPacketIO(t *testing.T) {
 			require.NoError(t, err)
 		},
 		func(t *testing.T, srv *PacketIO) {
-			var salt [40]byte
+			var salt [20]byte
 			var msg []byte
 			var err error
 
@@ -95,10 +95,7 @@ func TestPacketIO(t *testing.T) {
 			}
 
 			// send handshake
-			require.NoError(t, srv.WriteInitialHandshake(0, salt[:], AuthNativePassword, ServerVersion, 100))
-			// salt should not be long enough
-			require.ErrorIs(t, srv.WriteInitialHandshake(0, make([]byte, 4), AuthNativePassword, ServerVersion, 100), ErrSaltNotLongEnough)
-
+			require.NoError(t, srv.WriteInitialHandshake(0, salt, AuthNativePassword, ServerVersion, 100))
 			// expect correct and wrong capability flags
 			_, isSSL, err := srv.ReadSSLRequestOrHandshakeResp()
 			require.NoError(t, err)