x

Author: chenquan

zrpc/internal/rpcpubserver.go

@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/zeromicro/go-zero/core/discov"
-	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/netx"
 	"github.com/zeromicro/go-zero/internal/health"
 )
@@ -53,29 +52,34 @@ type keepAliveServer struct {
 }
 
 func (s keepAliveServer) Start(fn RegisterFn) error {
+	chErr := make(chan error)
+	defer close(chErr)
+
 	go func() {
-		// Wait for the service to start successfully, otherwise the registration service will fail.
-		ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
-		defer cancel()
-
-		ticker := time.NewTicker(time.Second / 2)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-ticker.C:
-				if health.IsReady() {
-					logx.Must(s.registerEtcd())
-					return
+		chErr <- s.Server.Start(fn)
+	}()
+
+	// Wait for the service to start successfully, otherwise the registration service will fail.
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if health.IsReady() {
+				err := s.registerEtcd()
+				if err != nil {
+					return err
 				}
-			case <-ctx.Done():
-				logx.Must(errNotReady)
-				return
+				return <-chErr
 			}
+		case <-ctx.Done():
+			return errNotReady
 		}
-	}()
-
-	return s.Server.Start(fn)
+	}
 }
 
 func figureOutListenOn(listenOn string) string {

zrpc/internal/rpcpubserver_test.go

@@ -19,7 +19,7 @@ func TestNewRpcPubServer(t *testing.T) {
 	}, "")
 	assert.NoError(t, err)
 	health.MarkReady()
-	assert.Panics(t, func() {
+	assert.NotPanics(t, func() {
 		s.Start(func(server *grpc.Server) {
 		})
 	})

zrpc/internal/serverinterceptors/shutdowninterceptor.go

@@ -43,6 +43,11 @@ func UnaryShutdownInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerI
 }
 
 func handleShutdown(f func() error) error {
+	//
+	// if health.IsReady() {
+	// 	return errShutdown
+	// }
+
 	if atomic.LoadInt32(&shutdown) == 1 {
 		return errShutdown
 	}