diff --git a/simplex/replication_state.go b/simplex/replication_state.go index ed9e2a28..8460374b 100644 --- a/simplex/replication_state.go +++ b/simplex/replication_state.go @@ -6,6 +6,7 @@ package simplex import ( "math/rand/v2" "sync" + "sync/atomic" "time" "github.com/ava-labs/simplex/common" @@ -30,7 +31,7 @@ type finalizedQuorumRound struct { } type ReplicationState struct { - enabled bool + enabled atomic.Bool logger common.Logger myNodeID common.NodeID rand *rand.Rand // Random number generator @@ -63,14 +64,12 @@ type ReplicationState struct { func NewReplicationState(logger common.Logger, comm Sender, myNodeID common.NodeID, maxRoundWindow uint64, enabled bool, start time.Time, lock *sync.Mutex, rng *rand.Rand) *ReplicationState { if !enabled { return &ReplicationState{ - enabled: enabled, - logger: logger, - rand: rng, + logger: logger, + rand: rng, } } r := &ReplicationState{ - enabled: enabled, myNodeID: myNodeID, logger: logger, rand: rng, @@ -86,6 +85,7 @@ func NewReplicationState(logger common.Logger, comm Sender, myNodeID common.Node sender: comm, epochLock: lock, } + r.enabled.Store(true) r.digestTimeouts = common.NewTimeoutHandler(logger, "digest", start, DefaultReplicationRequestTimeout, r.requestDigests) r.emptyRoundTimeouts = common.NewTimeoutHandler(logger, "empty round replication", start, DefaultReplicationRequestTimeout, r.requestEmptyRounds) @@ -94,7 +94,7 @@ func NewReplicationState(logger common.Logger, comm Sender, myNodeID common.Node } func (r *ReplicationState) AdvanceTime(now time.Time) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -161,7 +161,7 @@ func (r *ReplicationState) storeRound(qr *common.QuorumRound) { // StoreQuorumRound stores the quorum round into the replication state. func (r *ReplicationState) StoreQuorumRound(round *common.QuorumRound) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -193,7 +193,7 @@ func (r *ReplicationState) StoreQuorumRound(round *common.QuorumRound) { // receivedFutureFinalization notifies the replication state a finalization was created in a future round. func (r *ReplicationState) ReceivedFutureFinalization(finalization *common.Finalization, nextSeqToCommit uint64) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -209,7 +209,7 @@ func (r *ReplicationState) ReceivedFutureFinalization(finalization *common.Final // receivedFutureRound notifies the replication state of a future round. func (r *ReplicationState) ReceivedFutureRound(round, seq, currentRound uint64, signers []common.NodeID) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -224,7 +224,7 @@ func (r *ReplicationState) ReceivedFutureRound(round, seq, currentRound uint64, // ResendFinalizationRequest notifies the replication state that `seq` should be re-requested. func (r *ReplicationState) ResendFinalizationRequest(seq uint64, signers []common.NodeID) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -253,7 +253,7 @@ func (r *ReplicationState) CreateDependencyTasks(parent *common.Digest, parentSe } func (r *ReplicationState) clearBlockDependencyTasks(digest common.Digest, seq uint64, finalizationPersisted bool) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -275,7 +275,7 @@ func (r *ReplicationState) clearBlockDependencyTasks(digest common.Digest, seq u // MaybeAdvanceState attempts to collect future sequences if // there are more to be collected and the round has caught up for us to send the request. func (r *ReplicationState) MaybeAdvanceState(nextSequenceToCommit uint64, currentRound uint64, lastCommittedRound uint64) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -398,7 +398,7 @@ func (r *ReplicationState) requestEmptyRounds(emptyRounds []uint64) { } func (r *ReplicationState) DeleteRound(round uint64) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -409,7 +409,7 @@ func (r *ReplicationState) DeleteRound(round uint64) { } func (r *ReplicationState) DeleteSeq(seq uint64) { - if !r.enabled { + if !r.enabled.Load() { return } @@ -417,7 +417,7 @@ func (r *ReplicationState) DeleteSeq(seq uint64) { } func (r *ReplicationState) Close() { - if !r.enabled { + if !r.enabled.Load() { return } @@ -425,4 +425,5 @@ func (r *ReplicationState) Close() { r.emptyRoundTimeouts.Close() r.roundRequestor.close() r.finalizationRequestor.close() + r.enabled.Store(false) }