feat: Add shutdown interception and readiness check

Author: chenquan

internal/health/health.go

@@ -43,6 +43,11 @@ func AddProbe(probe Probe) {
 	defaultHealthManager.addProbe(probe)
 }
 
+// IsReady return global comboHealthManager status.
+func IsReady() bool {
+	return defaultHealthManager.IsReady()
+}
+
 // CreateHttpHandler create health http handler base on given probe.
 func CreateHttpHandler(healthResponse string) http.HandlerFunc {
 	return func(w http.ResponseWriter, _ *http.Request) {

zrpc/internal/config.go

@@ -23,6 +23,7 @@ type (
 		StatConf   StatConf `json:",optional"`
 		Prometheus bool     `json:",default=true"`
 		Breaker    bool     `json:",default=true"`
+		Shutdown   bool     `json:",default=true"`
 	}
 
 	// MethodTimeoutConf defines specified timeout for gRPC methods.

zrpc/internal/rpcpubserver.go

@@ -1,18 +1,25 @@
 package internal
 
 import (
+	"context"
+	"errors"
 	"os"
 	"strings"
+	"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"
 )
 
 const (
 	allEths  = "0.0.0.0"
 	envPodIp = "POD_IP"
 )
 
+var notReadyError = errors.New("service is not ready for a limited time")
+
 // NewRpcPubServer returns a Server.
 func NewRpcPubServer(etcd discov.EtcdConf, listenOn string,
 	opts ...ServerOption) (Server, error) {
@@ -46,9 +53,27 @@ type keepAliveServer struct {
 }
 
 func (s keepAliveServer) Start(fn RegisterFn) error {
-	if err := s.registerEtcd(); err != nil {
-		return err
-	}
+	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)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				if health.IsReady() {
+					logx.Must(s.registerEtcd())
+					return
+				}
+			case <-ctx.Done():
+				logx.Must(notReadyError)
+				return
+			}
+		}
+	}()
 
 	return s.Server.Start(fn)
 }

zrpc/internal/serverinterceptors/shutdowninterceptor.go

@@ -0,0 +1,50 @@
+package serverinterceptors
+
+import (
+	"context"
+	"sync/atomic"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/zeromicro/go-zero/core/proc"
+)
+
+var (
+	shutdown      int32
+	shutdownError = status.Error(codes.Unavailable, "server is shutting down")
+)
+
+func init() {
+	proc.AddWrapUpListener(func() {
+		atomic.StoreInt32(&shutdown, 1)
+	})
+}
+
+// StreamShutdownInterceptor returns a grpc.StreamServerInterceptor that handles server shutdown.
+func StreamShutdownInterceptor(svr any, stream grpc.ServerStream, _ *grpc.StreamServerInfo,
+	handler grpc.StreamHandler) error {
+
+	return handleShutdown(func() error {
+		return handler(svr, stream)
+	})
+}
+
+// UnaryShutdownInterceptor returns a grpc.UnaryServerInterceptor that handles server shutdown.
+func UnaryShutdownInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo,
+	handler grpc.UnaryHandler) (resp any, err error) {
+
+	err = handleShutdown(func() error {
+		resp, err = handler(ctx, req)
+		return err
+	})
+	return
+}
+
+func handleShutdown(f func() error) error {
+	if atomic.LoadInt32(&shutdown) == 1 {
+		return shutdownError
+	}
+	return f()
+}

zrpc/server.go

@@ -3,14 +3,15 @@ package zrpc
 import (
 	"time"
 
+	"google.golang.org/grpc"
+
 	"github.com/zeromicro/go-zero/core/load"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/core/stores/redis"
 	"github.com/zeromicro/go-zero/zrpc/internal"
 	"github.com/zeromicro/go-zero/zrpc/internal/auth"
 	"github.com/zeromicro/go-zero/zrpc/internal/serverinterceptors"
-	"google.golang.org/grpc"
 )
 
 // A RpcServer is a rpc server.
@@ -139,6 +140,9 @@ func setupStreamInterceptors(svr internal.Server, c RpcServerConf) {
 	if c.Middlewares.Breaker {
 		svr.AddStreamInterceptors(serverinterceptors.StreamBreakerInterceptor)
 	}
+	if c.Middlewares.Shutdown {
+		svr.AddStreamInterceptors(serverinterceptors.StreamShutdownInterceptor)
+	}
 }
 
 func setupUnaryInterceptors(svr internal.Server, c RpcServerConf, metrics *stat.Metrics) {
@@ -157,6 +161,9 @@ func setupUnaryInterceptors(svr internal.Server, c RpcServerConf, metrics *stat.
 	if c.Middlewares.Breaker {
 		svr.AddUnaryInterceptors(serverinterceptors.UnaryBreakerInterceptor)
 	}
+	if c.Middlewares.Shutdown {
+		svr.AddUnaryInterceptors(serverinterceptors.UnaryShutdownInterceptor)
+	}
 	if c.CpuThreshold > 0 {
 		shedder := load.NewAdaptiveShedder(load.WithCpuThreshold(c.CpuThreshold))
 		svr.AddUnaryInterceptors(serverinterceptors.UnarySheddingInterceptor(shedder, metrics))