From 503ec03f37a5a76d87ef95451ff9d2b36393365e Mon Sep 17 00:00:00 2001
From: Tony Chen <zhihui_chen@foxmail.com>
Date: Fri, 27 May 2022 01:15:21 +0800
Subject: [PATCH] app: fix instance nil when not registered (#2059)

* fix instance nil when not registered

* fix data race
---
 app.go          | 23 ++++++++++++-----------
 app_test.go     |  2 +-
 options.go      |  6 ++----
 options_test.go |  2 +-
 4 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/app.go b/app.go
index 7a1ed4463..df992250d 100644
--- a/app.go
+++ b/app.go
@@ -31,7 +31,7 @@ type App struct {
 	opts     options
 	ctx      context.Context
 	cancel   func()
-	lk       sync.Mutex
+	mu       sync.Mutex
 	instance *registry.ServiceInstance
 }
 
@@ -39,7 +39,6 @@ type App struct {
 func New(opts ...Option) *App {
 	o := options{
 		ctx:              context.Background(),
-		logger:           log.NewHelper(log.GetLogger()),
 		sigs:             []os.Signal{syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT},
 		registrarTimeout: 10 * time.Second,
 		stopTimeout:      10 * time.Second,
@@ -50,6 +49,9 @@ func New(opts ...Option) *App {
 	for _, opt := range opts {
 		opt(&o)
 	}
+	if o.logger != nil {
+		log.SetLogger(o.logger)
+	}
 	ctx, cancel := context.WithCancel(o.ctx)
 	return &App{
 		ctx:    ctx,
@@ -72,10 +74,10 @@ func (a *App) Metadata() map[string]string { return a.opts.metadata }
 
 // Endpoint returns endpoints.
 func (a *App) Endpoint() []string {
-	if a.instance == nil {
-		return []string{}
+	if a.instance != nil {
+		return a.instance.Endpoints
 	}
-	return a.instance.Endpoints
+	return nil
 }
 
 // Run executes all OnStart hooks registered with the application's Lifecycle.
@@ -84,6 +86,9 @@ func (a *App) Run() error {
 	if err != nil {
 		return err
 	}
+	a.mu.Lock()
+	a.instance = instance
+	a.mu.Unlock()
 	eg, ctx := errgroup.WithContext(NewContext(a.ctx, a))
 	wg := sync.WaitGroup{}
 	for _, srv := range a.opts.servers {
@@ -107,9 +112,6 @@ func (a *App) Run() error {
 		if err := a.opts.registrar.Register(rctx, instance); err != nil {
 			return err
 		}
-		a.lk.Lock()
-		a.instance = instance
-		a.lk.Unlock()
 	}
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, a.opts.sigs...)
@@ -120,7 +122,6 @@ func (a *App) Run() error {
 				return ctx.Err()
 			case <-c:
 				if err := a.Stop(); err != nil {
-					a.opts.logger.Errorf("failed to stop app: %v", err)
 					return err
 				}
 			}
@@ -134,9 +135,9 @@ func (a *App) Run() error {
 
 // Stop gracefully stops the application.
 func (a *App) Stop() error {
-	a.lk.Lock()
+	a.mu.Lock()
 	instance := a.instance
-	a.lk.Unlock()
+	a.mu.Unlock()
 	if a.opts.registrar != nil && instance != nil {
 		ctx, cancel := context.WithTimeout(NewContext(a.ctx, a), a.opts.registrarTimeout)
 		defer cancel()
diff --git a/app_test.go b/app_test.go
index 9165ce075..80831b541 100644
--- a/app_test.go
+++ b/app_test.go
@@ -154,7 +154,7 @@ func TestApp_Endpoint(t *testing.T) {
 				endpoint []string
 				metadata map[string]string
 			}{
-				id: "3", version: "v3", name: "kratos-v3", endpoint: []string{},
+				id: "3", version: "v3", name: "kratos-v3", endpoint: nil,
 				metadata: map[string]string{},
 			},
 		},
diff --git a/options.go b/options.go
index d18d75da1..8f9e9a55d 100644
--- a/options.go
+++ b/options.go
@@ -25,7 +25,7 @@ type options struct {
 	ctx  context.Context
 	sigs []os.Signal
 
-	logger           *log.Helper
+	logger           log.Logger
 	registrar        registry.Registrar
 	registrarTimeout time.Duration
 	stopTimeout      time.Duration
@@ -64,9 +64,7 @@ func Context(ctx context.Context) Option {
 
 // Logger with service logger.
 func Logger(logger log.Logger) Option {
-	return func(o *options) {
-		o.logger = log.NewHelper(logger)
-	}
+	return func(o *options) { o.logger = logger }
 }
 
 // Server with transport servers.
diff --git a/options_test.go b/options_test.go
index 953c52f54..4880bd195 100644
--- a/options_test.go
+++ b/options_test.go
@@ -79,7 +79,7 @@ func TestLogger(t *testing.T) {
 	o := &options{}
 	v := xlog.NewStdLogger(log.Writer())
 	Logger(v)(o)
-	if !reflect.DeepEqual(xlog.NewHelper(v), o.logger) {
+	if !reflect.DeepEqual(v, o.logger) {
 		t.Fatalf("o.logger:%v is not equal to xlog.NewHelper(v):%v", o.logger, xlog.NewHelper(v))
 	}
 }