Apply suggestion to use a CommitteeProvider interface

ssvlabs · Jul 16, 2024 · 04f606e · 04f606e
1 parent 2a5e1a2
commit 04f606e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/network/p2p/p2p_setup.go b/network/p2p/p2p_setup.go
@@ -28,7 +28,6 @@ import (
 	"github.com/ssvlabs/ssv/network/records"
 	"github.com/ssvlabs/ssv/network/streams"
 	"github.com/ssvlabs/ssv/network/topics"
-	"github.com/ssvlabs/ssv/registry/storage"
 	"github.com/ssvlabs/ssv/utils/commons"
 )
 
@@ -301,7 +300,7 @@ func (n *p2pNetwork) setupPubsub(logger *zap.Logger) error {
 	// run GC every 3 minutes to clear old messages
 	async.RunEvery(n.ctx, time.Minute*3, midHandler.GC)
 
-	_, tc, err := topics.NewPubSub(n.ctx, logger, cfg, n.metrics, func() []*storage.Committee { return n.nodeStorage.ValidatorStore().Committees() })
+	_, tc, err := topics.NewPubSub(n.ctx, logger, cfg, n.metrics, n.nodeStorage.ValidatorStore())
 	if err != nil {
 		return errors.Wrap(err, "could not setup pubsub")
 	}

diff --git a/network/topics/pubsub.go b/network/topics/pubsub.go
@@ -108,8 +108,12 @@ func (cfg *PubSubConfig) initScoring() {
 	}
 }
 
+type CommitteesProvider interface {
+	Committees() []*storage.Committee
+}
+
 // NewPubSub creates a new pubsub router and the necessary components
-func NewPubSub(ctx context.Context, logger *zap.Logger, cfg *PubSubConfig, metrics Metrics, getCommittees func() []*storage.Committee) (*pubsub.PubSub, Controller, error) {
+func NewPubSub(ctx context.Context, logger *zap.Logger, cfg *PubSubConfig, metrics Metrics, committeesProvider CommitteesProvider) (*pubsub.PubSub, Controller, error) {
 	if err := cfg.init(); err != nil {
 		return nil, nil, err
 	}
@@ -165,7 +169,7 @@ func NewPubSub(ctx context.Context, logger *zap.Logger, cfg *PubSubConfig, metri
 				return 100, 100, 10, nil
 			}
 		}
-		topicScoreFactory = topicScoreParams(logger, cfg, getCommittees)
+		topicScoreFactory = topicScoreParams(logger, cfg, committeesProvider)
 	}
 
 	if cfg.MsgIDHandler != nil {

diff --git a/network/topics/scoring.go b/network/topics/scoring.go
@@ -96,7 +96,7 @@ func scoreInspector(logger *zap.Logger, scoreIdx peers.ScoreIndex, logFrequency
 }
 
 // topicScoreParams factory for creating scoring params for topics
-func topicScoreParams(logger *zap.Logger, cfg *PubSubConfig, getCommittees func() []*storage.Committee) func(string) *pubsub.TopicScoreParams {
+func topicScoreParams(logger *zap.Logger, cfg *PubSubConfig, committeesProvider CommitteesProvider) func(string) *pubsub.TopicScoreParams {
 	return func(t string) *pubsub.TopicScoreParams {
 
 		// Get validator stats
@@ -110,7 +110,7 @@ func topicScoreParams(logger *zap.Logger, cfg *PubSubConfig, getCommittees func(
 		logger.Debug("got validator stats for score params")
 
 		// Get committees
-		committees := getCommittees()
+		committees := committeesProvider.Committees()
 		topicCommittees := filterCommitteesForTopic(t, committees)
 
 		// Log