feat: background worker routines to shutdown client for migration

cmd/zetaclientd-supervisor/lib.go

@@ -12,16 +12,15 @@ import (
 	"runtime"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/cosmos/cosmos-sdk/client/grpc/tmservice"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
 	"github.com/hashicorp/go-getter"
 	"github.com/rs/zerolog"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 
-	observertypes "github.com/zeta-chain/zetacore/x/observer/types"
 	"github.com/zeta-chain/zetacore/zetaclient/config"
 )
 
@@ -79,7 +78,7 @@ func newZetaclientdSupervisor(
 	logger = logger.With().Str("module", "zetaclientdSupervisor").Logger()
 	conn, err := grpc.Dial(
 		fmt.Sprintf("%s:9090", zetaCoreURL),
-		grpc.WithInsecure(),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("grpc dial: %w", err)
@@ -99,9 +98,6 @@ func newZetaclientdSupervisor(
 func (s *zetaclientdSupervisor) Start(ctx context.Context) {
 	go s.watchForVersionChanges(ctx)
 	go s.handleCoreUpgradePlan(ctx)
-	go s.handleNewKeygen(ctx)
-	go s.handleNewTSSKeyGeneration(ctx)
-	go s.handleTSSUpdate(ctx)
 }
 
 func (s *zetaclientdSupervisor) WaitForReloadSignal(ctx context.Context) {
@@ -177,125 +173,6 @@ func (s *zetaclientdSupervisor) watchForVersionChanges(ctx context.Context) {
 	}
 }
 
-func (s *zetaclientdSupervisor) handleTSSUpdate(ctx context.Context) {
-	maxRetries := 11
-	retryInterval := 5 * time.Second
-
-	// TODO : use retry library under pkg/retry
-	// https://github.com/zeta-chain/node/issues/2492
-	for i := 0; i < maxRetries; i++ {
-		client := observertypes.NewQueryClient(s.zetacoredConn)
-		tss, err := client.TSS(ctx, &observertypes.QueryGetTSSRequest{})
-		if err != nil {
-			s.logger.Warn().Err(err).Msg("unable to get original tss")
-			time.Sleep(retryInterval)
-			continue
-		}
-		i = 0
-		for {
-			select {
-			case <-time.After(time.Second):
-			case <-ctx.Done():
-				return
-			}
-			tssNew, err := client.TSS(ctx, &observertypes.QueryGetTSSRequest{})
-			if err != nil {
-				s.logger.Warn().Err(err).Msg("unable to get tss")
-				continue
-			}
-
-			if tssNew.TSS.TssPubkey == tss.TSS.TssPubkey {
-				continue
-			}
-
-			tss = tssNew
-			s.logger.Info().
-				Msgf("tss address is updated from %s to %s", tss.TSS.TssPubkey, tssNew.TSS.TssPubkey)
-			time.Sleep(6 * time.Second)
-			s.logger.Info().Msg("restarting zetaclientd to update tss address")
-			s.restartChan <- syscall.SIGHUP
-		}
-	}
-	s.logger.Warn().Msg("handleTSSUpdate exiting without success")
-}
-
-func (s *zetaclientdSupervisor) handleNewTSSKeyGeneration(ctx context.Context) {
-	maxRetries := 11
-	retryInterval := 5 * time.Second
-
-	// TODO : use retry library under pkg/retry
-	for i := 0; i < maxRetries; i++ {
-		client := observertypes.NewQueryClient(s.zetacoredConn)
-		alltss, err := client.TssHistory(ctx, &observertypes.QueryTssHistoryRequest{})
-		if err != nil {
-			s.logger.Warn().Err(err).Msg("unable to get tss original history")
-			time.Sleep(retryInterval)
-			continue
-		}
-		i = 0
-		tssLenCurrent := len(alltss.TssList)
-		for {
-			select {
-			case <-time.After(time.Second):
-			case <-ctx.Done():
-				return
-			}
-			tssListNew, err := client.TssHistory(ctx, &observertypes.QueryTssHistoryRequest{})
-			if err != nil {
-				s.logger.Warn().Err(err).Msg("unable to get tss new history")
-				continue
-			}
-			tssLenUpdated := len(tssListNew.TssList)
-
-			if tssLenUpdated == tssLenCurrent {
-				continue
-			}
-			if tssLenUpdated < tssLenCurrent {
-				tssLenCurrent = len(tssListNew.TssList)
-				continue
-			}
-
-			tssLenCurrent = tssLenUpdated
-			s.logger.Info().Msgf("tss list updated from %d to %d", tssLenCurrent, tssLenUpdated)
-			time.Sleep(5 * time.Second)
-			s.logger.Info().Msg("restarting zetaclientd to update tss list")
-			s.restartChan <- syscall.SIGHUP
-		}
-	}
-	s.logger.Warn().Msg("handleNewTSSKeyGeneration exiting without success")
-}
-
-func (s *zetaclientdSupervisor) handleNewKeygen(ctx context.Context) {
-	client := observertypes.NewQueryClient(s.zetacoredConn)
-	prevKeygenBlock := int64(0)
-	for {
-		select {
-		case <-time.After(time.Second):
-		case <-ctx.Done():
-			return
-		}
-		resp, err := client.Keygen(ctx, &observertypes.QueryGetKeygenRequest{})
-		if err != nil {
-			s.logger.Warn().Err(err).Msg("unable to get keygen")
-			continue
-		}
-		if resp.Keygen == nil {
-			s.logger.Warn().Err(err).Msg("keygen is nil")
-			continue
-		}
-
-		if resp.Keygen.Status != observertypes.KeygenStatus_PendingKeygen {
-			continue
-		}
-		keygenBlock := resp.Keygen.BlockNumber
-		if prevKeygenBlock == keygenBlock {
-			continue
-		}
-		prevKeygenBlock = keygenBlock
-		s.logger.Info().Msgf("got new keygen at block %d", keygenBlock)
-		s.restartChan <- syscall.SIGHUP
-	}
-}
 func (s *zetaclientdSupervisor) handleCoreUpgradePlan(ctx context.Context) {
 	client := upgradetypes.NewQueryClient(s.zetacoredConn)
 
@@ -345,16 +222,16 @@ func (s *zetaclientdSupervisor) downloadZetaclientd(ctx context.Context, plan *u
 	if plan.Info == "" {
 		return errors.New("upgrade info empty")
 	}
-	var config upgradeConfig
-	err := json.Unmarshal([]byte(plan.Info), &config)
+	var cfg upgradeConfig
+	err := json.Unmarshal([]byte(plan.Info), &cfg)
 	if err != nil {
 		return fmt.Errorf("unmarshal upgrade config: %w", err)
 	}
 
 	s.logger.Info().Msg("downloading zetaclientd")
 
 	binKey := fmt.Sprintf("%s-%s/%s", zetaclientdBinaryName, runtime.GOOS, runtime.GOARCH)
-	binURL, ok := config.Binaries[binKey]
+	binURL, ok := cfg.Binaries[binKey]
 	if !ok {
 		return fmt.Errorf("no binary found for: %s", binKey)
 	}

cmd/zetaclientd-supervisor/main.go

@@ -69,10 +69,18 @@ func main() {
 		cmd.Stdin = &passwordInputBuffer
 
 		eg, ctx := errgroup.WithContext(ctx)
-		eg.Go(cmd.Run)
+		eg.Go(func() error {
+			defer cancel()
+			if err := cmd.Run(); err != nil {
+				logger.Error().Err(err).Msg("zetaclient process exited with error")
+				return err
+			}
+
+			logger.Info().Msg("zetaclient process exited")
+			return nil
+		})
 		eg.Go(func() error {
 			supervisor.WaitForReloadSignal(ctx)
-			cancel()
 			return nil
 		})
 		eg.Go(func() error {

cmd/zetaclientd/start.go

@@ -210,6 +210,14 @@ func start(_ *cobra.Command, _ []string) error {
 	// Set P2P ID for telemetry
 	telemetryServer.SetP2PID(server.GetLocalPeerID())
 
+	// Create a notification context for background threads.
+	// These threads are responsible for sending shutdown signals to the main thread.
+	notifyCtx, cancel := context.WithCancelCause(ctx)
+
+	// Start background threads which monitors zeta core for state changes related to TSS migration
+	cancelBackgroundThreads := zetacoreClient.StartTssMigrationRoutines(notifyCtx, cancel, masterLogger)
+	defer cancelBackgroundThreads()
+
 	// Generate a new TSS if keygen is set and add it into the tss server
 	// If TSS has already been generated, and keygen was successful ; we use the existing TSS
 	err = GenerateTSS(ctx, masterLogger, zetacoreClient, server)
@@ -318,15 +326,16 @@ func start(_ *cobra.Command, _ []string) error {
 	}
 
 	// Orchestrator wraps the zetacore client and adds the observers and signer maps to it . This is the high level object used for CCTX interactions
-	orchestrator := orchestrator.NewOrchestrator(
+	cctxOrchestrator := orchestrator.NewOrchestrator(
 		ctx,
 		zetacoreClient,
 		signerMap,
 		observerMap,
 		masterLogger,
 		telemetryServer,
 	)
-	err = orchestrator.MonitorCore(ctx)
+
+	err = cctxOrchestrator.MonitorCore(ctx)
 	if err != nil {
 		startLogger.Error().Err(err).Msg("Orchestrator failed to start")
 		return err
@@ -346,18 +355,24 @@ func start(_ *cobra.Command, _ []string) error {
 	//	defer zetaSupplyChecker.Stop()
 	//}
 
-	startLogger.Info().Msgf("awaiting the os.Interrupt, syscall.SIGTERM signals...")
 	ch := make(chan os.Signal, 1)
 	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
-	sig := <-ch
-	startLogger.Info().Msgf("stop signal received: %s", sig)
 
-	// stop chain observers
+	startLogger.Info().Msgf("awaiting shutdown signals")
+	select {
+	case <-notifyCtx.Done():
+		cause := context.Cause(notifyCtx)
+		startLogger.Info().Msgf("shutdown signal received , cause : %s", cause)
+
+	case sig := <-ch:
+		startLogger.Info().Msgf("stop signal received: %s", sig)
+	}
+
+	//stop chain observers
 	for _, observer := range observerMap {
 		observer.Stop()
 	}
 	zetacoreClient.Stop()
-
 	return nil
 }
 

contrib/localnet/docker-compose.yml

@@ -124,6 +124,7 @@ services:
       - ETHDEV_ENDPOINT=http://eth:8545
       - HOTKEY_BACKEND=file
       - HOTKEY_PASSWORD=password # test purposes only
+    restart: always
     volumes:
       - ssh:/root/.ssh
       - preparams:/root/preparams
@@ -140,6 +141,7 @@ services:
       - ETHDEV_ENDPOINT=http://eth:8545
       - HOTKEY_BACKEND=file
       - HOTKEY_PASSWORD=password # test purposes only
+    restart: always
     volumes:
       - ssh:/root/.ssh
       - preparams:/root/preparams

pkg/bg/bg.go

@@ -11,6 +11,7 @@ import (
 type config struct {
 	name   string
 	logger zerolog.Logger
+	cancel context.CancelCauseFunc
 }
 
 type Opt func(*config)
@@ -19,6 +20,10 @@ func WithName(name string) Opt {
 	return func(cfg *config) { cfg.name = name }
 }
 
+func WithCancel(cancel context.CancelCauseFunc) Opt {
+	return func(cfg *config) { cfg.cancel = cancel }
+}
+
 func WithLogger(logger zerolog.Logger) Opt {
 	return func(cfg *config) { cfg.logger = logger }
 }
@@ -28,6 +33,7 @@ func Work(ctx context.Context, f func(context.Context) error, opts ...Opt) {
 	cfg := config{
 		name:   "",
 		logger: zerolog.Nop(),
+		cancel: nil,
 	}
 
 	for _, opt := range opts {
@@ -42,9 +48,15 @@ func Work(ctx context.Context, f func(context.Context) error, opts ...Opt) {
 			}
 		}()
 
-		if err := f(ctx); err != nil {
+		err := f(ctx)
+		if err != nil {
 			logError(err, cfg)
 		}
+		// Use cancel function if it is provided.
+		// This is used for stopping the main thread based on the outcome of the background task.
+		if cfg.cancel != nil {
+			cfg.cancel(fmt.Errorf("cancel function triggered by %s", cfg.name))
+		}
 	}()
 }
 

pkg/retry/retry.go

@@ -46,6 +46,11 @@ func DefaultBackoff() Backoff {
 	return backoff.WithMaxRetries(bo, 5)
 }
 
+func DefaultConstantBackoff() Backoff {
+	bo := backoff.NewConstantBackOff(5 * time.Second)
+	return backoff.WithMaxRetries(bo, 10)
+}
+
 // Do executes the callback function with the default backoff config.
 // It will retry a callback ONLY if error is retryable.
 func Do(cb Callback) error {

zetaclient/context/app.go

@@ -70,6 +70,10 @@ func (a *AppContext) Config() config.Config {
 	return a.config
 }
 
+func (a *AppContext) Logger() zerolog.Logger {
+	return a.logger
+}
+
 // GetBTCChainAndConfig returns btc chain and config if enabled
 func (a *AppContext) GetBTCChainAndConfig() (chains.Chain, config.BTCConfig, bool) {
 	btcConfig, configEnabled := a.Config().GetBTCConfig()

zetaclient/orchestrator/orchestrator.go

@@ -116,9 +116,6 @@ func (oc *Orchestrator) MonitorCore(ctx context.Context) error {
 	shutdownOrchestrator := func() {
 		// now stop orchestrator and all observers
 		close(oc.stop)
-		for _, c := range oc.observerMap {
-			c.Stop()
-		}
 	}
 
 	oc.zetacoreClient.OnBeforeStop(shutdownOrchestrator)

zetaclient/zetacore/client_query_observer.go

@@ -167,6 +167,14 @@ func (c *Client) GetBTCTSSAddress(ctx context.Context, chainID int64) (string, e
 	return resp.Btc, nil
 }
 
+func (c *Client) GetTSS(ctx context.Context) (types.TSS, error) {
+	resp, err := c.client.observer.TSS(ctx, &types.QueryGetTSSRequest{})
+	if err != nil {
+		return types.TSS{}, errors.Wrap(err, "failed to get tss")
+	}
+	return resp.TSS, nil
+}
+
 // GetTSSHistory returns the TSS history
 func (c *Client) GetTSSHistory(ctx context.Context) ([]types.TSS, error) {
 	resp, err := c.client.observer.TssHistory(ctx, &types.QueryTssHistoryRequest{})