From 39ee26f8453bf93807fc339f1f539c2853066042 Mon Sep 17 00:00:00 2001
From: rene <41963722+renaynay@users.noreply.github.com>
Date: Sun, 21 Aug 2022 12:05:09 +0200
Subject: [PATCH] refactor(node/state): Use fx annotation from #923

---
 node/state/core.go    | 14 ++++----------
 node/state/keyring.go |  1 +
 node/state/state.go   | 33 ++++++++++++++++++---------------
 3 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/node/state/core.go b/node/state/core.go
index 804c7f07d2..0fb045fd15 100644
--- 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
index e7da28d776..de22e6e164 100644
--- 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
index 7c568eb846..240a546310 100644
--- 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)
 }