Extract registry to parameter (#261)

NethermindEth · Jun 26, 2024 · dfb1874 · dfb1874
1 parent df5a777
commit dfb1874
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 25 deletions.
diff --git a/aggregator/aggregator.go b/aggregator/aggregator.go
@@ -138,12 +138,13 @@ var _ RpcAggregatorer = (*Aggregator)(nil)
 var _ RestAggregatorer = (*Aggregator)(nil)
 
 // NewAggregator creates a new Aggregator with the provided config.
-// TODO: Remove this context once OperatorPubkeysServiceInMemory's API is
-// changed and we can gracefully exit otherwise
-func NewAggregator(ctx context.Context, config *config.Config, logger logging.Logger) (*Aggregator, error) {
- // TODO: Pass the registry as a parameter (see https://github.com/NethermindEth/near-sffl/pull/211#pullrequestreview-2101946551)
- registry := prometheus.NewRegistry()
-
+func NewAggregator(
+ // TODO: Remove `ctx` once OperatorsInfoServiceInMemory's API is changed and we can gracefully exit otherwise
+ ctx context.Context,
+ config *config.Config,
+ registry *prometheus.Registry,
+ logger logging.Logger,
+) (*Aggregator, error) {
  ethHttpClient, err := core.CreateEthClientWithCollector(AggregatorNamespace, config.EthHttpRpcUrl, config.EnableMetrics, registry, logger)
  if err != nil {
  logger.Error("Cannot create http ethclient", "err", err)
@@ -615,11 +616,6 @@ func (agg *Aggregator) ProcessSignedOperatorSetUpdateMessage(signedOperatorSetUp
  return err
 }
 
-// May return nil
-func (agg *Aggregator) GetRegistry() *prometheus.Registry {
- return agg.registry
-}
-
 func (agg *Aggregator) GetAggregatedCheckpointMessages(fromTimestamp, toTimestamp uint64) (*messages.CheckpointMessages, error) {
  checkpointMessages, err := agg.msgDb.FetchCheckpointMessages(fromTimestamp, toTimestamp)
  if err != nil {

diff --git a/aggregator/cmd/main.go b/aggregator/cmd/main.go
@@ -8,6 +8,7 @@ import (
  "os"
 
  sdklogging "github.com/Layr-Labs/eigensdk-go/logging"
+ "github.com/prometheus/client_golang/prometheus"
  "github.com/urfave/cli"
 
  "github.com/NethermindEth/near-sffl/aggregator"
@@ -66,23 +67,26 @@ func aggregatorMain(ctx *cli.Context) error {
  }
 
  bgCtx := context.Background()
- agg, err := aggregator.NewAggregator(bgCtx, config, logger)
+ var optRegistry *prometheus.Registry
+ if config.EnableMetrics {
+ optRegistry = prometheus.NewRegistry()
+ }
+ agg, err := aggregator.NewAggregator(bgCtx, config, optRegistry, logger)
  if err != nil {
  return err
  }
 
- registry := agg.GetRegistry()
  rpcServer := rpcserver.NewRpcServer(config.AggregatorServerIpPortAddr, agg, logger)
- if registry != nil {
- if err = rpcServer.EnableMetrics(registry); err != nil {
+ if optRegistry != nil {
+ if err = rpcServer.EnableMetrics(optRegistry); err != nil {
  return err
  }
  }
  go rpcServer.Start()
 
  restServer := restserver.NewRestServer(config.AggregatorRestServerIpPortAddr, agg, logger)
- if registry != nil {
- if err = restServer.EnableMetrics(registry); err != nil {
+ if optRegistry != nil {
+ if err = restServer.EnableMetrics(optRegistry); err != nil {
  return err
  }
  }

diff --git a/aggregator/rpc_server/server.go b/aggregator/rpc_server/server.go
@@ -3,12 +3,13 @@ package rpc_server
 import (
  "errors"
  "fmt"
- "github.com/NethermindEth/near-sffl/core"
- "github.com/prometheus/client_golang/prometheus"
  "net/http"
  "net/rpc"
  "strings"
 
+ "github.com/NethermindEth/near-sffl/core"
+ "github.com/prometheus/client_golang/prometheus"
+
  "github.com/Layr-Labs/eigensdk-go/logging"
  eigentypes "github.com/Layr-Labs/eigensdk-go/types"
 

diff --git a/tests/integration/integration_test.go b/tests/integration/integration_test.go
@@ -25,6 +25,7 @@ import (
  "github.com/ethereum/go-ethereum/accounts/abi/bind"
  "github.com/ethereum/go-ethereum/common"
  "github.com/ethereum/go-ethereum/crypto"
+ "github.com/prometheus/client_golang/prometheus"
  "github.com/stretchr/testify/assert"
  "github.com/testcontainers/testcontainers-go"
  "github.com/testcontainers/testcontainers-go/modules/rabbitmq"
@@ -292,24 +293,27 @@ func startOperator(t *testing.T, ctx context.Context, nodeConfig optypes.NodeCon
 func startAggregator(t *testing.T, ctx context.Context, config *config.Config, logger sdklogging.Logger) *aggregator.Aggregator {
  t.Log("starting aggregator for integration tests")
 
- agg, err := aggregator.NewAggregator(ctx, config, logger)
+ var optRegistry *prometheus.Registry
+ if config.EnableMetrics {
+ optRegistry = prometheus.NewRegistry()
+ }
+ agg, err := aggregator.NewAggregator(ctx, config, nil, logger)
  if err != nil {
  t.Fatalf("Failed to create aggregator: %s", err.Error())
  }
 
- registry := agg.GetRegistry()
  rpcServer := rpcserver.NewRpcServer(config.AggregatorServerIpPortAddr, agg, logger)
- if registry != nil {
- err = rpcServer.EnableMetrics(registry)
+ if optRegistry != nil {
+ err = rpcServer.EnableMetrics(optRegistry)
  if err != nil {
  t.Fatalf("Failed to create metrics for rpc server: %s", err.Error())
  }
  }
  go rpcServer.Start()
 
  restServer := restserver.NewRestServer(config.AggregatorRestServerIpPortAddr, agg, logger)
- if registry != nil {
- err = restServer.EnableMetrics(registry)
+ if optRegistry != nil {
+ err = restServer.EnableMetrics(optRegistry)
  if err != nil {
  t.Fatalf("Failed to create metrics for rest server: %s", err.Error())
  }