fix(server): replace select{} in Serve() with graceful lifecycle management

server/serve_stop_test.go

@@ -0,0 +1,112 @@
+package server
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestStopBeforeServe verifies that Stop() is safe to call before Serve()
+// and does not panic.
+func TestStopBeforeServe(t *testing.T) {
+	srv, err := NewServer()
+	require.NoError(t, err)
+	// Should not panic when server has not been started
+	assert.NotPanics(t, func() {
+		srv.Stop()
+	})
+}
+
+// TestStopIdempotent verifies that calling Stop() multiple times is safe.
+func TestStopIdempotent(t *testing.T) {
+	srv, err := NewServer()
+	require.NoError(t, err)
+
+	// Start Serve in a goroutine; it will block on stopCh.
+	serveErrCh := make(chan error, 1)
+	go func() {
+		serveErrCh <- srv.Serve()
+	}()
+
+	// Give Serve() time to reach the blocking phase.
+	time.Sleep(200 * time.Millisecond)
+
+	// First Stop should succeed.
+	assert.NotPanics(t, func() {
+		srv.Stop()
+	})
+
+	// Wait for Serve to return.
+	select {
+	case err := <-serveErrCh:
+		assert.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("Serve() did not return after Stop()")
+	}
+
+	// Second Stop should be safe (no-op).
+	assert.NotPanics(t, func() {
+		srv.Stop()
+	})
+}
+
+// TestServeStopGracefulShutdown verifies that Stop() causes Serve() to return gracefully.
+func TestServeStopGracefulShutdown(t *testing.T) {
+	srv, err := NewServer()
+	require.NoError(t, err)
+
+	serveErrCh := make(chan error, 1)
+	go func() {
+		serveErrCh <- srv.Serve()
+	}()
+
+	// Allow Serve() to proceed past initialization and reach the blocking phase.
+	time.Sleep(200 * time.Millisecond)
+
+	// Trigger shutdown.
+	srv.Stop()
+
+	// Serve should return without error.
+	select {
+	case err := <-serveErrCh:
+		assert.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("Serve() did not return after Stop()")
+	}
+}
+
+// TestServeAfterStop verifies that Serve() can be called again after Stop().
+func TestServeAfterStop(t *testing.T) {
+	srv, err := NewServer()
+	require.NoError(t, err)
+
+	// First cycle
+	serveErrCh1 := make(chan error, 1)
+	go func() {
+		serveErrCh1 <- srv.Serve()
+	}()
+	time.Sleep(200 * time.Millisecond)
+	srv.Stop()
+	select {
+	case err := <-serveErrCh1:
+		assert.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("first Serve() did not return after Stop()")
+	}
+
+	// Second cycle
+	serveErrCh2 := make(chan error, 1)
+	go func() {
+		serveErrCh2 <- srv.Serve()
+	}()
+	time.Sleep(200 * time.Millisecond)
+	srv.Stop()
+	select {
+	case err := <-serveErrCh2:
+		assert.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("second Serve() did not return after Stop()")
+	}
+}

server/server.go

@@ -57,6 +57,9 @@ type Server struct {
 	interfaceNameServices map[string]*ServiceOptions
 	// indicate whether the server is already started
 	serve bool
+	// stopCh is used to signal graceful shutdown of Serve().
+	// Closing this channel causes Serve() to return instead of blocking forever with select{}.
+	stopCh chan struct{}
 }
 
 // ServiceInfo Deprecated： common.ServiceInfo type alias, just for compatible with old generate pb.go file
@@ -323,6 +326,13 @@ func (s *Server) Serve() error {
 	}
 	// prevent multiple calls to Serve
 	s.serve = true
+	// re-create stopCh if it was closed by a previous Stop(),
+	// allowing Serve() to be called again after graceful shutdown.
+	select {
+	case <-s.stopCh:
+		s.stopCh = make(chan struct{})
+	default:
+	}
 
 	// release lock in case causing deadlock
 	s.mu.Unlock()
@@ -356,7 +366,24 @@ func (s *Server) Serve() error {
 	probe.SetStartupComplete(true)
 	probe.SetReady(true)
 
-	select {}
+	// Block until Stop() is called or the stopCh is closed,
+	// enabling graceful shutdown instead of an unrecoverable hard spin.
+	<-s.stopCh
+	return nil
+}
+
+// Stop signals the server to shut down gracefully.
+// It is safe to call multiple times; only the first call has effect.
+// After Stop, Serve() will return and the server can be started again
+// by calling Serve() (if desired).
+func (s *Server) Stop() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if !s.serve {
+		return
+	}
+	s.serve = false
+	close(s.stopCh)
 }
 
 // In order to expose internal services
@@ -463,6 +490,7 @@ func NewServer(opts ...ServerOption) (*Server, error) {
 		cfg:                   newSrvOpts,
 		svcOptsMap:            make(map[string]*ServiceOptions),
 		interfaceNameServices: make(map[string]*ServiceOptions),
+		stopCh:                make(chan struct{}),
 	}
 	return srv, nil
 }