From b3fd5af8f1d567cb647b896de5553555751e7018 Mon Sep 17 00:00:00 2001
From: Fangyu Gai <greferry@gmail.com>
Date: Wed, 9 Oct 2024 03:08:24 +0800
Subject: [PATCH] fix e2e

---
 finality-provider/service/app.go        |  8 ++-
 finality-provider/service/fp_manager.go | 69 ++++++++++++-------------
 itest/test_manager.go                   |  4 +-
 3 files changed, 39 insertions(+), 42 deletions(-)

diff --git a/finality-provider/service/app.go b/finality-provider/service/app.go
index 0b12ffb8..52b9c6d7 100644
--- a/finality-provider/service/app.go
+++ b/finality-provider/service/app.go
@@ -321,11 +321,9 @@ func (app *FinalityProviderApp) Stop() error {
 		app.wg.Wait()
 
 		app.logger.Debug("Stopping finality providers")
-		if app.fpManager.isStarted.Load() {
-			if err := app.fpManager.Stop(); err != nil {
-				stopErr = err
-				return
-			}
+		if err := app.fpManager.Stop(); err != nil {
+			stopErr = err
+			return
 		}
 
 		app.logger.Debug("Stopping EOTS manager")
diff --git a/finality-provider/service/fp_manager.go b/finality-provider/service/fp_manager.go
index c52d5b85..0d78f356 100644
--- a/finality-provider/service/fp_manager.go
+++ b/finality-provider/service/fp_manager.go
@@ -8,7 +8,6 @@ import (
 
 	"github.com/avast/retry-go/v4"
 	bbntypes "github.com/babylonlabs-io/babylon/types"
-	"go.uber.org/atomic"
 	"go.uber.org/zap"
 
 	"github.com/babylonlabs-io/finality-provider/clientcontroller"
@@ -34,9 +33,9 @@ func (ce *CriticalError) Error() string {
 // FinalityProviderManager is responsible to initiate and start the given finality
 // provider instance, monitor its running status
 type FinalityProviderManager struct {
-	isStarted *atomic.Bool
+	startOnce sync.Once
+	stopOnce  sync.Once
 
-	// mutex to acess map of fp instances (fpIns)
 	wg sync.WaitGroup
 
 	fpIns *FinalityProviderInstance
@@ -67,7 +66,6 @@ func NewFinalityProviderManager(
 ) (*FinalityProviderManager, error) {
 	return &FinalityProviderManager{
 		criticalErrChan: make(chan *CriticalError),
-		isStarted:       atomic.NewBool(false),
 		fps:             fps,
 		pubRandStore:    pubRandStore,
 		config:          config,
@@ -88,8 +86,7 @@ func (fpm *FinalityProviderManager) monitorCriticalErr() {
 
 	var criticalErr *CriticalError
 
-	exitLoop := false
-	for !exitLoop {
+	for {
 		select {
 		case criticalErr = <-fpm.criticalErrChan:
 			fpi, err := fpm.GetFinalityProviderInstance()
@@ -97,7 +94,6 @@ func (fpm *FinalityProviderManager) monitorCriticalErr() {
 				fpm.logger.Debug("the finality-provider instance is already shutdown",
 					zap.String("pk", criticalErr.fpBtcPk.MarshalHex()))
 
-				exitLoop = true
 				continue
 			}
 			if errors.Is(criticalErr.err, ErrFinalityProviderSlashed) {
@@ -105,7 +101,6 @@ func (fpm *FinalityProviderManager) monitorCriticalErr() {
 				fpm.logger.Debug("the finality-provider has been slashed",
 					zap.String("pk", criticalErr.fpBtcPk.MarshalHex()))
 
-				exitLoop = true
 				continue
 			}
 			if errors.Is(criticalErr.err, ErrFinalityProviderJailed) {
@@ -113,7 +108,6 @@ func (fpm *FinalityProviderManager) monitorCriticalErr() {
 				fpm.logger.Debug("the finality-provider has been jailed",
 					zap.String("pk", criticalErr.fpBtcPk.MarshalHex()))
 
-				exitLoop = true
 				continue
 			}
 			fpm.logger.Fatal(instanceTerminatingMsg,
@@ -122,10 +116,6 @@ func (fpm *FinalityProviderManager) monitorCriticalErr() {
 			return
 		}
 	}
-
-	if err := fpm.Stop(); err != nil {
-		fpm.logger.Fatal("failed to stop the finality provider manager", zap.Error(err))
-	}
 }
 
 // monitorStatusUpdate periodically check the status of each managed finality providers and update
@@ -242,42 +232,49 @@ func (fpm *FinalityProviderManager) setFinalityProviderJailed(fpi *FinalityProvi
 }
 
 func (fpm *FinalityProviderManager) StartFinalityProvider(fpPk *bbntypes.BIP340PubKey, passphrase string) error {
-	if !fpm.isStarted.Load() {
-		fpm.isStarted.Store(true)
-
-		if err := fpm.startFinalityProviderInstance(fpPk, passphrase); err != nil {
-			return err
-		}
-
-		fpm.wg.Add(1)
+	fpm.startOnce.Do(func() {
+		fpm.wg.Add(2)
 		go fpm.monitorCriticalErr()
-
-		fpm.wg.Add(1)
 		go fpm.monitorStatusUpdate()
+	})
+
+	fpm.logger.Info("starting finality provider", zap.String("pk", fpPk.MarshalHex()))
+
+	if err := fpm.startFinalityProviderInstance(fpPk, passphrase); err != nil {
+		return err
 	}
 
+	fpm.logger.Info("finality provider is started", zap.String("pk", fpPk.MarshalHex()))
+
 	return nil
 }
 
 func (fpm *FinalityProviderManager) Stop() error {
-	if !fpm.isStarted.Swap(false) {
-		return fmt.Errorf("the finality-provider manager has already stopped")
-	}
-
-	defer func() {
+	var stopErr error
+	fpm.stopOnce.Do(func() {
 		close(fpm.quit)
 		fpm.wg.Wait()
-	}()
 
-	if fpm.fpIns == nil {
-		return nil
-	}
+		if fpm.fpIns == nil {
+			return
+		}
 
-	if !fpm.fpIns.IsRunning() {
-		return nil
-	}
+		if !fpm.fpIns.IsRunning() {
+			return
+		}
+
+		pkHex := fpm.fpIns.GetBtcPkHex()
+		fpm.logger.Info("stopping finality provider", zap.String("pk", pkHex))
+
+		if err := fpm.fpIns.Stop(); err != nil {
+			stopErr = err
+			return
+		}
+
+		fpm.logger.Info("finality provider is stopped", zap.String("pk", pkHex))
+	})
 
-	return fpm.fpIns.Stop()
+	return stopErr
 }
 
 func (fpm *FinalityProviderManager) GetFinalityProviderInstance() (*FinalityProviderInstance, error) {
diff --git a/itest/test_manager.go b/itest/test_manager.go
index 70c4abf1..96d75984 100644
--- a/itest/test_manager.go
+++ b/itest/test_manager.go
@@ -252,7 +252,9 @@ func StartManagerWithFinalityProvider(t *testing.T) (*TestManager, *service.Fina
 }
 
 func (tm *TestManager) Stop(t *testing.T) {
-	err := tm.manager.ClearResources()
+	err := tm.Fpa.Stop()
+	require.NoError(t, err)
+	err = tm.manager.ClearResources()
 	require.NoError(t, err)
 	err = os.RemoveAll(tm.baseDir)
 	require.NoError(t, err)