GROW-705 - Decouple EpochUpdater from ProviderStateQuery (#445)

lavanet · Apr 27, 2023 · 7aac888 · 7aac888
1 parent 43c7af7
commit 7aac888
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 13 deletions.
diff --git a/protocol/statetracker/epoch_updater.go b/protocol/statetracker/epoch_updater.go
@@ -15,10 +15,10 @@ type EpochUpdatable interface {
 type EpochUpdater struct {
 	epochUpdatables []*EpochUpdatable
 	currentEpoch    uint64
-	stateQuery      *ProviderStateQuery
+	stateQuery      *EpochStateQuery
 }
 
-func NewEpochUpdater(stateQuery *ProviderStateQuery) *EpochUpdater {
+func NewEpochUpdater(stateQuery *EpochStateQuery) *EpochUpdater {
 	return &EpochUpdater{epochUpdatables: []*EpochUpdatable{}, stateQuery: stateQuery}
 }
 

diff --git a/protocol/statetracker/provider_state_tracker.go b/protocol/statetracker/provider_state_tracker.go
@@ -35,7 +35,7 @@ func NewProviderStateTracker(ctx context.Context, txFactory tx.Factory, clientCt
 }
 
 func (pst *ProviderStateTracker) RegisterForEpochUpdates(ctx context.Context, epochUpdatable EpochUpdatable) {
-	epochUpdater := NewEpochUpdater(pst.stateQuery)
+	epochUpdater := NewEpochUpdater(&pst.stateQuery.EpochStateQuery)
 	epochUpdaterRaw := pst.StateTracker.RegisterForUpdates(ctx, epochUpdater)
 	epochUpdater, ok := epochUpdaterRaw.(*EpochUpdater)
 	if !ok {

diff --git a/protocol/statetracker/state_query.go b/protocol/statetracker/state_query.go
@@ -110,13 +110,33 @@ func (csq *ConsumerStateQuery) GetMaxCUForUser(ctx context.Context, chainID stri
 	return UserEntryRes.GetMaxCU(), nil
 }
 
+type EpochStateQuery struct {
+	StateQuery
+}
+
+func (esq *EpochStateQuery) CurrentEpochStart(ctx context.Context) (uint64, error) {
+	epochDetails, err := esq.EpochStorageQueryClient.EpochDetails(ctx, &epochstoragetypes.QueryGetEpochDetailsRequest{})
+	if err != nil {
+		return 0, utils.LavaFormatError("Failed Querying EpochDetails", err)
+	}
+	details := epochDetails.GetEpochDetails()
+	return details.StartBlock, nil
+}
+
+func NewEpochStateQuery(stateQuery *StateQuery) *EpochStateQuery {
+	return &EpochStateQuery{StateQuery: *stateQuery}
+}
+
 type ProviderStateQuery struct {
 	StateQuery
+	EpochStateQuery
 	clientCtx client.Context
 }
 
 func NewProviderStateQuery(ctx context.Context, clientCtx client.Context) *ProviderStateQuery {
-	csq := &ProviderStateQuery{StateQuery: *NewStateQuery(ctx, clientCtx), clientCtx: clientCtx}
+	sq := NewStateQuery(ctx, clientCtx)
+	esq := NewEpochStateQuery(sq)
+	csq := &ProviderStateQuery{StateQuery: *sq, EpochStateQuery: *esq, clientCtx: clientCtx}
 	return csq
 }
 
@@ -150,15 +170,6 @@ func (psq *ProviderStateQuery) entryKey(consumerAddress string, chainID string,
 	return consumerAddress + chainID + strconv.FormatUint(epoch, 10) + providerAddress
 }
 
-func (psq *ProviderStateQuery) CurrentEpochStart(ctx context.Context) (uint64, error) {
-	epochDetails, err := psq.EpochStorageQueryClient.EpochDetails(ctx, &epochstoragetypes.QueryGetEpochDetailsRequest{})
-	if err != nil {
-		return 0, utils.LavaFormatError("Failed Querying EpochDetails", err)
-	}
-	details := epochDetails.GetEpochDetails()
-	return details.StartBlock, nil
-}
-
 func (psq *ProviderStateQuery) PaymentEvents(ctx context.Context, latestBlock int64) (payments []*rewardserver.PaymentRequest, err error) {
 	blockResults, err := psq.clientCtx.Client.BlockResults(ctx, &latestBlock)
 	if err != nil {