Disable replication on Close

simplex/replication_state.go

@@ -6,6 +6,7 @@
 import (
 	"math/rand/v2"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/ava-labs/simplex/common"
@@ -30,7 +31,7 @@
 }
 
 type ReplicationState struct {
-	enabled  bool
+	enabled  atomic.Bool
 	logger   common.Logger
 	myNodeID common.NodeID
 	rand     *rand.Rand // Random number generator
@@ -63,14 +64,12 @@
 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 @@
 		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 (r *ReplicationState) AdvanceTime(now time.Time) {
-	if !r.enabled {
+	if !r.enabled.Load() {
 		return
 	}
 
@@ -173,7 +173,7 @@
 
 // StoreQuorumRound stores the quorum round into the replication state.
 func (r *ReplicationState) StoreQuorumRound(round *common.QuorumRound) {
-	if !r.enabled {
+	if !r.enabled.Load() {
 		return
 	}
 
@@ -205,7 +205,7 @@
 
 // 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
 	}
 
@@ -221,7 +221,7 @@
 
 // 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
 	}
 
@@ -236,7 +236,7 @@
 
 // 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
 	}
 
@@ -252,7 +252,7 @@
 
 // ResendRoundRequest notifies the replication state that `round` should be re-requested.
 func (r *ReplicationState) ResendRoundRequest(round uint64, signers []common.NodeID) {
 	if !r.enabled {
 		return
 	}
 
@@ -279,7 +279,7 @@
 }
 
 func (r *ReplicationState) clearBlockDependencyTasks(digest common.Digest, seq uint64, finalizationPersisted bool) {
-	if !r.enabled {
+	if !r.enabled.Load() {
 		return
 	}
 
@@ -301,7 +301,7 @@
 // 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
 	}
 
@@ -426,7 +426,7 @@
 }
 
 func (r *ReplicationState) DeleteRound(round uint64) {
-	if !r.enabled {
+	if !r.enabled.Load() {
 		return
 	}
 
@@ -437,20 +437,21 @@
 }
 
 func (r *ReplicationState) DeleteSeq(seq uint64) {
-	if !r.enabled {
+	if !r.enabled.Load() {
 		return
 	}
 
 	delete(r.seqs, seq)
 }
 
 func (r *ReplicationState) Close() {
-	if !r.enabled {
+	if !r.enabled.Load() {
 		return
 	}
 
 	r.digestTimeouts.Close()
 	r.emptyRoundTimeouts.Close()
 	r.roundRequestor.close()
 	r.finalizationRequestor.close()
+	r.enabled.Store(false)
 }