feat: add clearOnStageTimerHandler method to GrandpaRound

LimeChain · Feb 7, 2025 · 234cf35 · 234cf35
1 parent 9f17578
commit 234cf35
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/src/main/java/com/limechain/grandpa/round/GrandpaRound.java b/src/main/java/com/limechain/grandpa/round/GrandpaRound.java
@@ -657,4 +657,11 @@ public void broadcastCommitMessage() {
 
         peerMessageCoordinator.sendCommitMessageToPeers(commitMessage);
     }
+
+    public void clearOnStageTimerHandler() {
+        if (onStageTimerHandler != null && onStageTimerHandler.isShutdown()) {
+            onStageTimerHandler.shutdown();
+            onStageTimerHandler = null;
+        }
+    }
 }
diff --git a/src/main/java/com/limechain/grandpa/round/PreCommitStage.java b/src/main/java/com/limechain/grandpa/round/PreCommitStage.java
@@ -32,20 +32,16 @@ public void start(GrandpaRound round) {
 
     @Override
     public void end(GrandpaRound round) {
-        log.fine(String.format("Round %d ended pre-commit stage.", round.getRoundNumber()));
 
-        if (round.getOnStageTimerHandler() != null && !round.getOnStageTimerHandler().isShutdown()) {
-            round.getOnStageTimerHandler().shutdown();
-            round.setOnStageTimerHandler(null);
-        }
+        round.clearOnStageTimerHandler();
 
         try {
 
             Vote grandpaGhost = Vote.fromBlockHeader(round.getGrandpaGhost());
             round.broadcastVoteMessage(grandpaGhost, SubRound.PRE_COMMIT);
             round.setState(new FinalizeStage());
             round.getState().start(round);
-            log.fine(String.format("Round %d ended start stage.", round.getRoundNumber()));
+            log.fine(String.format("Round %d ended pre-commit stage.", round.getRoundNumber()));
 
         } catch (GrandpaGenericException e) {
             log.fine(String.format("Round %d cannot end now: %s", round.getRoundNumber(), e.getMessage()));