feat: 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/rpcpubserver.go

@@ -1,18 +1,24 @@
 package internal
 
 import (
+	"context"
+	"errors"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/zeromicro/go-zero/core/discov"
 	"github.com/zeromicro/go-zero/core/netx"
+	"github.com/zeromicro/go-zero/internal/health"
 )
 
 const (
 	allEths  = "0.0.0.0"
 	envPodIp = "POD_IP"
 )
 
+var errNotReady = 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,11 +52,45 @@ type keepAliveServer struct {
 }
 
 func (s keepAliveServer) Start(fn RegisterFn) error {
-	if err := s.registerEtcd(); err != nil {
-		return err
+	errCh := make(chan error)
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	go func() {
+		defer close(errCh)
+		select {
+		case errCh <- s.Server.Start(fn):
+		case <-stopCh:
+			// prevent goroutine leak
+		}
+	}()
+
+	// 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)
+
+l:
+	for {
+		select {
+		case <-ticker.C:
+			if health.IsReady() {
+				err := s.registerEtcd()
+				if err != nil {
+					return err
+				}
+				// break for loop
+				break l
+			}
+		case <-ctx.Done():
+			return errNotReady
+		case err := <-errCh:
+			return err
+		}
 	}
+	ticker.Stop()
 
-	return s.Server.Start(fn)
+	return <-errCh
 }
 
 func figureOutListenOn(listenOn string) string {

zrpc/internal/rpcpubserver_test.go

@@ -4,6 +4,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+
 	"github.com/zeromicro/go-zero/core/discov"
 	"github.com/zeromicro/go-zero/core/netx"
 )
@@ -16,7 +18,8 @@ func TestNewRpcPubServer(t *testing.T) {
 	}, "")
 	assert.NoError(t, err)
 	assert.NotPanics(t, func() {
-		s.Start(nil)
+		s.Start(func(server *grpc.Server) {
+		})
 	})
 }
 

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.