From 9f8ccc84d3bad384d4949c21cc522bde8640e149 Mon Sep 17 00:00:00 2001
From: Sam Calder-Mason <sam@puritydev.io>
Date: Tue, 2 Apr 2024 12:47:55 +1000
Subject: [PATCH] fix(sentry): Fetch beacon committees just before epoch
 transition

---
 pkg/sentry/beacon_committees.go        |  1 -
 pkg/sentry/ethereum/services/duties.go | 36 +++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/pkg/sentry/beacon_committees.go b/pkg/sentry/beacon_committees.go
index eba37b00..01e5a097 100644
--- a/pkg/sentry/beacon_committees.go
+++ b/pkg/sentry/beacon_committees.go
@@ -33,7 +33,6 @@ func (s *Sentry) startBeaconCommitteesWatcher(ctx context.Context) error {
 		// Grab the current epoch (+1) committees.
 		now := s.beacon.Metadata().Wallclock().Epochs().Current()
 		epochs := []phase0.Epoch{
-			phase0.Epoch(now.Number() + 1),
 			phase0.Epoch(now.Number()),
 		}
 
diff --git a/pkg/sentry/ethereum/services/duties.go b/pkg/sentry/ethereum/services/duties.go
index bb950dfd..6521be03 100644
--- a/pkg/sentry/ethereum/services/duties.go
+++ b/pkg/sentry/ethereum/services/duties.go
@@ -83,20 +83,39 @@ func (m *DutiesService) Start(ctx context.Context) error {
 	}()
 
 	m.metadata.Wallclock().OnEpochChanged(func(epoch ethwallclock.Epoch) {
-		if err := m.fetchRequiredEpochDuties(ctx, true); err != nil {
-			m.log.WithError(err).Warn("Failed to fetch required epoch duties")
-		}
-
 		// Sleep for a bit to give the beacon node a chance to run its epoch transition.
 		// We don't really care about nice-to-have duties so the sleep here is fine.
 		// "Required" duties (aka the current epoch) will be refetched the moment that epoch
 		// starts.
+
+		time.Sleep(500 * time.Millisecond)
+
+		if err := m.fetchRequiredEpochDuties(ctx, true); err != nil {
+			m.log.WithError(err).Warn("Failed to fetch required epoch duties after an epoch change")
+		}
+
 		time.Sleep(15 * time.Second)
 
 		//nolint:errcheck // We don't care about the error here
 		m.fetchNiceToHaveEpochDuties(ctx)
 	})
 
+	m.metadata.Wallclock().OnEpochChanged(func(epoch ethwallclock.Epoch) {
+		m.log.
+			WithField("current_epoch", epoch.Number()).
+			WithField("next_epoch", epoch.Number()+1).
+			Debug("Fetching beacon committees for next epoch")
+		// Sleep until just before the start of the next epoch to fetch the next epoch's duties.
+		time.Sleep(epoch.TimeWindow().EndsIn() - 2*time.Second)
+
+		if err := m.fetchBeaconCommittee(ctx, phase0.Epoch(epoch.Number()+1), true); err != nil {
+			m.log.WithError(err).Warn("Failed to fetch required epoch duties in anticipation of an epoch change")
+		}
+
+		//nolint:errcheck // We don't care about the error here
+		m.fetchNiceToHaveEpochDuties(ctx)
+	})
+
 	m.beacon.OnChainReOrg(ctx, func(ctx context.Context, ev *v1.ChainReorgEvent) error {
 		m.log.Info("Chain reorg detected - refetching beacon committees")
 
@@ -168,7 +187,6 @@ func (m *DutiesService) NiceToHaveEpochDuties(ctx context.Context) []phase0.Epoc
 
 	epochs := []phase0.Epoch{
 		phase0.Epoch(epochNumber - 1),
-		phase0.Epoch(epochNumber + 1),
 	}
 
 	final := map[phase0.Epoch]struct{}{}
@@ -246,7 +264,13 @@ func (m *DutiesService) fetchBeaconCommittee(ctx context.Context, epoch phase0.E
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	m.log.WithField("epoch", epoch).WithField("override_cache", overrideCache).Debug("Fetching beacon committee")
+	wallclockEpoch := m.metadata.Wallclock().Epochs().Current()
+
+	m.log.
+		WithField("epoch", epoch).
+		WithField("override_cache", overrideCache).
+		WithField("wallclock_epoch", wallclockEpoch).
+		Debug("Fetching beacon committee")
 
 	committees, err := m.beacon.FetchBeaconCommittees(ctx, "head", epoch)
 	if err != nil {