CuratorSingletonServiceTest.java

/**
 * 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 org.apache.aurora.scheduler.discovery;

import java.net.InetSocketAddress;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;

import org.apache.aurora.common.zookeeper.SingletonService;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.easymock.Capture;
import org.easymock.IMocksControl;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.Timeout;

import static org.easymock.EasyMock.capture;
import static org.easymock.EasyMock.createControl;
import static org.easymock.EasyMock.expectLastCall;
import static org.easymock.EasyMock.newCapture;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

public class CuratorSingletonServiceTest extends BaseCuratorDiscoveryTest {

  private IMocksControl control;
  // This test has a lot of blocking, this ensures we don't deadlock.
  private final Timeout timeout = new Timeout(1, TimeUnit.MINUTES);

  @Rule
  public Timeout getTimeout() {
    return timeout;
  }

  @Before
  public void setUpSingletonService() throws Exception {
    control = createControl();
    addTearDown(control::verify);
  }

  private SingletonService.LeadershipListener createMockLeadershipListener() {
    return control.createMock(SingletonService.LeadershipListener.class);
  }

  private void newLeader(
      CuratorFramework client,
      String hostName,
      SingletonService.LeadershipListener listener)
      throws Exception {

    CuratorSingletonService singletonService =
        new CuratorSingletonService(client, GROUP_PATH, MEMBER_TOKEN, CODEC);
    InetSocketAddress leaderEndpoint = InetSocketAddress.createUnresolved(hostName, PRIMARY_PORT);
    singletonService.lead(leaderEndpoint, ImmutableMap.of(), listener);
  }

  @Test
  public void testLeadAdvertise() throws Exception {
    SingletonService.LeadershipListener listener = createMockLeadershipListener();
    Capture<SingletonService.LeaderControl> capture = newCapture();
    listener.onLeading(capture(capture));
    expectLastCall();

    control.replay();

    startGroupMonitor();

    // Can't be leader until we try to lead.
    assertFalse(capture.hasCaptured());

    newLeader(getClient(), "host1", listener);

    // Wait to become leader.
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    awaitCapture(capture);

    // Leadership nodes should not be seen as service group nodes.
    assertEquals(ImmutableSet.of(), getGroupMonitor().get());

    capture.getValue().advertise();

    // Verify we've advertised as leader.
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    assertEquals(ImmutableSet.of(serviceInstance("host1")), getGroupMonitor().get());
  }

  @Test
  public void testAbdicateTransition() throws Exception {
    SingletonService.LeadershipListener host1Listener = createMockLeadershipListener();
    Capture<SingletonService.LeaderControl> host1OnLeadingCapture = newCapture();
    host1Listener.onLeading(capture(host1OnLeadingCapture));
    expectLastCall();

    SingletonService.LeadershipListener host2Listener = createMockLeadershipListener();
    Capture<SingletonService.LeaderControl> host2OnLeadingCapture = newCapture();
    host2Listener.onLeading(capture(host2OnLeadingCapture));
    expectLastCall();

    control.replay();

    startGroupMonitor();

    // Have host1 become leader.
    newLeader(getClient(), "host1", host1Listener);
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    awaitCapture(host1OnLeadingCapture);

    host1OnLeadingCapture.getValue().advertise();
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    assertEquals(ImmutableSet.of(serviceInstance("host1")), getGroupMonitor().get());

    newLeader(getClient(), "host2", host2Listener);
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    assertFalse(host2OnLeadingCapture.hasCaptured());

    // Now have host1 abdicate.
    host1OnLeadingCapture.getValue().leave();

    // Should see both the leadership and service group member nodes get cleaned up by host1.
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_REMOVED);
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_REMOVED);

    awaitCapture(host2OnLeadingCapture);
  }

  @Test
  public void testDefeatTransition() throws Exception {
    SingletonService.LeadershipListener host1Listener = createMockLeadershipListener();
    Capture<SingletonService.LeaderControl> host1OnLeadingCapture = newCapture();
    host1Listener.onLeading(capture(host1OnLeadingCapture));
    expectLastCall();

    CountDownLatch host1Defeated = new CountDownLatch(1);
    host1Listener.onDefeated();
    expectLastCall().andAnswer(() -> {
      host1Defeated.countDown();
      return null;
    });

    SingletonService.LeadershipListener host2Listener = createMockLeadershipListener();
    Capture<SingletonService.LeaderControl> host2OnLeadingCapture = newCapture();
    host2Listener.onLeading(capture(host2OnLeadingCapture));
    expectLastCall();

    control.replay();

    startGroupMonitor();

    // Have host1 become leader.
    newLeader(getClient(), "host1", host1Listener);
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    awaitCapture(host1OnLeadingCapture);

    host1OnLeadingCapture.getValue().advertise();
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    assertEquals(ImmutableSet.of(serviceInstance("host1")), getGroupMonitor().get());

    newLeader(startNewClient(), "host2", host2Listener);
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_ADDED);
    assertFalse(host2OnLeadingCapture.hasCaptured());

    // Simulate a session timeout - the ephemeral leader node goes away and host1 should be
    // defeated.
    expireSession(getClient());

    // Should see both the leadership and service group member nodes go away as part of session
    // expiration for ephemeral nodes.
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_REMOVED);
    expectGroupEvent(PathChildrenCacheEvent.Type.CHILD_REMOVED);

    awaitCapture(host2OnLeadingCapture);

    // No advertisement by host2 yet, even though it won leadership, but the host1 service group
    // node should have been cleaned up as tested above.
    assertEquals(ImmutableSet.of(), getGroupMonitor().get());

    // Eventually host1 should notice its been defeated.
    host1Defeated.await();
  }

  @Test
  public void testZKDisconnection() throws Exception {
    CountDownLatch leading = new CountDownLatch(1);
    AtomicBoolean leader = new AtomicBoolean(false);
    CountDownLatch defeated = new CountDownLatch(1);

    // The listener is executed in an internal thread of Curator, where it executes the leader
    // listener callbacks. The exceptions there are not propagated out, so we have our own
    // listener to validate behaviour.
    SingletonService.LeadershipListener listener = new SingletonService.LeadershipListener() {
      @Override
      public void onLeading(SingletonService.LeaderControl leaderControl) {
        leader.set(true);
        leading.countDown();
      }

      @Override
      public void onDefeated() {
        leader.set(false);
        defeated.countDown();
      }
    };

    control.replay();

    startGroupMonitor();

    CuratorFramework client = getClient();
    newLeader(client, "host1", listener);
    leading.await();

    causeDisconnection();
    assertTrue(leader.get());

    expireSession(client);
    defeated.await();

    assertFalse(leader.get());
  }

  private void awaitCapture(Capture<?> capture) throws InterruptedException {
    while (!capture.hasCaptured()) {
      Thread.sleep(1L);
    }
  }
}