stateroot: set networking callback in a more straightforward way

pkg/core/stateroot_test.go

@@ -83,7 +83,7 @@ func TestStateRoot(t *testing.T) {
 	defer os.RemoveAll(tmpDir)
 	w := createAndWriteWallet(t, accs[0], path.Join(tmpDir, "w"), "pass")
 	cfg := createStateRootConfig(w.Path(), "pass")
-	srv, err := stateroot.New(cfg, zaptest.NewLogger(t), bc)
+	srv, err := stateroot.New(cfg, zaptest.NewLogger(t), bc, nil)
 	require.NoError(t, err)
 	require.EqualValues(t, 0, srv.CurrentValidatedHeight())
 	r, err := srv.GetStateRoot(bc.BlockHeight())
@@ -153,7 +153,7 @@ func TestStateRootInitNonZeroHeight(t *testing.T) {
 		defer os.RemoveAll(tmpDir)
 		w := createAndWriteWallet(t, accs[0], path.Join(tmpDir, "w"), "pass")
 		cfg := createStateRootConfig(w.Path(), "pass")
-		srv, err := stateroot.New(cfg, zaptest.NewLogger(t), bc)
+		srv, err := stateroot.New(cfg, zaptest.NewLogger(t), bc, nil)
 		require.NoError(t, err)
 		r, err := srv.GetStateRoot(2)
 		require.NoError(t, err)
@@ -199,17 +199,16 @@ func TestStateRootFull(t *testing.T) {
 	h, pubs, accs := newMajorityMultisigWithGAS(t, 2)
 	w := createAndWriteWallet(t, accs[1], path.Join(tmpDir, "wallet2"), "two")
 	cfg := createStateRootConfig(w.Path(), "two")
-	srv, err := stateroot.New(cfg, zaptest.NewLogger(t), bc)
-	require.NoError(t, err)
-	srv.Run()
-	t.Cleanup(srv.Shutdown)
 
 	var lastValidated atomic.Value
 	var lastHeight atomic.Uint32
-	srv.SetRelayCallback(func(ep *payload.Extensible) {
+	srv, err := stateroot.New(cfg, zaptest.NewLogger(t), bc, func(ep *payload.Extensible) {
 		lastHeight.Store(ep.ValidBlockStart)
 		lastValidated.Store(ep)
 	})
+	require.NoError(t, err)
+	srv.Run()
+	t.Cleanup(srv.Shutdown)
 
 	bc.setNodesByRole(t, true, noderoles.StateValidator, pubs)
 	transferTokenFromMultisigAccount(t, bc, h, bc.contracts.GAS.Hash, 1_0000_0000)

pkg/network/server.go

@@ -177,7 +177,7 @@ func newServerFromConstructors(config ServerConfig, chain blockchainer.Blockchai
 		return nil, errors.New("`StateRootInHeader` should be disabled when state service is enabled")
 	}
 
-	sr, err := stateroot.New(config.StateRootCfg, s.log, chain)
+	sr, err := stateroot.New(config.StateRootCfg, s.log, chain, s.handleNewPayload)
 	if err != nil {
 		return nil, fmt.Errorf("can't initialize StateRoot service: %w", err)
 	}
@@ -221,10 +221,6 @@ func newServerFromConstructors(config ServerConfig, chain blockchainer.Blockchai
 
 	s.consensus = srv
 
-	if config.StateRootCfg.Enabled {
-		s.stateRoot.SetRelayCallback(s.handleNewPayload)
-	}
-
 	if s.MinPeers < 0 {
 		s.log.Info("bad MinPeers configured, using the default value",
 			zap.Int("configured", s.MinPeers),

pkg/services/stateroot/network.go

@@ -94,18 +94,5 @@ func (s *service) sendValidatedRoot(r *state.MPTRoot, priv *keys.PrivateKey) {
 	buf := io.NewBufBinWriter()
 	emit.Bytes(buf.BinWriter, sig)
 	ep.Witness.InvocationScript = buf.Bytes()
-	s.getRelayCallback()(ep)
-}
-
-func (s *service) getRelayCallback() RelayCallback {
-	s.cbMtx.RLock()
-	defer s.cbMtx.RUnlock()
-	return s.onValidatedRoot
-}
-
-// SetRelayCallback sets callback to pool and broadcast tx.
-func (s *service) SetRelayCallback(cb RelayCallback) {
-	s.cbMtx.Lock()
-	defer s.cbMtx.Unlock()
-	s.onValidatedRoot = cb
+	s.onValidatedRoot(ep)
 }

pkg/services/stateroot/service.go

@@ -24,7 +24,6 @@ type (
 		OnPayload(p *payload.Extensible) error
 		AddSignature(height uint32, validatorIndex int32, sig []byte) error
 		GetConfig() config.StateRoot
-		SetRelayCallback(RelayCallback)
 		Run()
 		Shutdown()
 	}
@@ -46,7 +45,6 @@ type (
 		srMtx           sync.Mutex
 		incompleteRoots map[uint32]*incompleteRoot
 
-		cbMtx           sync.RWMutex
 		onValidatedRoot RelayCallback
 		blockCh         chan *block.Block
 		done            chan struct{}
@@ -59,7 +57,7 @@ const (
 )
 
 // New returns new state root service instance using underlying module.
-func New(cfg config.StateRoot, log *zap.Logger, bc blockchainer.Blockchainer) (Service, error) {
+func New(cfg config.StateRoot, log *zap.Logger, bc blockchainer.Blockchainer, cb RelayCallback) (Service, error) {
 	s := &service{
 		StateRoot:       bc.GetStateModule(),
 		Network:         bc.GetConfig().Magic,
@@ -68,6 +66,7 @@ func New(cfg config.StateRoot, log *zap.Logger, bc blockchainer.Blockchainer) (S
 		incompleteRoots: make(map[uint32]*incompleteRoot),
 		blockCh:         make(chan *block.Block),
 		done:            make(chan struct{}),
+		onValidatedRoot: cb,
 	}
 
 	s.MainCfg = cfg

pkg/services/stateroot/validators.go

@@ -91,7 +91,7 @@ func (s *service) signAndSend(r *state.MPTRoot) error {
 	buf := io.NewBufBinWriter()
 	emit.Bytes(buf.BinWriter, sig)
 	e.Witness.InvocationScript = buf.Bytes()
-	s.getRelayCallback()(e)
+	s.onValidatedRoot(e)
 	return nil
 }