fix(engine): IsRunning returns true when it's ready (#1301)

* added Listener() method for network transport
* check Listener() in IsRunning()

Author: xiaost

pkg/app/client/client_test.go

@@ -111,25 +111,10 @@ func nextUnixSock() string {
 func waitEngineRunning(e *route.Engine) {
 	for i := 0; i < 100; i++ {
 		if e.IsRunning() {
-			break
+			return
 		}
 		time.Sleep(10 * time.Millisecond)
 	}
-	opts := e.GetOptions()
-	network, addr := opts.Network, opts.Addr
-	if network == "" {
-		network = "tcp"
-	}
-	for i := 0; i < 100; i++ {
-		conn, err := net.Dial(network, addr)
-		if err != nil {
-			time.Sleep(10 * time.Millisecond)
-			continue
-		}
-		conn.Close()
-		return
-	}
-
 	panic("not running")
 }
 

pkg/app/server/hertz_test.go

@@ -53,30 +53,15 @@ import (
 
 type routeEngine interface {
 	IsRunning() bool
-	GetOptions() *config.Options
 }
 
 func waitEngineRunning(e routeEngine) {
 	for i := 0; i < 100; i++ {
 		if e.IsRunning() {
-			break
+			return
 		}
 		time.Sleep(10 * time.Millisecond)
 	}
-	opts := e.GetOptions()
-	network, addr := opts.Network, opts.Addr
-	if network == "" {
-		network = "tcp"
-	}
-	for i := 0; i < 100; i++ {
-		conn, err := net.Dial(network, addr)
-		if err != nil {
-			time.Sleep(10 * time.Millisecond)
-			continue
-		}
-		conn.Close()
-		return
-	}
 	panic("not running")
 }
 
@@ -840,15 +825,9 @@ type CloseWithoutResetBuffer interface {
 }
 
 func TestOnprepare(t *testing.T) {
-	n := int32(0)
 	h1 := New(
 		WithHostPorts("localhost:9333"),
 		WithOnConnect(func(ctx context.Context, conn network.Conn) context.Context {
-			if atomic.AddInt32(&n, 1) == 1 {
-				// the 1st connection is from waitEngineRunning
-				conn.Close()
-				return ctx
-			}
 			b, err := conn.Peek(3)
 			assert.Nil(t, err)
 			assert.DeepEqual(t, string(b), "GET")

pkg/network/netpoll/dial_test.go

@@ -27,12 +27,17 @@ import (
 	"github.com/cloudwego/hertz/pkg/common/config"
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
 	"github.com/cloudwego/hertz/pkg/common/test/mock"
+	"github.com/cloudwego/hertz/pkg/network"
 )
 
+func getListenerAddr(trans network.Transporter) string {
+	return trans.(*transporter).Listener().Addr().String()
+}
+
 func TestDial(t *testing.T) {
 	t.Run("NetpollDial", func(t *testing.T) {
 		const nw = "tcp"
-		const addr = "localhost:10100"
+		var addr = "127.0.0.1:0"
 		transporter := NewTransporter(&config.Options{
 			Addr:    addr,
 			Network: nw,
@@ -48,6 +53,7 @@ func TestDial(t *testing.T) {
 		_, err := dial.DialConnection("tcp", "localhost:10101", time.Second, nil) // wrong addr
 		assert.NotNil(t, err)
 
+		addr = getListenerAddr(transporter)
 		nwConn, err := dial.DialConnection(nw, addr, time.Second, nil)
 		assert.Nil(t, err)
 		defer nwConn.Close()

pkg/network/netpoll/transport.go

@@ -47,18 +47,19 @@ func cancelContext(ctx context.Context) context.Context {
 }
 
 type transporter struct {
-	sync.RWMutex
 	senseClientDisconnection bool
 	network                  string
 	addr                     string
 	keepAliveTimeout         time.Duration
 	readTimeout              time.Duration
 	writeTimeout             time.Duration
-	listener                 net.Listener
-	eventLoop                netpoll.EventLoop
 	listenConfig             *net.ListenConfig
 	OnAccept                 func(conn net.Conn) context.Context
 	OnConnect                func(ctx context.Context, conn network.Conn) context.Context
+
+	mu sync.RWMutex
+	ln net.Listener
+	el netpoll.EventLoop
 }
 
 // For transporter switch
@@ -70,23 +71,31 @@ func NewTransporter(options *config.Options) network.Transporter {
 		keepAliveTimeout:         options.KeepAliveTimeout,
 		readTimeout:              options.ReadTimeout,
 		writeTimeout:             options.WriteTimeout,
-		listener:                 nil,
-		eventLoop:                nil,
 		listenConfig:             options.ListenConfig,
 		OnAccept:                 options.OnAccept,
 		OnConnect:                options.OnConnect,
 	}
 }
 
+func (t *transporter) Listener() net.Listener {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.ln
+}
+
 // ListenAndServe binds listen address and keep serving, until an error occurs
 // or the transport shutdowns
 func (t *transporter) ListenAndServe(onReq network.OnData) (err error) {
 	network.UnlinkUdsFile(t.network, t.addr) //nolint:errcheck
+
+	t.mu.Lock()
 	if t.listenConfig != nil {
-		t.listener, err = t.listenConfig.Listen(context.Background(), t.network, t.addr)
+		t.ln, err = t.listenConfig.Listen(context.Background(), t.network, t.addr)
 	} else {
-		t.listener, err = net.Listen(t.network, t.addr)
+		t.ln, err = net.Listen(t.network, t.addr)
 	}
+	ln := t.ln
+	t.mu.Unlock()
 
 	if err != nil {
 		panic("create netpoll listener fail: " + err.Error())
@@ -127,20 +136,19 @@ func (t *transporter) ListenAndServe(onReq network.OnData) (err error) {
 	}
 
 	// Create EventLoop
-	t.Lock()
-	t.eventLoop, err = netpoll.NewEventLoop(func(ctx context.Context, connection netpoll.Connection) error {
+	t.mu.Lock()
+	t.el, err = netpoll.NewEventLoop(func(ctx context.Context, connection netpoll.Connection) error {
 		return onReq(ctx, newConn(connection))
 	}, opts...)
-	t.Unlock()
+	eventLoop := t.el
+	t.mu.Unlock()
 	if err != nil {
 		panic("create netpoll event-loop fail")
 	}
 
 	// Start Server
-	hlog.SystemLogger().Infof("HTTP server listening on address=%s", t.listener.Addr().String())
-	t.RLock()
-	err = t.eventLoop.Serve(t.listener)
-	t.RUnlock()
+	hlog.SystemLogger().Infof("HTTP server listening on address=%s", ln.Addr().String())
+	err = eventLoop.Serve(ln)
 	if err != nil {
 		panic("netpoll server exit")
 	}
@@ -160,11 +168,11 @@ func (t *transporter) Close() error {
 func (t *transporter) Shutdown(ctx context.Context) error {
 	defer func() {
 		network.UnlinkUdsFile(t.network, t.addr) //nolint:errcheck
-		t.RUnlock()
+		t.mu.RUnlock()
 	}()
-	t.RLock()
-	if t.eventLoop == nil {
+	t.mu.RLock()
+	if t.el == nil {
 		return nil
 	}
-	return t.eventLoop.Shutdown(ctx)
+	return t.el.Shutdown(ctx)
 }

pkg/network/netpoll/transport_test.go

@@ -34,8 +34,8 @@ import (
 
 func TestTransport(t *testing.T) {
 	const nw = "tcp"
-	const addr = "localhost:10103"
 	t.Run("TestDefault", func(t *testing.T) {
+		var addr = "127.0.0.1:0"
 		var onConnFlag, onAcceptFlag, onDataFlag int32
 		transporter := NewTransporter(&config.Options{
 			Addr:    addr,
@@ -57,6 +57,8 @@ func TestTransport(t *testing.T) {
 		defer transporter.Close()
 		time.Sleep(100 * time.Millisecond)
 
+		addr = getListenerAddr(transporter)
+
 		dial := NewDialer()
 		conn, err := dial.DialConnection(nw, addr, time.Second, nil)
 		assert.Nil(t, err)
@@ -70,6 +72,7 @@ func TestTransport(t *testing.T) {
 	})
 
 	t.Run("TestSenseClientDisconnection", func(t *testing.T) {
+		var addr = "127.0.0.1:0"
 		var onReqFlag int32
 		transporter := NewTransporter(&config.Options{
 			Addr:                     addr,
@@ -86,6 +89,8 @@ func TestTransport(t *testing.T) {
 		defer transporter.Close()
 		time.Sleep(100 * time.Millisecond)
 
+		addr = getListenerAddr(transporter)
+
 		dial := NewDialer()
 		conn, err := dial.DialConnection(nw, addr, time.Second, nil)
 		assert.Nil(t, err)
@@ -106,7 +111,7 @@ func TestTransport(t *testing.T) {
 			})
 		}}
 		transporter := NewTransporter(&config.Options{
-			Addr:         addr,
+			Addr:         "127.0.0.1:0",
 			Network:      nw,
 			ListenConfig: listenCfg,
 		})

pkg/network/standard/dial_test.go

@@ -31,11 +31,16 @@ import (
 
 	"github.com/cloudwego/hertz/pkg/common/config"
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
+	"github.com/cloudwego/hertz/pkg/network"
 )
 
+func getListenerAddr(trans network.Transporter) string {
+	return trans.(*transport).Listener().Addr().String()
+}
+
 func TestDial(t *testing.T) {
 	const nw = "tcp"
-	const addr = "localhost:10104"
+	var addr = "127.0.0.1:0"
 	transporter := NewTransporter(&config.Options{
 		Addr:    addr,
 		Network: nw,
@@ -47,6 +52,8 @@ func TestDial(t *testing.T) {
 	defer transporter.Close()
 	time.Sleep(time.Millisecond * 100)
 
+	addr = getListenerAddr(transporter)
+
 	dial := NewDialer()
 	_, err := dial.DialConnection(nw, addr, time.Second, nil)
 	assert.Nil(t, err)
@@ -58,9 +65,9 @@ func TestDial(t *testing.T) {
 
 func TestDialTLS(t *testing.T) {
 	const nw = "tcp"
-	const addr = "localhost:10105"
+	var addr = "127.0.0.1:0"
 	data := []byte("abcdefg")
-	listened := make(chan struct{})
+	listened := make(chan net.Listener)
 	go func() {
 		mockTLSServe(nw, addr, func(conn net.Conn) {
 			defer conn.Close()
@@ -70,7 +77,8 @@ func TestDialTLS(t *testing.T) {
 	}()
 
 	select {
-	case <-listened:
+	case ln := <-listened:
+		addr = ln.Addr().String()
 	case <-time.After(time.Second * 5):
 		t.Fatalf("timeout")
 	}
@@ -99,7 +107,7 @@ func TestDialTLS(t *testing.T) {
 	assert.DeepEqual(t, string(data), string(buf))
 }
 
-func mockTLSServe(nw, addr string, handle func(conn net.Conn), listened chan struct{}) (err error) {
+func mockTLSServe(nw, addr string, handle func(conn net.Conn), listened chan net.Listener) (err error) {
 	certData, keyData, err := generateTestCertificate("")
 	if err != nil {
 		return
@@ -119,7 +127,7 @@ func mockTLSServe(nw, addr string, handle func(conn net.Conn), listened chan str
 	}
 	defer ln.Close()
 
-	listened <- struct{}{}
+	listened <- ln
 	for {
 		conn, err := ln.Accept()
 		if err != nil {

pkg/network/standard/transport.go

@@ -44,26 +44,27 @@ type transport struct {
 	keepAliveTimeout time.Duration
 	readTimeout      time.Duration
 	handler          network.OnData
-	ln               net.Listener
 	tls              *tls.Config
 	listenConfig     *net.ListenConfig
-	lock             sync.Mutex
 	OnAccept         func(conn net.Conn) context.Context
 	OnConnect        func(ctx context.Context, conn network.Conn) context.Context
 
 	// active connections. it +1 after accept and -1 after handler returns
 	active int32
+
+	mu sync.RWMutex
+	ln net.Listener
 }
 
-func (t *transport) listener() net.Listener {
-	t.lock.Lock()
-	defer t.lock.Unlock()
+func (t *transport) Listener() net.Listener {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
 	return t.ln
 }
 
 func (t *transport) serve() (err error) {
 	network.UnlinkUdsFile(t.network, t.addr) //nolint:errcheck
-	t.lock.Lock()
+	t.mu.Lock()
 	if t.listenConfig != nil {
 		t.ln, err = t.listenConfig.Listen(context.Background(), t.network, t.addr)
 	} else {
@@ -72,7 +73,7 @@ func (t *transport) serve() (err error) {
 	// fix concurrency issue
 	// normally listener must not be changed during serve()
 	ln := t.ln
-	t.lock.Unlock()
+	t.mu.Unlock()
 	if err != nil {
 		return err
 	}
@@ -133,11 +134,9 @@ func (t *transport) Shutdown(ctx context.Context) error {
 	defer func() {
 		network.UnlinkUdsFile(t.network, t.addr) //nolint:errcheck
 	}()
-	t.lock.Lock()
-	if t.ln != nil {
-		_ = t.ln.Close()
+	if ln := t.Listener(); ln != nil {
+		_ = ln.Close()
 	}
-	t.lock.Unlock()
 
 	tk := time.NewTicker(shutdownTicker)
 	defer tk.Stop()

pkg/network/standard/transport_test.go

@@ -55,12 +55,12 @@ func TestTransporter(t *testing.T) {
 		<-handlerExit
 		return nil
 	})
-	for trans.listener() == nil { // wait server up
+	for trans.Listener() == nil { // wait server up
 		time.Sleep(5 * time.Millisecond)
 	}
 
 	// dial and test
-	c, err := net.Dial("tcp", trans.listener().Addr().String())
+	c, err := net.Dial("tcp", trans.Listener().Addr().String())
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/route/engine.go

@@ -47,6 +47,7 @@ import (
 	"fmt"
 	"html/template"
 	"io"
+	"net"
 	"path/filepath"
 	"reflect"
 	"runtime"
@@ -258,7 +259,18 @@ func (engine *Engine) IsStreamRequestBody() bool {
 }
 
 func (engine *Engine) IsRunning() bool {
-	return atomic.LoadUint32(&engine.status) == statusRunning
+	if atomic.LoadUint32(&engine.status) != statusRunning {
+		return false
+	}
+	// double check listener
+	type ListenerIface interface {
+		Listener() net.Listener
+	}
+	v, ok := engine.transport.(ListenerIface)
+	if ok {
+		return v.Listener() != nil
+	}
+	return true // default behavior if no ListenerIface
 }
 
 func (engine *Engine) HijackConnHandle(c network.Conn, h app.HijackHandler) {