refactor(node/state): Use fx annotation from celestiaorg#923

renaynay · Sep 19, 2022 · 39ee26f · 39ee26f
1 parent bea5f08
commit 39ee26f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 25 deletions.
diff --git a/node/state/core.go b/node/state/core.go
@@ -1,9 +1,8 @@
 package state
 
 import (
-	"go.uber.org/fx"
-
 	apptypes "github.com/celestiaorg/celestia-app/x/payment/types"
+
 	"github.com/celestiaorg/celestia-node/header"
 	"github.com/celestiaorg/celestia-node/service/state"
 )
@@ -14,13 +13,8 @@ func CoreAccessor(
 	coreIP,
 	coreRPC,
 	coreGRPC string,
-) func(fx.Lifecycle, *apptypes.KeyringSigner, header.Store) (state.Accessor, error) {
-	return func(lc fx.Lifecycle, signer *apptypes.KeyringSigner, getter header.Store) (state.Accessor, error) {
-		ca := state.NewCoreAccessor(signer, getter, coreIP, coreRPC, coreGRPC)
-		lc.Append(fx.Hook{
-			OnStart: ca.Start,
-			OnStop:  ca.Stop,
-		})
-		return ca, nil
+) func(*apptypes.KeyringSigner, header.Store) state.Accessor {
+	return func(signer *apptypes.KeyringSigner, getter header.Store) state.Accessor {
+		return state.NewCoreAccessor(signer, getter, coreIP, coreRPC, coreGRPC)
 	}
 }
diff --git a/node/state/keyring.go b/node/state/keyring.go
@@ -10,6 +10,7 @@ import (
 	"github.com/celestiaorg/celestia-app/app"
 	"github.com/celestiaorg/celestia-app/app/encoding"
 	apptypes "github.com/celestiaorg/celestia-app/x/payment/types"
+
 	"github.com/celestiaorg/celestia-node/libs/keystore"
 	"github.com/celestiaorg/celestia-node/node/key"
 	"github.com/celestiaorg/celestia-node/params"

diff --git a/node/state/state.go b/node/state/state.go
@@ -23,27 +23,30 @@ func Module(coreCfg core.Config, keyCfg key.Config) fx.Option {
 	return fx.Module(
 		"state",
 		fx.Provide(Keyring(keyCfg)),
-		fx.Provide(CoreAccessor(coreCfg.IP, coreCfg.RPCPort, coreCfg.GRPCPort)),
-		fx.Provide(Service),
+		fx.Provide(fx.Annotate(CoreAccessor(coreCfg.IP, coreCfg.RPCPort, coreCfg.GRPCPort),
+			fx.OnStart(func(ctx context.Context, accessor state.Accessor) error {
+				return accessor.Start(ctx)
+			}),
+			fx.OnStop(func(ctx context.Context, accessor state.Accessor) error {
+				return accessor.Stop(ctx)
+			}),
+		)),
+		fx.Provide(fx.Annotate(Service,
+			fx.OnStart(func(ctx context.Context, lc fx.Lifecycle, fservice fraud.Service, serv *state.Service) error {
+				lifecycleCtx := fxutil.WithLifecycle(ctx, lc)
+				return services.FraudLifecycle(ctx, lifecycleCtx, fraud.BadEncoding, fservice, serv.Start, serv.Stop)
+			}),
+			fx.OnStop(func(ctx context.Context, serv *state.Service) error {
+				return serv.Stop(ctx)
+			}),
+		)),
 	)
 }
 
 // Service constructs a new state.Service.
 func Service(
-	ctx context.Context,
-	lc fx.Lifecycle,
 	accessor state.Accessor,
 	store header.Store,
-	fservice fraud.Service,
 ) *state.Service {
-	serv := state.NewService(accessor, store)
-	lifecycleCtx := fxutil.WithLifecycle(ctx, lc)
-	lc.Append(fx.Hook{
-		OnStart: func(startCtx context.Context) error {
-			return services.FraudLifecycle(startCtx, lifecycleCtx, fraud.BadEncoding, fservice, serv.Start, serv.Stop)
-		},
-		OnStop: serv.Stop,
-	})
-
-	return serv
+	return state.NewService(accessor, store)
 }