TestZkLeaderElection.java

/*
 * Copyright (C) 2012 Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License 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.facebook.zookeeper.election;

import com.facebook.testing.MockExecutor;
import com.facebook.zookeeper.mock.MockWatcher;
import com.facebook.zookeeper.mock.MockZkConnectionManager;
import com.facebook.zookeeper.mock.MockZooKeeper;
import com.facebook.zookeeper.mock.MockZooKeeperDataStore;
import com.facebook.zookeeper.mock.MockZooKeeperFactory;
import com.facebook.zookeeper.VariablePayload;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.testng.Assert;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import java.util.List;

public class TestZkLeaderElection {
  private static final String testData1 = "testData1";
  private static final String testData2 = "testData2";
  private static final String testData3 = "testData3";
  private static final String electionRoot = "/root/work/election";
  private static final String candidateBaseName = "candidate";
  private MockZooKeeperDataStore dataStore;
  private MockZooKeeperFactory mockZooKeeperFactory;
  private MockZkConnectionManager mockZkConnectionManager1;
  private MockZkConnectionManager mockZkConnectionManager2;
  private MockZkConnectionManager mockZkConnectionManager3;
  private ZkLeaderElection zkLeaderElection1;
  private ZkLeaderElection zkLeaderElection2;
  private ZkLeaderElection zkLeaderElection3;
  private MockLeaderElectionCallback mockLeaderElectionCallback1;
  private MockLeaderElectionCallback mockLeaderElectionCallback2;
  private MockLeaderElectionCallback mockLeaderElectionCallback3;
  private VariablePayload variablePayload1;
  private VariablePayload variablePayload2;
  private VariablePayload variablePayload3;
  private MockExecutor mockExecutor1;
  private MockExecutor mockExecutor2;
  private MockExecutor mockExecutor3;
  private MockZooKeeper zk1;
  private MockZooKeeper zk2;
  private MockZooKeeper zk3;

  @BeforeMethod(alwaysRun = true)
  public void setUp() throws Exception {
    dataStore = new MockZooKeeperDataStore();
    mockZooKeeperFactory = new MockZooKeeperFactory(dataStore);
    mockZkConnectionManager1 = new MockZkConnectionManager(mockZooKeeperFactory);
    mockZkConnectionManager1.refreshClient(); // Generate a new client
    zk1 = mockZooKeeperFactory.getLastZooKeeper();
    zk1.triggerConnect();
    mockZkConnectionManager2 = new MockZkConnectionManager(mockZooKeeperFactory);
    mockZkConnectionManager2.refreshClient(); // Generate a new client
    zk2 = mockZooKeeperFactory.getLastZooKeeper();
    zk2.triggerConnect();
    mockZkConnectionManager3 = new MockZkConnectionManager(mockZooKeeperFactory);
    mockZkConnectionManager3.refreshClient(); // Generate a new client
    zk3 = mockZooKeeperFactory.getLastZooKeeper();
    zk3.triggerConnect();

    // Create ZkLeaderElection1
    mockExecutor1 = new MockExecutor();
    variablePayload1 = new VariablePayload(testData1);
    mockLeaderElectionCallback1 = new MockLeaderElectionCallback();
    zkLeaderElection1 =
      new ZkLeaderElection(
        mockZkConnectionManager1,
        electionRoot,
        candidateBaseName,
        variablePayload1,
        mockLeaderElectionCallback1,
        mockExecutor1
      );

    // Create ZkLeaderElection2
    mockExecutor2 = new MockExecutor();
    variablePayload2 = new VariablePayload(testData2);
    mockLeaderElectionCallback2 = new MockLeaderElectionCallback();
    zkLeaderElection2 =
      new ZkLeaderElection(
        mockZkConnectionManager2,
        electionRoot,
        candidateBaseName,
        variablePayload2,
        mockLeaderElectionCallback2,
        mockExecutor2      
      );

    // Create ZkLeaderElection3
    mockExecutor3 = new MockExecutor();
    variablePayload3 = new VariablePayload(testData3);
    mockLeaderElectionCallback3 = new MockLeaderElectionCallback();
    zkLeaderElection3 =
      new ZkLeaderElection(
        mockZkConnectionManager3,
        electionRoot,
        candidateBaseName,
        variablePayload3,
        mockLeaderElectionCallback3,
        mockExecutor3
      );

    // Set up election root
    zk1.create("/root", null, null, CreateMode.PERSISTENT);
    zk1.create("/root/work", null, null, CreateMode.PERSISTENT);
    zk1.create("/root/work/election", null, null, CreateMode.PERSISTENT);
  }

  @Test(groups = "fast", expectedExceptions = KeeperException.ConnectionLossException.class)
  public void testNoConnection() throws Exception {
    zk1.triggerDisconnect();
    zkLeaderElection1.enter();
  }

  @Test(groups = "fast", expectedExceptions = KeeperException.NoNodeException.class)
  public void testMissingRoot() throws Exception {
    zk1.delete(electionRoot, -1);
    zkLeaderElection1.enter();
  }

  @Test(groups = "fast")
  public void testEnter() throws Exception {
    // Verify no candidates
    MockWatcher mockWatcher = new MockWatcher();
    List<String> candidates = zk1.getChildren(electionRoot, mockWatcher);
    Assert.assertTrue(candidates.isEmpty());

    // Enter the candidate and run any callbacks
    zkLeaderElection1.enter();
    mockExecutor1.drain();

    // Check that candidate is created
    candidates = zk1.getChildren(electionRoot, null);
    Assert.assertEquals(candidates.size(), 1);

    // Check that candidate is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();
    Assert.assertEquals(zkLeaderElection1.getLeader(), candidates.get(0));

    // Check data contents
    String data = VariablePayload.decode(
      zk1.getData(electionRoot + "/" + candidates.get(0), null, null)
    );
    Assert.assertEquals(data, testData1);

    // Check that external watch was notified of candidate creation
    Assert.assertEquals(mockWatcher.getEventQueue().size(), 1);
    WatchedEvent watchedEvent = mockWatcher.getEventQueue().remove();
    Assert.assertEquals(watchedEvent.getType(), EventType.NodeChildrenChanged);
    Assert.assertEquals(watchedEvent.getPath(), electionRoot);
  }

  @Test(groups = "fast")
  public void testMultipleEnter() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate is notified of election win again
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    // Only one candidate should be entered
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 1);
  }

  @Test(groups = "fast")
  public void testPrematureWithdraw() throws Exception {
    // No exception should be thrown if the candidate does not exist
    zkLeaderElection1.withdraw();
    mockExecutor1.drain();
  }

  @Test(groups = "fast")
  public void testWithdraw() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();

    // Verify candidate
    MockWatcher mockWatcher = new MockWatcher();
    List<String> candidates = zk1.getChildren(electionRoot, mockWatcher);
    Assert.assertEquals(candidates.size(), 1);
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection1.withdraw();
    mockExecutor1.drain();

    candidates = zk1.getChildren(electionRoot, null);
    Assert.assertTrue(candidates.isEmpty());
    Assert.assertEquals(zkLeaderElection1.getLeader(), null);
    // Manual withdraw should not trigger a "removed" callback
    Assert.assertFalse(mockLeaderElectionCallback1.isRemoved());
    Assert.assertEquals(mockWatcher.getEventQueue().size(), 1);
    WatchedEvent watchedEvent = mockWatcher.getEventQueue().remove();
    Assert.assertEquals(watchedEvent.getType(), EventType.NodeChildrenChanged);
    Assert.assertEquals(watchedEvent.getPath(), electionRoot);
  }

  @Test(groups = "fast")
  public void testMultipleWithdraw() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();

    zkLeaderElection1.withdraw();
    mockExecutor1.drain();
    // Check candidate is removed
    Assert.assertTrue(zk1.getChildren(electionRoot, null).isEmpty());

    zkLeaderElection1.withdraw();
    mockExecutor1.drain();
    // Check that there is still no candidate
    Assert.assertTrue(zk1.getChildren(electionRoot, null).isEmpty());
  }

  @Test(groups = "fast")
  public void testCycle() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection1.cycle();
    mockExecutor1.drain();
    // Check that candidate is re-elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();
    // Check no remove was signaled
    Assert.assertFalse(mockLeaderElectionCallback1.isRemoved());
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 1);
  }

  @Test(groups = "fast")
  public void testEarlyCycle() throws Exception {
    zkLeaderElection1.cycle();
    mockExecutor1.drain();
    // Check that candidate is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();
    // Check no remove was signaled
    Assert.assertFalse(mockLeaderElectionCallback1.isRemoved());
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 1);
  }

  @Test(groups = "fast")
  public void testMultiEnter() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate1 is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection2.enter();
    mockExecutor2.drain();
    // Check that candidate2 was not elected
    Assert.assertFalse(mockLeaderElectionCallback2.isElected());

    // Check that both candidates are in the election
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 2);
    // Check that they both agree on the winner
    Assert.assertEquals(zkLeaderElection1.getLeader(), zkLeaderElection2.getLeader());
    // Check that the reported winner is candidate1
    Assert.assertEquals(
      VariablePayload.decode(zk1.getData(
        electionRoot + "/" + zkLeaderElection1.getLeader(), null, null
      )),
      testData1
    );
  }

  @Test(groups = "fast")
  public void testMultiCycle() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate1 is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection2.enter();
    mockExecutor2.drain();
    // Check that candidate2 was not elected
    Assert.assertFalse(mockLeaderElectionCallback2.isElected());

    zkLeaderElection1.cycle();
    mockExecutor1.drain();
    mockExecutor2.drain();
    // Check that candidate2 was elected and not candidate1
    Assert.assertTrue(mockLeaderElectionCallback2.isElected());
    mockLeaderElectionCallback2.resetElected();
    Assert.assertFalse(mockLeaderElectionCallback1.isElected());

    // Check that both candidates are in the election
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 2);
    // Check that they both agree on the winner
    Assert.assertEquals(zkLeaderElection1.getLeader(), zkLeaderElection2.getLeader());
    // Check that the reported winner is candidate2
    Assert.assertEquals(
      VariablePayload.decode(zk1.getData(
        electionRoot + "/" + zkLeaderElection1.getLeader(), null, null
      )),
      testData2
    );
  }

  @Test(groups = "fast")
  public void testMultiWithExpire() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate1 is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection2.enter();
    mockExecutor2.drain();
    // Check that candidate2 was not elected
    Assert.assertFalse(mockLeaderElectionCallback2.isElected());

    zk1.triggerSessionExpiration();
    mockExecutor1.drain();
    mockExecutor2.drain();
    // Check that candidate2 was elected
    Assert.assertTrue(mockLeaderElectionCallback2.isElected());
    mockLeaderElectionCallback2.resetElected();

    // Check that only candidate2 is in the election
    Assert.assertEquals(zk2.getChildren(electionRoot, null).size(), 1);
    // Check that the reported winner is candidate2
    Assert.assertEquals(
      VariablePayload.decode(zk2.getData(
        electionRoot + "/" + zkLeaderElection2.getLeader(), null, null
      )),
      testData2
    );
  }

  @Test(groups = "fast")
  public void testAdminRemove() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate1 is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection2.enter();
    mockExecutor2.drain();
    // Check that candidate2 was not elected
    Assert.assertFalse(mockLeaderElectionCallback2.isElected());

    // Delete candidate1 via external forces
    zk1.delete(electionRoot + "/" + zkLeaderElection1.getLeader(), -1);
    mockExecutor1.drain();
    mockExecutor2.drain();
    //Check that candidate1 got a removed signal
    Assert.assertTrue(mockLeaderElectionCallback1.isRemoved());
    mockLeaderElectionCallback1.resetRemoved();
    // Check that candidate2 was elected
    Assert.assertTrue(mockLeaderElectionCallback2.isElected());
    mockLeaderElectionCallback2.resetElected();

    // Check that only candidate2 is in the election
    Assert.assertEquals(zk2.getChildren(electionRoot, null).size(), 1);
    // Check that the reported winner is candidate2
    Assert.assertEquals(
      VariablePayload.decode(zk2.getData(
        electionRoot + "/" + zkLeaderElection2.getLeader(), null, null
      )),
      testData2
    );
  }

  @Test(groups = "fast")
  public void testMultiSuccession() throws Exception {
    zkLeaderElection1.enter();
    mockExecutor1.drain();
    // Check that candidate1 is elected
    Assert.assertTrue(mockLeaderElectionCallback1.isElected());
    mockLeaderElectionCallback1.resetElected();

    zkLeaderElection2.enter();
    mockExecutor2.drain();
    // Check that candidate2 was not elected
    Assert.assertFalse(mockLeaderElectionCallback2.isElected());

    zkLeaderElection3.enter();
    mockExecutor3.drain();
    // Check that candidate3 was not elected
    Assert.assertFalse(mockLeaderElectionCallback3.isElected());

    // Check that all candidates are in the election
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 3);
    // Check that they all agree on the winner
    Assert.assertEquals(zkLeaderElection1.getLeader(), zkLeaderElection2.getLeader());
    Assert.assertEquals(zkLeaderElection2.getLeader(), zkLeaderElection3.getLeader());

    zkLeaderElection2.withdraw();
    mockExecutor2.drain();
    mockExecutor3.drain();
    // Check that no election callbacks happened
    Assert.assertFalse(mockLeaderElectionCallback1.isElected());
    Assert.assertFalse(mockLeaderElectionCallback2.isElected());
    Assert.assertFalse(mockLeaderElectionCallback3.isElected());

    // Check that there are now two candidates in the election
    Assert.assertEquals(zk1.getChildren(electionRoot, null).size(), 2);

    zkLeaderElection1.withdraw();
    mockExecutor1.drain();
    mockExecutor3.drain();
    // Check that candidate3 got the election callback
    Assert.assertTrue(mockLeaderElectionCallback3.isElected());
    mockLeaderElectionCallback3.resetElected();

    // Check that the reported winner is candidate3
    Assert.assertEquals(
      VariablePayload.decode(zk1.getData(
        electionRoot + "/" + zkLeaderElection1.getLeader(), null, null
      )),
      testData3
    );
  }
}