x

Author: chenquan

zrpc/internal/rpcpubserver.go

@@ -52,15 +52,21 @@ type keepAliveServer struct {
 }
 
 func (s keepAliveServer) Start(fn RegisterFn) error {
-	chErr := make(chan error)
+	errCh := make(chan error)
+	stopCh := make(chan struct{})
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
 	go func() {
-		defer close(chErr)
-		chErr <- s.Server.Start(fn)
+		defer close(errCh)
+		select {
+		case <-ctx.Done():
+		case errCh <- s.Server.Start(fn):
+		case <-stopCh:
+		}
 	}()
 
 	// 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:
@@ -70,20 +76,21 @@ l:
 			if health.IsReady() {
 				err := s.registerEtcd()
 				if err != nil {
+					close(stopCh)
 					return err
 				}
 				// break for loop
 				break l
 			}
 		case <-ctx.Done():
 			return errNotReady
-		case err := <-chErr:
+		case err := <-errCh:
 			return err
 		}
 	}
 	ticker.Stop()
 
-	return <-chErr
+	return <-errCh
 }
 
 func figureOutListenOn(listenOn string) string {