15
15
import java .time .Duration ;
16
16
import java .util .concurrent .CountDownLatch ;
17
17
import java .util .concurrent .TimeUnit ;
18
+ import java .util .function .Consumer ;
18
19
import org .junit .Test ;
19
20
20
21
/** Leader Election tests using "simulated" locks created by {@link LockSmith} */
@@ -53,11 +54,64 @@ public void testLeaderGracefulShutdown() throws Exception {
53
54
leaderElector2 .close ();
54
55
}
55
56
57
+ @ Test (timeout = 20000L )
58
+ public void testLeaderTransitionHook () throws InterruptedException {
59
+ LockSmith lockSmith = new LockSmith ();
60
+
61
+ CountDownLatch startBeingLeader1 = new CountDownLatch (1 );
62
+ CountDownLatch stopBeingLeader1 = new CountDownLatch (1 );
63
+
64
+ LeaderElector leaderElector1 =
65
+ makeAndRunLeaderElectorAsync (
66
+ lockSmith , "candidate1" , startBeingLeader1 , stopBeingLeader1 , (id ) -> {});
67
+
68
+ // wait for candidate1 to become leader
69
+ startBeingLeader1 .await ();
70
+
71
+ // start candidate2 and the transition hook should be called on the start
72
+ CountDownLatch startBeingLeader2 = new CountDownLatch (1 );
73
+ CountDownLatch stopBeingLeader2 = new CountDownLatch (1 );
74
+ CountDownLatch notifiedLeader = new CountDownLatch (1 );
75
+ String expectedLeader = "candidate1" ;
76
+ LeaderElector leaderElector2 =
77
+ makeAndRunLeaderElectorAsync (
78
+ lockSmith ,
79
+ "candidate2" ,
80
+ startBeingLeader2 ,
81
+ stopBeingLeader2 ,
82
+ (id ) -> {
83
+ if (expectedLeader .equals (id )) {
84
+ notifiedLeader .countDown ();
85
+ }
86
+ });
87
+
88
+ notifiedLeader .await ();
89
+
90
+ // start candidate1
91
+ leaderElector1 .close ();
92
+
93
+ // ensure stopBeingLeader hook is called
94
+ stopBeingLeader1 .await ();
95
+
96
+ // wait for candidate2 to become leader
97
+ startBeingLeader2 .await ();
98
+ }
99
+
56
100
private LeaderElector makeAndRunLeaderElectorAsync (
57
101
LockSmith lockSmith ,
58
102
String lockIdentity ,
59
103
CountDownLatch startBeingLeader ,
60
104
CountDownLatch stopBeingLeader ) {
105
+ return makeAndRunLeaderElectorAsync (
106
+ lockSmith , lockIdentity , startBeingLeader , stopBeingLeader , (id ) -> {});
107
+ }
108
+
109
+ private LeaderElector makeAndRunLeaderElectorAsync (
110
+ LockSmith lockSmith ,
111
+ String lockIdentity ,
112
+ CountDownLatch startBeingLeader ,
113
+ CountDownLatch stopBeingLeader ,
114
+ Consumer <String > onNewLeaderHook ) {
61
115
LeaderElectionConfig leaderElectionConfig =
62
116
new LeaderElectionConfig (
63
117
lockSmith .makeLock (lockIdentity ),
@@ -70,7 +124,9 @@ private LeaderElector makeAndRunLeaderElectorAsync(
70
124
new Thread (
71
125
() ->
72
126
leaderElector .run (
73
- () -> startBeingLeader .countDown (), () -> stopBeingLeader .countDown ()),
127
+ () -> startBeingLeader .countDown (),
128
+ () -> stopBeingLeader .countDown (),
129
+ onNewLeaderHook ),
74
130
String .format ("%s-leader-elector-main" , lockIdentity ));
75
131
thread .setDaemon (true );
76
132
thread .start ();
0 commit comments