Java tutorial
// ================================================================================================= // Copyright 2011 Twitter, Inc. // ------------------------------------------------------------------------------------------------- // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this work except in compliance with the License. // You may obtain a copy of the License in the LICENSE file, or at: // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // ================================================================================================= package com.twitter.common.zookeeper; import java.io.IOException; import java.net.InetSocketAddress; import java.util.List; import java.util.Map; import java.util.concurrent.CountDownLatch; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.testing.TearDown; import org.apache.zookeeper.CreateMode; import org.apache.zookeeper.WatchedEvent; import org.apache.zookeeper.Watcher; import org.apache.zookeeper.Watcher.Event.EventType; import org.apache.zookeeper.Watcher.Event.KeeperState; import org.easymock.Capture; import org.easymock.EasyMock; import org.easymock.IAnswer; import org.easymock.IExpectationSetters; import org.easymock.IMocksControl; import org.junit.Before; import org.junit.Test; import com.twitter.common.base.ExceptionalCommand; import com.twitter.common.zookeeper.Candidate.Leader; import com.twitter.common.zookeeper.Group.JoinException; import com.twitter.common.zookeeper.ServerSet.EndpointStatus; import com.twitter.common.zookeeper.SingletonService.DefeatOnDisconnectLeader; import com.twitter.common.zookeeper.SingletonService.LeaderControl; import com.twitter.common.zookeeper.SingletonService.LeadershipListener; import com.twitter.common.zookeeper.testing.BaseZooKeeperTest; import static com.twitter.common.testing.easymock.EasyMockTest.createCapture; import static org.easymock.EasyMock.capture; import static org.easymock.EasyMock.createControl; import static org.easymock.EasyMock.expect; import static org.easymock.EasyMock.expectLastCall; import static org.junit.Assert.fail; public class SingletonServiceTest extends BaseZooKeeperTest { private static final int PORT_A = 1234; private static final int PORT_B = 8080; private static final InetSocketAddress PRIMARY_ENDPOINT = InetSocketAddress.createUnresolved("foo", PORT_A); private static final Map<String, InetSocketAddress> AUX_ENDPOINTS = ImmutableMap.of("http-admin", InetSocketAddress.createUnresolved("foo", PORT_B)); private IMocksControl control; private SingletonService.LeadershipListener listener; private ServerSet serverSet; private ServerSet.EndpointStatus endpointStatus; private Candidate candidate; private ExceptionalCommand<Group.JoinException> abdicate; private SingletonService service; @Before @SuppressWarnings("unchecked") public void mySetUp() throws IOException { control = createControl(); addTearDown(new TearDown() { @Override public void tearDown() { control.verify(); } }); listener = control.createMock(SingletonService.LeadershipListener.class); serverSet = control.createMock(ServerSet.class); candidate = control.createMock(Candidate.class); endpointStatus = control.createMock(ServerSet.EndpointStatus.class); abdicate = control.createMock(ExceptionalCommand.class); service = new SingletonService(serverSet, candidate); } private void newLeader(final String hostName, Capture<Leader> leader, LeadershipListener listener) throws Exception { service.lead(InetSocketAddress.createUnresolved(hostName, PORT_A), ImmutableMap.of("http-admin", InetSocketAddress.createUnresolved(hostName, PORT_B)), listener); // This actually elects the leader. leader.getValue().onElected(abdicate); } private void newLeader(String hostName, Capture<Leader> leader) throws Exception { newLeader(hostName, leader, listener); } private IExpectationSetters<EndpointStatus> expectJoin() throws Exception { return expect(serverSet.join(PRIMARY_ENDPOINT, AUX_ENDPOINTS)); } @Test public void testLeadAdvertise() throws Exception { Capture<Leader> leaderCapture = createCapture(); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); Capture<LeaderControl> controlCapture = createCapture(); listener.onLeading(capture(controlCapture)); expectJoin().andReturn(endpointStatus); endpointStatus.leave(); abdicate.execute(); control.replay(); newLeader("foo", leaderCapture); controlCapture.getValue().advertise(); controlCapture.getValue().leave(); } @Test public void teatLeadLeaveNoAdvertise() throws Exception { Capture<Leader> leaderCapture = createCapture(); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); abdicate.execute(); Capture<LeaderControl> controlCapture = createCapture(); listener.onLeading(capture(controlCapture)); control.replay(); newLeader("foo", leaderCapture); controlCapture.getValue().leave(); } @Test public void testLeadJoinFailure() throws Exception { Capture<Leader> leaderCapture = new Capture<Leader>(); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); Capture<LeaderControl> controlCapture = createCapture(); listener.onLeading(capture(controlCapture)); expectJoin().andThrow(new Group.JoinException("Injected join failure.", new Exception())); abdicate.execute(); control.replay(); newLeader("foo", leaderCapture); try { controlCapture.getValue().advertise(); fail("Join should have failed."); } catch (JoinException e) { // Expected. } controlCapture.getValue().leave(); } @Test(expected = IllegalStateException.class) public void testMultipleAdvertise() throws Exception { Capture<Leader> leaderCapture = createCapture(); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); Capture<LeaderControl> controlCapture = createCapture(); listener.onLeading(capture(controlCapture)); expectJoin().andReturn(endpointStatus); control.replay(); newLeader("foo", leaderCapture); controlCapture.getValue().advertise(); controlCapture.getValue().advertise(); } @Test(expected = IllegalStateException.class) public void testMultipleLeave() throws Exception { Capture<Leader> leaderCapture = createCapture(); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); Capture<LeaderControl> controlCapture = createCapture(); listener.onLeading(capture(controlCapture)); expectJoin().andReturn(endpointStatus); endpointStatus.leave(); abdicate.execute(); control.replay(); newLeader("foo", leaderCapture); controlCapture.getValue().advertise(); controlCapture.getValue().leave(); controlCapture.getValue().leave(); } @Test(expected = IllegalStateException.class) public void testAdvertiseAfterLeave() throws Exception { Capture<Leader> leaderCapture = createCapture(); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); Capture<LeaderControl> controlCapture = createCapture(); listener.onLeading(capture(controlCapture)); abdicate.execute(); control.replay(); newLeader("foo", leaderCapture); controlCapture.getValue().leave(); controlCapture.getValue().advertise(); } @Test public void testLeadMulti() throws Exception { List<Capture<Leader>> leaderCaptures = Lists.newArrayList(); List<Capture<LeaderControl>> leaderControlCaptures = Lists.newArrayList(); for (int i = 0; i < 5; i++) { Capture<Leader> leaderCapture = new Capture<Leader>(); leaderCaptures.add(leaderCapture); Capture<LeaderControl> controlCapture = createCapture(); leaderControlCaptures.add(controlCapture); expect(candidate.offerLeadership(capture(leaderCapture))).andReturn(null); listener.onLeading(capture(controlCapture)); InetSocketAddress primary = InetSocketAddress.createUnresolved("foo" + i, PORT_A); Map<String, InetSocketAddress> aux = ImmutableMap.of("http-admin", InetSocketAddress.createUnresolved("foo" + i, PORT_B)); expect(serverSet.join(primary, aux)).andReturn(endpointStatus); endpointStatus.leave(); abdicate.execute(); } control.replay(); for (int i = 0; i < 5; i++) { final String leaderName = "foo" + i; newLeader(leaderName, leaderCaptures.get(i)); leaderControlCaptures.get(i).getValue().advertise(); leaderControlCaptures.get(i).getValue().leave(); } } @Test public void testLeaderLeaves() throws Exception { control.replay(); shutdownNetwork(); } private static IAnswer<?> countDownAnswer(final CountDownLatch latch) { return new IAnswer<Void>() { @Override public Void answer() { latch.countDown(); return null; } }; } @Test public void testLeaderDisconnect() throws Exception { Capture<LeaderControl> controlCapture = createCapture(); CountDownLatch leading = new CountDownLatch(1); listener.onLeading(capture(controlCapture)); expectLastCall().andAnswer(countDownAnswer(leading)); CountDownLatch defeated = new CountDownLatch(1); listener.onDefeated(null); expectLastCall().andAnswer(countDownAnswer(defeated)); control.replay(); ZooKeeperClient zkClient = createZkClient(); serverSet = new ServerSetImpl(zkClient, "/fake/path"); candidate = new CandidateImpl(new Group(zkClient, ZooKeeperUtils.OPEN_ACL_UNSAFE, "/fake/path")); DefeatOnDisconnectLeader leader = new DefeatOnDisconnectLeader(zkClient, listener); service = new SingletonService(serverSet, candidate); service.lead(InetSocketAddress.createUnresolved("foo", PORT_A), ImmutableMap.of("http-admin", InetSocketAddress.createUnresolved("foo", PORT_B)), leader); leading.await(); shutdownNetwork(); defeated.await(); } @Test public void testNonLeaderDisconnect() throws Exception { CountDownLatch elected = new CountDownLatch(1); listener.onLeading(EasyMock.<LeaderControl>anyObject()); expectLastCall().andAnswer(countDownAnswer(elected)); listener.onDefeated(null); expectLastCall().anyTimes(); control.replay(); ZooKeeperClient zkClient = createZkClient(); String path = "/fake/path"; // Create a fake leading candidate node to ensure that the leader in this test is never // elected. ZooKeeperUtils.ensurePath(zkClient, ZooKeeperUtils.OPEN_ACL_UNSAFE, path); String leaderNode = zkClient.get().create(path + "/" + SingletonService.LEADER_ELECT_NODE_PREFIX, "fake_leader".getBytes(), ZooKeeperUtils.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT_SEQUENTIAL); serverSet = new ServerSetImpl(zkClient, path); candidate = SingletonService.createSingletonCandidate(zkClient, path, ZooKeeperUtils.OPEN_ACL_UNSAFE); DefeatOnDisconnectLeader leader = new DefeatOnDisconnectLeader(zkClient, listener); service = new SingletonService(serverSet, candidate); service.lead(InetSocketAddress.createUnresolved("foo", PORT_A), ImmutableMap.of("http-admin", InetSocketAddress.createUnresolved("foo", PORT_B)), leader); final CountDownLatch disconnected = new CountDownLatch(1); zkClient.register(new Watcher() { @Override public void process(WatchedEvent event) { if ((event.getType() == EventType.None) && (event.getState() == KeeperState.Disconnected)) { disconnected.countDown(); } } }); shutdownNetwork(); disconnected.await(); restartNetwork(); zkClient.get().delete(leaderNode, ZooKeeperUtils.ANY_VERSION); // Upon deletion of the fake leader node, the candidate should become leader. elected.await(); } }