[bug] fix ut TestAddressManager (#16002)

fix ut TestAddressManager

Approved by: @zhangxu19830126

pkg/proxy/handler.go

@@ -176,6 +176,8 @@ func (h *handler) handle(c goetty.IOSession) error {
 	h.logger.Info("build connection",
 		zap.String("client->proxy", fmt.Sprintf("%s -> %s", cc.RawConn().RemoteAddr(), cc.RawConn().LocalAddr())),
 		zap.String("proxy->server", fmt.Sprintf("%s -> %s", sc.RawConn().LocalAddr(), sc.RawConn().RemoteAddr())),
+		zap.Uint32("conn ID", cc.ConnID()),
+		zap.Uint64("session ID", c.ID()),
 	)
 
 	st := stopper.NewStopper("proxy-conn-handle", stopper.WithLogger(h.logger.RawLogger()))

pkg/util/address/address.go

@@ -77,7 +77,7 @@ func replaceHost(
 // by port base and the port slot.
 type AddressManager interface {
 	// Register registers a service by its name and port slot.
-	Register(portSlot int)
+	Register(portSlot int) int
 	// ListenAddress returns the service address of the service.
 	ListenAddress(slot int) string
 	// ServiceAddress returns the service address of the service.
@@ -110,16 +110,17 @@ func NewAddressManager(serviceAddress string, portBase int) AddressManager {
 }
 
 // Register implements the AddressManager interface.
-func (m *addressManager) Register(portSlot int) {
+func (m *addressManager) Register(portSlot int) int {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	if m.portBase == 0 {
-		return
+		return 0
 	}
 	if _, ok := m.mu.services[portSlot]; ok {
-		return
+		return m.mu.services[portSlot]
 	}
 	m.mu.services[portSlot] = m.portAdvanceLocked()
+	return m.mu.services[portSlot]
 }
 
 // ListenAddress implements the AddressManager interface.

pkg/util/address/address_test.go

@@ -15,6 +15,7 @@
 package address
 
 import (
+	"fmt"
 	"net"
 	"testing"
 	"time"
@@ -85,25 +86,38 @@ func TestAdjust(t *testing.T) {
 
 func TestAddressManager(t *testing.T) {
 	serviceAddr := "127.0.0.1"
-	portBase := 39000
+	portBase := 59320
 	m := NewAddressManager(serviceAddr, portBase)
 	assert.NotNil(t, m)
-	m.Register(2)
-	m.Register(1)
-	assert.Equal(t, "0.0.0.0:39000", m.ListenAddress(2))
-	assert.Equal(t, "127.0.0.1:39000", m.ServiceAddress(2))
-	assert.Equal(t, "0.0.0.0:39001", m.ListenAddress(1))
-	assert.Equal(t, "127.0.0.1:39001", m.ServiceAddress(1))
 
-	l, err := net.Listen("tcp4", "0.0.0.0:39002")
-	assert.NoError(t, err)
+	p1 := m.Register(2)
+	assert.Equal(t, fmt.Sprintf("0.0.0.0:%d", p1), m.ListenAddress(2))
+	assert.Equal(t, fmt.Sprintf("127.0.0.1:%d", p1), m.ServiceAddress(2))
+
+	p2 := m.Register(1)
+	assert.Equal(t, fmt.Sprintf("0.0.0.0:%d", p2), m.ListenAddress(1))
+	assert.Equal(t, fmt.Sprintf("127.0.0.1:%d", p2), m.ServiceAddress(1))
+
+	p3 := p2 + 1
+
+	var err error
+	var l net.Listener
+	for {
+		l, err = net.Listen("tcp4", fmt.Sprintf("0.0.0.0:%d", p3))
+		if err == nil {
+			break
+		}
+		p3++
+	}
 	defer func() {
-		err = l.Close()
-		assert.NoError(t, err)
+		if l != nil {
+			err = l.Close()
+			assert.NoError(t, err)
+		}
 	}()
-	m.Register(3)
-	assert.Equal(t, "0.0.0.0:39003", m.ListenAddress(3))
-	assert.Equal(t, "127.0.0.1:39003", m.ServiceAddress(3))
+	p4 := m.Register(3)
+	assert.Equal(t, fmt.Sprintf("0.0.0.0:%d", p4), m.ListenAddress(3))
+	assert.Equal(t, fmt.Sprintf("127.0.0.1:%d", p4), m.ServiceAddress(3))
 }
 
 func TestRemoteAddressAvail(t *testing.T) {