proxy: listen on additional addrs (#393)

Signed-off-by: xhe <[email protected]>

Author: xhebox

lib/config/proxy.go

@@ -168,7 +168,6 @@ func (cfg *Config) Clone() *Config {
 }
 
 func (cfg *Config) Check() error {
-
 	if cfg.Workdir == "" {
 		d, err := os.Getwd()
 		if err != nil {

pkg/manager/config/manager.go

@@ -136,16 +136,17 @@ func (e *ConfigManager) Close() error {
 		e.cancel()
 		e.cancel = nil
 	}
-	if e.wch != nil {
-		wcherr = e.wch.Close()
-		e.wch = nil
-	}
 	e.sts.Lock()
 	for _, ch := range e.sts.listeners {
 		close(ch)
 	}
 	e.sts.listeners = nil
 	e.sts.Unlock()
 	e.wg.Wait()
+	// close after all goroutines are done
+	if e.wch != nil {
+		wcherr = e.wch.Close()
+		e.wch = nil
+	}
 	return wcherr
 }

pkg/manager/infosync/info.go

@@ -160,7 +160,8 @@ func (is *InfoSyncer) getTopologyInfo(cfg *config.Config) (*TopologyInfo, error)
 		s = ""
 	}
 	dir := path.Dir(s)
-	ip, port, err := net.SplitHostPort(cfg.Proxy.Addr)
+	addrs := strings.Split(cfg.Proxy.Addr, ",")
+	ip, port, err := net.SplitHostPort(addrs[0])
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

pkg/metrics/metrics_test.go

@@ -20,7 +20,7 @@ import (
 // Test that the metrics are pushed or not pushed with different configurations.
 func TestPushMetrics(t *testing.T) {
 	proxyAddr := "0.0.0.0:6000"
-	labelName := fmt.Sprintf("%s_%s_connections", ModuleProxy, LabelServer)
+	labelName := fmt.Sprintf("%s_%s_maxprocs", ModuleProxy, LabelServer)
 	hostname, err := os.Hostname()
 	require.NoError(t, err)
 	expectedPath := fmt.Sprintf("/metrics/job/tiproxy/instance/%s_6000", hostname)

pkg/proxy/backend/handshake_handler.go

@@ -18,6 +18,7 @@ type ConnContextKey string
 const (
 	ConnContextKeyTLSState ConnContextKey = "tls-state"
 	ConnContextKeyConnID   ConnContextKey = "conn-id"
+	ConnContextKeyConnAddr ConnContextKey = "conn-addr"
 )
 
 type ErrorSource int

pkg/proxy/client/client_conn.go

@@ -23,8 +23,9 @@ type ClientConnection struct {
 }
 
 func NewClientConnection(logger *zap.Logger, conn net.Conn, frontendTLSConfig *tls.Config, backendTLSConfig *tls.Config,
-	hsHandler backend.HandshakeHandler, connID uint64, bcConfig *backend.BCConfig) *ClientConnection {
+	hsHandler backend.HandshakeHandler, connID uint64, addr string, bcConfig *backend.BCConfig) *ClientConnection {
 	bemgr := backend.NewBackendConnManager(logger.Named("be"), hsHandler, connID, bcConfig)
+	bemgr.SetValue(backend.ConnContextKeyConnAddr, addr)
 	opts := make([]pnet.PacketIOption, 0, 2)
 	opts = append(opts, pnet.WithWrapError(backend.ErrClientConn))
 	if bcConfig.ProxyProtocol {

pkg/proxy/proxy.go

@@ -6,6 +6,7 @@ package proxy
 import (
 	"context"
 	"net"
+	"strings"
 	"sync"
 	"time"
 
@@ -37,7 +38,8 @@ type serverState struct {
 }
 
 type SQLServer struct {
-	listener          net.Listener
+	listeners         []net.Listener
+	addrs             []string
 	logger            *zap.Logger
 	certMgr           *cert.CertManager
 	hsHandler         backend.HandshakeHandler
@@ -65,9 +67,13 @@ func NewSQLServer(logger *zap.Logger, cfg config.ProxyServer, certMgr *cert.Cert
 
 	s.reset(&cfg.ProxyServerOnline)
 
-	s.listener, err = net.Listen("tcp", cfg.Addr)
-	if err != nil {
-		return nil, err
+	s.addrs = strings.Split(cfg.Addr, ",")
+	s.listeners = make([]net.Listener, len(s.addrs))
+	for i, addr := range s.addrs {
+		s.listeners[i], err = net.Listen("tcp", addr)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	return s, nil
@@ -104,31 +110,34 @@ func (s *SQLServer) Run(ctx context.Context, cfgch <-chan *config.Config) {
 		}
 	})
 
-	s.wg.Run(func() {
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			default:
-				conn, err := s.listener.Accept()
-				if err != nil {
-					if errors.Is(err, net.ErrClosed) {
-						return
+	for i := range s.listeners {
+		j := i
+		s.wg.Run(func() {
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					conn, err := s.listeners[j].Accept()
+					if err != nil {
+						if errors.Is(err, net.ErrClosed) {
+							return
+						}
+
+						s.logger.Error("accept failed", zap.Error(err))
+						continue
 					}
 
-					s.logger.Error("accept failed", zap.Error(err))
-					continue
+					s.wg.Run(func() {
+						util.WithRecovery(func() { s.onConn(ctx, conn, s.addrs[j]) }, nil, s.logger)
+					})
 				}
-
-				s.wg.Run(func() {
-					util.WithRecovery(func() { s.onConn(ctx, conn) }, nil, s.logger)
-				})
 			}
-		}
-	})
+		})
+	}
 }
 
-func (s *SQLServer) onConn(ctx context.Context, conn net.Conn) {
+func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) {
 	s.mu.Lock()
 	conns := uint64(len(s.mu.clients))
 	maxConns := s.mu.maxConnections
@@ -149,9 +158,9 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn) {
 	connID := s.mu.connID
 	s.mu.connID++
 	logger := s.logger.With(zap.Uint64("connID", connID), zap.String("client_addr", conn.RemoteAddr().String()),
-		zap.Bool("proxy-protocol", s.mu.proxyProtocol))
+		zap.Bool("proxy-protocol", s.mu.proxyProtocol), zap.String("addr", addr))
 	clientConn := client.NewClientConnection(logger.Named("conn"), conn, s.certMgr.ServerTLS(), s.certMgr.SQLTLS(),
-		s.hsHandler, connID, &backend.BCConfig{
+		s.hsHandler, connID, addr, &backend.BCConfig{
 			ProxyProtocol:      s.mu.proxyProtocol,
 			RequireBackendTLS:  s.requireBackendTLS,
 			HealthyKeepAlive:   s.mu.healthyKeepAlive,
@@ -232,8 +241,8 @@ func (s *SQLServer) Close() error {
 		s.cancelFunc = nil
 	}
 	errs := make([]error, 0, 4)
-	if s.listener != nil {
-		errs = append(errs, s.listener.Close())
+	for i := range s.listeners {
+		errs = append(errs, s.listeners[i].Close())
 	}
 
 	s.mu.RLock()

pkg/proxy/proxy_test.go

@@ -6,6 +6,7 @@ package proxy
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"net"
 	"strings"
 	"testing"
@@ -48,13 +49,13 @@ func TestGracefulShutdown(t *testing.T) {
 	createClientConn := func() *client.ClientConnection {
 		server.mu.Lock()
 		go func() {
-			conn, err := net.Dial("tcp", server.listener.Addr().String())
+			conn, err := net.Dial("tcp", server.listeners[0].Addr().String())
 			require.NoError(t, err)
 			require.NoError(t, conn.Close())
 		}()
-		conn, err := server.listener.Accept()
+		conn, err := server.listeners[0].Accept()
 		require.NoError(t, err)
-		clientConn := client.NewClientConnection(lg, conn, nil, nil, hsHandler, 0, &backend.BCConfig{})
+		clientConn := client.NewClientConnection(lg, conn, nil, nil, hsHandler, 0, "", &backend.BCConfig{})
 		server.mu.clients[1] = clientConn
 		server.mu.Unlock()
 		return clientConn
@@ -107,18 +108,40 @@ func TestGracefulShutdown(t *testing.T) {
 	}
 }
 
-func TestRecoverPanic(t *testing.T) {
-	lg, text := logger.CreateLoggerForTest(t)
+func TestMultiAddr(t *testing.T) {
+	lg, _ := logger.CreateLoggerForTest(t)
 	certManager := cert.NewCertManager()
 	err := certManager.Init(&config.Config{}, lg, nil)
 	require.NoError(t, err)
 	server, err := NewSQLServer(lg, config.ProxyServer{
-		Addr: "0.0.0.0:6000",
+		Addr: "0.0.0.0:0,0.0.0.0:0",
 	}, certManager, &panicHsHandler{})
 	require.NoError(t, err)
 	server.Run(context.Background(), nil)
 
-	mdb, err := sql.Open("mysql", "root@tcp(localhost:6000)/test")
+	require.Len(t, server.listeners, 2)
+	for _, listener := range server.listeners {
+		conn, err := net.Dial("tcp", listener.Addr().String())
+		require.NoError(t, err)
+		require.NoError(t, conn.Close())
+	}
+
+	require.NoError(t, server.Close())
+	certManager.Close()
+}
+
+func TestRecoverPanic(t *testing.T) {
+	lg, text := logger.CreateLoggerForTest(t)
+	certManager := cert.NewCertManager()
+	err := certManager.Init(&config.Config{}, lg, nil)
+	require.NoError(t, err)
+	server, err := NewSQLServer(lg, config.ProxyServer{}, certManager, &panicHsHandler{})
+	require.NoError(t, err)
+	server.Run(context.Background(), nil)
+
+	_, port, err := net.SplitHostPort(server.listeners[0].Addr().String())
+	require.NoError(t, err)
+	mdb, err := sql.Open("mysql", fmt.Sprintf("root@tcp(localhost:%s)/test", port))
 	require.NoError(t, err)
 	// The first connection encounters panic.
 	require.ErrorContains(t, mdb.Ping(), "invalid connection")