From a5ce4b70c5fdff43f5b62ccb1d9d3ee053f66f6e Mon Sep 17 00:00:00 2001
From: Kacper Sawicki <kacper@ksawicki.pl>
Date: Mon, 22 Jan 2024 12:16:46 +0100
Subject: [PATCH] Optimize reward sync

---
 collector/collector.go | 1 +
 collector/mesh.go      | 2 ++
 storage/storage.go     | 7 +++++--
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/collector/collector.go b/collector/collector.go
index 5d3ff8d..848d311 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -44,6 +44,7 @@ type Listener interface {
 	GetEpochNumLayers() uint32
 	GetTransactions(parent context.Context, query *bson.D, opts ...*options.FindOptions) ([]model.Transaction, error)
 	UpdateTransactionState(parent context.Context, id string, state int32) error
+	UpdateEpochStats(layer uint32)
 }
 
 type Collector struct {
diff --git a/collector/mesh.go b/collector/mesh.go
index f6078b2..c53daa3 100644
--- a/collector/mesh.go
+++ b/collector/mesh.go
@@ -175,6 +175,8 @@ func (c *Collector) syncLayer(lid types.LayerID) error {
 		c.listener.OnReward(r)
 	}
 
+	c.listener.UpdateEpochStats(layer.Number.Number)
+
 	return nil
 }
 
diff --git a/storage/storage.go b/storage/storage.go
index 4a05a68..f910e21 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -252,8 +252,11 @@ func (s *Storage) OnReward(in *pb.Reward) {
 	}
 
 	s.requestBalanceUpdate(reward.Layer, reward.Coinbase)
-	s.setChangedEpoch(reward.Layer)
-	s.updateEpochs() // trigger epoch stat recalculation todo: optimize this
+}
+
+func (s *Storage) UpdateEpochStats(layer uint32) {
+	s.setChangedEpoch(layer)
+	s.updateEpochs()
 }
 
 func (s *Storage) OnTransactionResult(res *pb.TransactionResult, state *pb.TransactionState) {