Repository: cloudwego/hertz
Pull request #1299 (2025-03-13 16:43:00)

fix(network/standard): Shutdown checks active conns
#### What type of PR is this?
fix

<!--
Add one of the following kinds:

build: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)
ci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)
docs: Documentation only changes
feat: A new feature
optimize: A new optimization
fix: A bug fix
perf: A code change that improves performance
refactor: A code change that neither fixes a bug nor adds a feature
style: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)
test: Adding missing tests or correcting existing tests
chore: Changes to the build process or auxiliary tools and libraries such as documentation generation
-->

#### Check the PR title.
<!--
The description of the title will be attached in Release Notes, 
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [ ] This PR title match the format: \<type\>(optional scope): \<description\>
- [ ] The description of this PR title is user-oriented and clear enough for others to understand.
- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)


#### (Optional) Translate the PR title into Chinese.


#### (Optional) More detailed description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).
-->
en:
zh(optional):

#### (Optional) Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->

#### (Optional) The PR that updates user documentation:
<!--
If the current PR requires user awareness at the usage level, please submit a PR to update user docs. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)
-->


Patch:
diff --git a/pkg/network/standard/transport.go b/pkg/network/standard/transport.go
index b5f169e..aa4bce0 100644
--- a/pkg/network/standard/transport.go
+++ b/pkg/network/standard/transport.go
@@ -19,8 +19,10 @@ package standard
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/cloudwego/hertz/pkg/common/config"
@@ -48,6 +50,15 @@ type transport struct {
 	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
+}
+
+func (t *transport) listener() net.Listener {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+	return t.ln
 }
 
 func (t *transport) serve() (err error) {
@@ -58,24 +69,28 @@ func (t *transport) serve() (err error) {
 	} else {
 		t.ln, err = net.Listen(t.network, t.addr)
 	}
+	// fix concurrency issue
+	// normally listener must not be changed during serve()
+	ln := t.ln
 	t.lock.Unlock()
 	if err != nil {
 		return err
 	}
-	hlog.SystemLogger().Infof("HTTP server listening on address=%s", t.ln.Addr().String())
+	hlog.SystemLogger().Infof("HTTP server listening on address=%s", ln.Addr().String())
 	for {
 		ctx := context.Background()
-		conn, err := t.ln.Accept()
-		var c network.Conn
+		conn, err := ln.Accept()
 		if err != nil {
 			hlog.SystemLogger().Errorf("Error=%s", err.Error())
 			return err
 		}
+		t.updateActive(1)
 
 		if t.OnAccept != nil {
 			ctx = t.OnAccept(conn)
 		}
 
+		var c network.Conn
 		if t.tls != nil {
 			c = newTLSConn(tls.Server(conn, t.tls), t.readBufferSize)
 		} else {
@@ -85,10 +100,17 @@ func (t *transport) serve() (err error) {
 		if t.OnConnect != nil {
 			ctx = t.OnConnect(ctx, c)
 		}
-		go t.handler(ctx, c)
+		go func(ctx context.Context, conn network.Conn) {
+			t.handler(ctx, conn)
+			t.updateActive(-1)
+		}(ctx, c)
 	}
 }
 
+func (t *transport) updateActive(delta int32) int32 {
+	return atomic.AddInt32(&t.active, delta)
+}
+
 func (t *transport) ListenAndServe(onData network.OnData) (err error) {
 	t.handler = onData
 	return t.serve()
@@ -100,6 +122,13 @@ func (t *transport) Close() error {
 	return t.Shutdown(ctx)
 }
 
+var (
+	shutdownTimeout = 30 * time.Second
+	shutdownTicker  = 10 * time.Millisecond
+
+	errShutdownTimeout = errors.New("shutdown timeout")
+)
+
 func (t *transport) Shutdown(ctx context.Context) error {
 	defer func() {
 		network.UnlinkUdsFile(t.network, t.addr) //nolint:errcheck
@@ -109,8 +138,33 @@ func (t *transport) Shutdown(ctx context.Context) error {
 		_ = t.ln.Close()
 	}
 	t.lock.Unlock()
-	<-ctx.Done()
-	return nil
+
+	tk := time.NewTicker(shutdownTicker)
+	defer tk.Stop()
+
+	// make sure t.active is updated correctly under concurrency
+	<-tk.C
+
+	// luckily the server is idle, no more active connections
+	if t.updateActive(0) <= 0 {
+		return nil
+	}
+
+	// check periodically to see if all connections closed
+	t0 := time.Now()
+	for {
+		select {
+		case now := <-tk.C:
+			if t.updateActive(0) <= 0 {
+				return nil
+			}
+			if now.Sub(t0) > shutdownTimeout {
+				return errShutdownTimeout
+			}
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
 }
 
 // For transporter switch
diff --git a/pkg/protocol/http1/client.go b/pkg/protocol/http1/client.go
index 5bac3e9..eb57d0d 100644
--- a/pkg/protocol/http1/client.go
+++ b/pkg/protocol/http1/client.go
@@ -363,6 +363,10 @@ func (c *HostClient) DoRedirects(ctx context.Context, req *protocol.Request, res
 // It is recommended obtaining req and resp via AcquireRequest
 // and AcquireResponse in performance-critical code.
 func (c *HostClient) Do(ctx context.Context, req *protocol.Request, resp *protocol.Response) error {
+	if ctx == nil {
+		panic("ctx is nil")
+	}
+
 	var (
 		err                error
 		canIdempotentRetry bool


Test patch:
diff --git a/pkg/network/standard/transport_test.go b/pkg/network/standard/transport_test.go
new file mode 100644
index 0000000..0f11fe0
--- /dev/null
+++ b/pkg/network/standard/transport_test.go
@@ -0,0 +1,108 @@
+/*
+ * Copyright 2025 CloudWeGo Authors
+ *
+ * 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 standard
+
+import (
+	"context"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/cloudwego/hertz/pkg/common/config"
+	"github.com/cloudwego/hertz/pkg/network"
+)
+
+func assertWriteRead(t *testing.T, c io.ReadWriter, w string, r string) {
+	t.Helper()
+	_, err := c.Write([]byte(w))
+	if err != nil {
+		t.Fatal("write err", err)
+	}
+	b := make([]byte, len(r))
+	_, err = io.ReadFull(c, b)
+	if err != nil {
+		t.Fatal("read err", err)
+	}
+	if s := string(b); s != r {
+		t.Fatal("read", r)
+	}
+}
+
+func TestTransporter(t *testing.T) {
+	handlerExit := make(chan struct{})
+	req := "hello"
+	resp := "world"
+	trans := NewTransporter(&config.Options{Network: "tcp", Addr: "127.0.0.1:0"}).(*transport)
+	go trans.ListenAndServe(func(ctx context.Context, conn interface{}) error {
+		c := conn.(network.Conn)
+		defer c.Close()
+		assertWriteRead(t, c, resp, req)
+		<-handlerExit
+		return nil
+	})
+	for trans.listener() == nil { // wait server up
+		time.Sleep(5 * time.Millisecond)
+	}
+
+	// dial and test
+	c, err := net.Dial("tcp", trans.listener().Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	assertWriteRead(t, c, req, resp)
+
+	checkActiveConn := func(n int) {
+		if v := trans.updateActive(0); v != int32(n) {
+			t.Helper()
+			t.Fatal("trans active conn", v)
+		}
+	}
+	checkActiveConn(1)
+
+	// make sure shutdownTimeout will be reset after this test
+	defer func(old time.Duration) { shutdownTimeout = old }(shutdownTimeout)
+	defer func(old time.Duration) { shutdownTicker = old }(shutdownTicker)
+
+	shutdownTicker = time.Millisecond // shorter for saving test time
+
+	// case: wait util shutdownTimeout
+	shutdownTimeout = time.Millisecond
+	if err := trans.Shutdown(context.Background()); err != errShutdownTimeout {
+		t.Fatal(err)
+	}
+
+	// case: ctx done
+	shutdownTimeout = time.Second // long enough
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = trans.Shutdown(ctx)
+	if err != ctx.Err() {
+		t.Fatal(err)
+	}
+
+	// case: even after listener is closed, handler may still active
+	checkActiveConn(1)
+
+	// case: Shutdown blocks at ticker and check periodically.
+	shutdownTimeout = time.Second // long enough
+	go trans.Shutdown(context.Background())
+	time.Sleep(30 * time.Millisecond) // make sure Shutdown blocks at for loop
+	close(handlerExit)                // signal handler to return
+	time.Sleep(10 * time.Millisecond) // wait handler returns, and active conn to be updated.
+	checkActiveConn(0)
+}
diff --git a/pkg/protocol/http1/client_test.go b/pkg/protocol/http1/client_test.go
index 2ecb764..353cf61 100644
--- a/pkg/protocol/http1/client_test.go
+++ b/pkg/protocol/http1/client_test.go
@@ -656,3 +656,13 @@ func TestDialTimeout(t *testing.T) {
 
 	c.Do(context.Background(), req, resp)
 }
+
+func TestContextNil(t *testing.T) {
+	defer func() {
+		v := recover()
+		assert.NotNil(t, v)
+		assert.True(t, fmt.Sprint(v) == "ctx is nil")
+	}()
+	c := &HostClient{}
+	c.Do(nil, nil, nil) //nolint:staticcheck // SA1012: do not pass a nil Context
+}
diff --git a/pkg/route/engine_test.go b/pkg/route/engine_test.go
index a5da5dc..2c8c97f 100644
--- a/pkg/route/engine_test.go
+++ b/pkg/route/engine_test.go
@@ -49,6 +49,7 @@ import (
 	"io/ioutil"
 	"net"
 	"net/http"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -871,48 +872,91 @@ func (e mockDeregsitryErr) Deregister(*registry.Info) error {
 	return errTestDeregsitry
 }
 
+type mockStandardTransporter struct {
+	network.Transporter
+}
+
+func (m *mockStandardTransporter) Shutdown(ctx context.Context) error {
+	// FIXME: standard.Transporter mindlessly blocks on ctx.Done()
+	// This change help tests run faster
+	newctx, cancel := context.WithCancel(ctx)
+	cancel()
+	return m.Transporter.Shutdown(newctx)
+}
+
+func newMockStandardTransporter(opt *config.Options) network.Transporter {
+	return &mockStandardTransporter{standard.NewTransporter(opt)}
+}
+
 func TestEngineShutdown(t *testing.T) {
-	defaultTransporter = standard.NewTransporter
 	mockCtxCallback := func(ctx context.Context) {}
-	// Test case 1: serve not running error
+
 	opt := config.NewOptions(nil)
-	opt.Addr = "127.0.0.1:10027"
-	engine := NewEngine(opt)
-	ctx1, cancel1 := context.WithTimeout(context.Background(), time.Second)
-	defer cancel1()
-	err := engine.Shutdown(ctx1)
-	assert.DeepEqual(t, errStatusNotRunning, err)
-
-	// Test case 2: serve successfully running and shutdown
+	opt.Addr = "127.0.0.1:0"
+	opt.TransporterNewer = newMockStandardTransporter
+
+	var wg sync.WaitGroup
+	var engine *Engine
+
+	runEngine := func() {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			engine.Run()
+		}()
+		// wait for engine to start
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	shutdownEngine := func(ctx context.Context, expectErr error, expectStatus uint32) {
+		t.Helper()
+		err := engine.Shutdown(ctx)
+		if expectErr == nil {
+			assert.Nil(t, err)
+		} else {
+			assert.DeepEqual(t, expectErr, err)
+		}
+		if expectStatus != 0 {
+			if expectStatus == statusShutdown {
+				assert.DeepEqual(t, expectStatus, atomic.LoadUint32(&engine.status))
+				// make sure engine.Run() returns
+				// in case registry fails, it blocks
+				engine.transport.Shutdown(ctx)
+				expectStatus = statusClosed
+			}
+			wg.Wait() // wait engine.Run() returns
+		}
+		assert.DeepEqual(t, expectStatus, atomic.LoadUint32(&engine.status))
+	}
+
+	// case: serve not running error
 	engine = NewEngine(opt)
-	engine.OnShutdown = []CtxCallback{mockCtxCallback}
-	go func() {
-		engine.Run()
-	}()
-	// wait for engine to start
-	time.Sleep(1 * time.Second)
+	shutdownEngine(context.Background(), errStatusNotRunning, 0)
 
-	ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
-	defer cancel2()
-	err = engine.Shutdown(ctx2)
-	assert.Nil(t, err)
-	assert.DeepEqual(t, statusClosed, atomic.LoadUint32(&engine.status))
+	// case: serve successfully running and shutdown
+	engine = NewEngine(opt)
+	engine.OnShutdown = []CtxCallback{mockCtxCallback}
+	runEngine()
+	shutdownEngine(context.Background(), nil, statusClosed)
 
-	// Test case 3: serve successfully running and shutdown with deregistry error
+	// case: serve successfully running and shutdown with deregistry error
 	engine = NewEngine(opt)
 	engine.OnShutdown = []CtxCallback{mockCtxCallback}
 	engine.options.Registry = &mockDeregsitryErr{}
-	go func() {
-		engine.Run()
-	}()
-	// wait for engine to start
-	time.Sleep(1 * time.Second)
-
-	ctx3, cancel3 := context.WithTimeout(context.Background(), time.Second)
-	defer cancel3()
-	err = engine.Shutdown(ctx3)
-	assert.DeepEqual(t, errTestDeregsitry, err)
-	assert.DeepEqual(t, statusShutdown, atomic.LoadUint32(&engine.status))
+	runEngine()
+	shutdownEngine(context.Background(), errTestDeregsitry, statusShutdown)
+	engine.options.Registry = nil
+
+	// case: ctx cancelled when Shutdown
+	engine = NewEngine(opt)
+
+	// make sure callback is in progress but ctx cancelled
+	engine.OnShutdown = []CtxCallback{func(ctx context.Context) { time.Sleep(50 * time.Millisecond) }}
+
+	runEngine()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	shutdownEngine(ctx, nil, statusClosed)
 }
 
 type mockStreamer struct{}