Pass Epoch Lock To NotarizationTime

epoch.go

@@ -214,23 +214,20 @@ func (e *Epoch) init() error {
 
 func (e *Epoch) initOldestNotFinalizedNotarization() {
 	rebroadcastFinalizationVotes := func() {
-		e.lock.Lock()
-		defer e.lock.Unlock()
-
 		if err := e.rebroadcastPastFinalizeVotes(); err != nil {
 			e.Logger.Error("Could not rebroadcast past finalization votes", zap.Error(err))
 		}
 	}
 	e.oldestNotFinalizedNotarization = NewNotarizationTime(
 		e.FinalizeRebroadcastTimeout,
 		e.haveNotFinalizedNotarizedRound,
-		rebroadcastFinalizationVotes, e.getRound)
+		rebroadcastFinalizationVotes,
+		e.getRound,
+		&e.lock,
+	)
 }
 
 func (e *Epoch) getRound() uint64 {
-	e.lock.Lock()
-	defer e.lock.Unlock()
-
 	return e.round
 }
 
@@ -2869,9 +2866,6 @@ func (e *Epoch) locateQuorumRecordByRound(targetRound uint64) *VerifiedQuorumRou
 }
 
 func (e *Epoch) haveNotFinalizedNotarizedRound() (uint64, bool) {
-	e.lock.Lock()
-	defer e.lock.Unlock()
-
 	var minRoundNum uint64
 	var found bool
 	for _, round := range e.rounds {

util.go

@@ -279,24 +279,32 @@ type NotarizationTime struct {
 	latestRound             uint64
 	lastRebroadcastTime     time.Time
 	oldestNotFinalizedRound uint64
+
+	// epoch lock
+	lock *sync.Mutex
 }
 
 func NewNotarizationTime(
 	finalizeVoteRebroadcastTimeout time.Duration,
 	haveUnFinalizedNotarization func() (uint64, bool),
 	rebroadcastFinalizationVotes func(),
 	getRound func() uint64,
+	lock *sync.Mutex,
 ) NotarizationTime {
 	return NotarizationTime{
 		finalizeVoteRebroadcastTimeout: finalizeVoteRebroadcastTimeout,
 		haveUnFinalizedNotarization:    haveUnFinalizedNotarization,
 		rebroadcastFinalizationVotes:   rebroadcastFinalizationVotes,
 		getRound:                       getRound,
 		checkInterval:                  finalizeVoteRebroadcastTimeout / 3,
+		lock:                           lock,
 	}
 }
 
 func (nt *NotarizationTime) CheckForNotFinalizedNotarizedBlocks(now time.Time) {
+	nt.lock.Lock()
+	defer nt.lock.Unlock()
+
 	// If we have recently checked, don't check again
 	if !nt.lastSampleTime.IsZero() && nt.lastSampleTime.Add(nt.checkInterval).After(now) {
 		return
@@ -305,7 +313,6 @@ func (nt *NotarizationTime) CheckForNotFinalizedNotarizedBlocks(now time.Time) {
 	nt.lastSampleTime = now
 
 	round := nt.getRound()
-
 	// As long as we make some progress, we don't check for a round not finalized.
 	if round > nt.latestRound {
 		nt.latestRound = round

util_test.go

@@ -6,6 +6,7 @@ package simplex_test
 import (
 	"context"
 	"fmt"
+	"sync"
 	"testing"
 	"time"
 
@@ -400,13 +401,16 @@ func TestNotarizationTime(t *testing.T) {
 	rebroadcastFinalizationVotes := func() {
 		invoked++
 	}
+	lock := &sync.Mutex{}
 	nt := NewNotarizationTime(
 		defaultFinalizeVoteRebroadcastTimeout,
 		haveNotFinalizedRound,
 		rebroadcastFinalizationVotes,
 		func() uint64 {
 			return round
-		})
+		},
+		lock,
+	)
 
 	// First call should set the time and the round.
 	have = true