fix: Remove ctx from grpc client struct

archway-network · Dec 20, 2023 · 053417d · 053417d
1 parent 8b315d4
commit 053417d
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 24 deletions.
diff --git a/cmd/validator-exporter/main.go b/cmd/validator-exporter/main.go
@@ -1,9 +1,11 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/caarlos0/env/v10"
 	"github.com/prometheus/client_golang/prometheus"
@@ -28,7 +30,14 @@ func main() {
 		log.Fatal(err.Error())
 	}
 
-	_, err := grpc.LatestBlockHeight(cfg)
+	ctx, cancel := context.WithTimeout(
+		context.Background(),
+		time.Duration(cfg.Timeout)*time.Second,
+	)
+
+	defer cancel()
+
+	_, err := grpc.LatestBlockHeight(ctx, cfg)
 	if err != nil {
 		log.Fatal(err.Error())
 	}

diff --git a/pkg/collector/collector.go b/pkg/collector/collector.go
@@ -1,7 +1,9 @@
 package collector
 
 import (
+	"context"
 	"fmt"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
@@ -36,7 +38,14 @@ func (vc ValidatorsCollector) Describe(ch chan<- *prometheus.Desc) {
 }
 
 func (vc ValidatorsCollector) Collect(ch chan<- prometheus.Metric) {
-	vals, err := grpc.SigningValidators(vc.Cfg)
+	ctx, cancel := context.WithTimeout(
+		context.Background(),
+		time.Duration(vc.Cfg.Timeout)*time.Second,
+	)
+
+	defer cancel()
+
+	vals, err := grpc.SigningValidators(ctx, vc.Cfg)
 	if err != nil {
 		log.Error(fmt.Sprintf("error getting signing validators: %s", err))
 	} else {

diff --git a/pkg/grpc/grpc.go b/pkg/grpc/grpc.go
@@ -3,7 +3,6 @@ package grpc
 import (
 	"context"
 	"fmt"
-	"time"
 
 	base "cosmossdk.io/api/cosmos/base/tendermint/v1beta1"
 
@@ -23,8 +22,6 @@ import (
 const valConsStr = "valcons"
 
 type Client struct {
-	ctx       context.Context
-	ctxCancel context.CancelFunc
 	cfg       config.Config
 	conn      *grpc.ClientConn
 	connClose func()
@@ -35,14 +32,6 @@ func NewClient(cfg config.Config) (Client, error) {
 		cfg: cfg,
 	}
 
-	ctx, cancel := context.WithTimeout(
-		context.Background(),
-		time.Duration(client.cfg.Timeout)*time.Second,
-	)
-
-	client.ctx = ctx
-	client.ctxCancel = cancel
-
 	conn, err := cfg.GRPCConn()
 	if err != nil {
 		return Client{}, err
@@ -58,15 +47,15 @@ func NewClient(cfg config.Config) (Client, error) {
 	return client, nil
 }
 
-func (c Client) SignigInfos() ([]slashing.ValidatorSigningInfo, error) {
+func (c Client) SignigInfos(ctx context.Context) ([]slashing.ValidatorSigningInfo, error) {
 	infos := []slashing.ValidatorSigningInfo{}
 	key := []byte{}
 	client := slashing.NewQueryClient(c.conn)
 
 	for {
 		request := &slashing.QuerySigningInfosRequest{Pagination: &query.PageRequest{Key: key}}
 
-		slashRes, err := client.SigningInfos(c.ctx, request)
+		slashRes, err := client.SigningInfos(ctx, request)
 		if err != nil {
 			return nil, err
 		}
@@ -93,7 +82,7 @@ func (c Client) SignigInfos() ([]slashing.ValidatorSigningInfo, error) {
 	return infos, nil
 }
 
-func (c Client) Validators() ([]staking.Validator, error) {
+func (c Client) Validators(ctx context.Context) ([]staking.Validator, error) {
 	vals := []staking.Validator{}
 	key := []byte{}
 
@@ -106,7 +95,7 @@ func (c Client) Validators() ([]staking.Validator, error) {
 	for {
 		request := &staking.QueryValidatorsRequest{Pagination: &query.PageRequest{Key: key}}
 
-		stakingRes, err := client.Validators(c.ctx, request)
+		stakingRes, err := client.Validators(ctx, request)
 		if err != nil {
 			return nil, err
 		}
@@ -160,7 +149,7 @@ func (c Client) valConsMap(vals []staking.Validator) (map[string]staking.Validat
 	return vMap, nil
 }
 
-func SigningValidators(cfg config.Config) ([]types.Validator, error) {
+func SigningValidators(ctx context.Context, cfg config.Config) ([]types.Validator, error) {
 	sVals := []types.Validator{}
 
 	client, err := NewClient(cfg)
@@ -171,16 +160,15 @@ func SigningValidators(cfg config.Config) ([]types.Validator, error) {
 	}
 
 	defer client.connClose()
-	defer client.ctxCancel()
 
-	sInfos, err := client.SignigInfos()
+	sInfos, err := client.SignigInfos(ctx)
 	if err != nil {
 		log.Error(err.Error())
 
 		return []types.Validator{}, err
 	}
 
-	vals, err := client.Validators()
+	vals, err := client.Validators(ctx)
 	if err != nil {
 		log.Error(err.Error())
 
@@ -210,7 +198,7 @@ func SigningValidators(cfg config.Config) ([]types.Validator, error) {
 	return sVals, nil
 }
 
-func LatestBlockHeight(cfg config.Config) (int64, error) {
+func LatestBlockHeight(ctx context.Context, cfg config.Config) (int64, error) {
 	client, err := NewClient(cfg)
 	if err != nil {
 		log.Error(err.Error())
@@ -219,12 +207,11 @@ func LatestBlockHeight(cfg config.Config) (int64, error) {
 	}
 
 	defer client.connClose()
-	defer client.ctxCancel()
 
 	request := &base.GetLatestBlockRequest{}
 	baseClient := base.NewServiceClient(client.conn)
 
-	blockResp, err := baseClient.GetLatestBlock(client.ctx, request)
+	blockResp, err := baseClient.GetLatestBlock(ctx, request)
 	if err != nil {
 		log.Error(err.Error())