fix: Improve server stop handling with graceful shutdowns (go-kratos#3525)

lftk · web-flow · commit 20970020a5ef · 2025-01-27T17:52:00.000+08:00
* Improve server stop handling with graceful shutdowns

- Remove default stop timeout
- Add context handling for server stop
- Implement graceful stop for gRPC server
- Enhance HTTP server shutdown logic
- Use shared context for server operations

* Remove unnecessary error logging in gRPC server test
diff --git a/app.go b/app.go
@@ -41,7 +41,6 @@ func New(opts ...Option) *App {
 		ctx:              context.Background(),
 		sigs:             []os.Signal{syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT},
 		registrarTimeout: 10 * time.Second,
-		stopTimeout:      10 * time.Second,
 	}
 	if id, err := uuid.NewUUID(); err == nil {
 		o.id = id.String()
@@ -98,18 +97,23 @@ func (a *App) Run() error {
 			return err
 		}
 	}
+	octx := NewContext(a.opts.ctx, a)
 	for _, srv := range a.opts.servers {
 		server := srv
 		eg.Go(func() error {
 			<-ctx.Done() // wait for stop signal
-			stopCtx, cancel := context.WithTimeout(NewContext(a.opts.ctx, a), a.opts.stopTimeout)
-			defer cancel()
+			stopCtx := octx
+			if a.opts.stopTimeout > 0 {
+				var cancel context.CancelFunc
+				stopCtx, cancel = context.WithTimeout(stopCtx, a.opts.stopTimeout)
+				defer cancel()
+			}
 			return server.Stop(stopCtx)
 		})
 		wg.Add(1)
 		eg.Go(func() error {
 			wg.Done() // here is to ensure server start has begun running before register, so defer is not needed
-			return server.Start(NewContext(a.opts.ctx, a))
+			return server.Start(octx)
 		})
 	}
 	wg.Wait()
diff --git a/transport/grpc/server.go b/transport/grpc/server.go
@@ -223,13 +223,25 @@ func (s *Server) Start(ctx context.Context) error {
 }
 
 // Stop stop the gRPC server.
-func (s *Server) Stop(_ context.Context) error {
+func (s *Server) Stop(ctx context.Context) error {
 	if s.adminClean != nil {
 		s.adminClean()
 	}
 	s.health.Shutdown()
-	s.GracefulStop()
-	log.Info("[gRPC] server stopping")
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		log.Info("[gRPC] server stopping")
+		s.Server.GracefulStop()
+	}()
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		log.Warn("[gRPC] server couldn't stop gracefully in time, doing force stop")
+		s.Server.Stop()
+	}
 	return nil
 }
 
diff --git a/transport/grpc/server_test.go b/transport/grpc/server_test.go
@@ -1,13 +1,15 @@
 package grpc
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
 	"net"
 	"net/url"
 	"reflect"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -17,6 +19,7 @@ import (
 	"github.com/go-kratos/kratos/v2/errors"
 	"github.com/go-kratos/kratos/v2/internal/matcher"
 	pb "github.com/go-kratos/kratos/v2/internal/testdata/helloworld"
+	"github.com/go-kratos/kratos/v2/log"
 	"github.com/go-kratos/kratos/v2/middleware"
 	"github.com/go-kratos/kratos/v2/transport"
 )
@@ -371,3 +374,127 @@ func TestListener(t *testing.T) {
 		t.Errorf("expect not empty")
 	}
 }
+
+func TestStop(t *testing.T) {
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+
+	tests := []struct {
+		name          string
+		ctx           context.Context
+		cancel        context.CancelFunc
+		wantForceStop bool
+	}{
+		{
+			name:          "normal",
+			ctx:           context.Background(),
+			cancel:        func() {},
+			wantForceStop: false,
+		},
+		{
+			name:          "timeout",
+			ctx:           timeoutCtx,
+			cancel:        cancel,
+			wantForceStop: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l, err := net.Listen("tcp", ":0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer l.Close()
+
+			old := log.GetLogger()
+			defer log.SetLogger(old)
+
+			// Create a logger to capture logs
+			var logs safeBytesBuffer
+			log.SetLogger(log.NewStdLogger(&logs))
+
+			s := NewServer(Listener(l))
+			pb.RegisterGreeterServer(s, &server{})
+
+			go func() {
+				err := s.Start(context.Background()) //nolint
+				if err != nil {
+					log.Fatal(err)
+				}
+			}()
+
+			time.Sleep(100 * time.Millisecond)
+
+			conn, err := DialInsecure(
+				context.Background(),
+				WithEndpoint(l.Addr().String()),
+				WithOptions(grpc.WithBlock()),
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer conn.Close()
+
+			go func() {
+				client := pb.NewGreeterClient(conn)
+				if tt.wantForceStop {
+					// Simulate a long-running request
+					s, err := client.SayHelloStream(context.Background()) //nolint
+					if err != nil {
+						log.Fatal(err)
+					}
+					// Keep the stream open
+					for {
+						// Intentionally do not send messages, only receive messages
+						_, err := s.Recv()
+						if err != nil {
+							break
+						}
+					}
+				} else {
+					_, err := client.SayHello(context.Background(), &pb.HelloRequest{Name: "test"}) //nolint
+					if err != nil {
+						log.Error(err)
+					}
+				}
+			}()
+
+			time.Sleep(100 * time.Millisecond)
+
+			err = s.Stop(tt.ctx)
+			if err != nil {
+				t.Errorf("Expected no error, got %v", err)
+				return
+			}
+
+			// Check if the stop was forced or graceful
+			if tt.wantForceStop {
+				if !strings.Contains(logs.String(), "force stop") {
+					t.Errorf("Expected force stop\n%s", logs.String())
+				}
+			} else {
+				if strings.Contains(logs.String(), "force stop") {
+					t.Errorf("Expected graceful stop\n%s", logs.String())
+				}
+			}
+		})
+	}
+}
+
+type safeBytesBuffer struct {
+	mu  sync.Mutex
+	buf bytes.Buffer
+}
+
+func (b *safeBytesBuffer) Write(p []byte) (n int, err error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.Write(p)
+}
+
+func (b *safeBytesBuffer) String() string {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.String()
+}
diff --git a/transport/http/resolver_test.go b/transport/http/resolver_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"reflect"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -86,6 +87,8 @@ type mockWatch struct {
 
 	nextErr bool
 	stopErr bool
+
+	lock sync.Mutex
 }
 
 func (m *mockWatch) Next() ([]*registry.ServiceInstance, error) {
@@ -94,6 +97,8 @@ func (m *mockWatch) Next() ([]*registry.ServiceInstance, error) {
 		return nil, m.ctx.Err()
 	default:
 	}
+	m.lock.Lock()
+	defer m.lock.Unlock()
 	if m.nextErr {
 		return nil, errors.New("mock test error")
 	}
@@ -115,6 +120,8 @@ func (m *mockWatch) Next() ([]*registry.ServiceInstance, error) {
 }
 
 func (m *mockWatch) Stop() error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
 	if m.stopErr {
 		return errors.New("mock test error")
 	}
@@ -130,46 +137,67 @@ func TestResolver(t *testing.T) {
 		return
 	}
 
+	cancelCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	// 异步 无需报错
-	_, err = newResolver(context.Background(), &mockDiscoveries{true, false, false}, ta, &mockRebalancer{}, false, false, 25)
+	r, err := newResolver(cancelCtx, &mockDiscoveries{true, false, false}, ta, &mockRebalancer{}, false, false, 25)
 	if err != nil {
 		t.Errorf("expect %v, got %v", nil, err)
 	}
+	if r != nil {
+		_ = r.Close()
+	}
 
 	// 同步 一切正常运行
-	_, err = newResolver(context.Background(), &mockDiscoveries{false, false, false}, ta, &mockRebalancer{}, true, true, 25)
+	r, err = newResolver(cancelCtx, &mockDiscoveries{false, false, false}, ta, &mockRebalancer{}, true, true, 25)
 	if err != nil {
 		t.Errorf("expect %v, got %v", nil, err)
 	}
+	if r != nil {
+		_ = r.Close()
+	}
 
 	// 同步 但是 next 出错 以及 stop 出错
-	_, err = newResolver(context.Background(), &mockDiscoveries{false, true, true}, ta, &mockRebalancer{}, true, true, 25)
+	r, err = newResolver(cancelCtx, &mockDiscoveries{false, true, true}, ta, &mockRebalancer{}, true, true, 25)
 	if err == nil {
 		t.Errorf("expect err, got nil")
 	}
+	if r != nil {
+		_ = r.Close()
+	}
 
 	// 同步 service name watch 失败
-	_, err = newResolver(context.Background(), &mockDiscoveries{false, true, true}, &Target{
+	r, err = newResolver(cancelCtx, &mockDiscoveries{false, true, true}, &Target{
 		Scheme:   "discovery",
 		Endpoint: errServiceName,
 	}, &mockRebalancer{}, true, true, 25)
 	if err == nil {
 		t.Errorf("expect err, got nil")
 	}
+	if r != nil {
+		_ = r.Close()
+	}
 
-	cancelCtx, cancel := context.WithCancel(context.Background())
 	cancel()
 
 	// 此处应该打印出来 context.Canceled
-	r, err := newResolver(cancelCtx, &mockDiscoveries{false, false, false}, ta, &mockRebalancer{}, false, false, 25)
+	r, err = newResolver(cancelCtx, &mockDiscoveries{false, false, false}, ta, &mockRebalancer{}, false, false, 25)
 	if err != nil {
 		t.Errorf("expect %v, got %v", nil, err)
 	}
-	_ = r.Close()
+	if r != nil {
+		_ = r.Close()
+	}
 
 	// 同步 但是服务取消，此时需要报错
-	_, err = newResolver(cancelCtx, &mockDiscoveries{false, false, true}, ta, &mockRebalancer{}, true, true, 25)
+	r, err = newResolver(cancelCtx, &mockDiscoveries{false, false, true}, ta, &mockRebalancer{}, true, true, 25)
 	if err == nil {
 		t.Errorf("expect ctx cancel err, got nil")
 	}
+	if r != nil {
+		_ = r.Close()
+	}
+
+	time.Sleep(100 * time.Millisecond)
 }
diff --git a/transport/http/server.go b/transport/http/server.go
@@ -343,7 +343,14 @@ func (s *Server) Start(ctx context.Context) error {
 // Stop stop the HTTP server.
 func (s *Server) Stop(ctx context.Context) error {
 	log.Info("[HTTP] server stopping")
-	return s.Shutdown(ctx)
+	err := s.Shutdown(ctx)
+	if err != nil {
+		if ctx.Err() != nil {
+			log.Warn("[HTTP] server couldn't stop gracefully in time, doing force stop")
+			err = s.Server.Close()
+		}
+	}
+	return err
 }
 
 func (s *Server) listenAndEndpoint() error {
diff --git a/transport/http/server_test.go b/transport/http/server_test.go
