helix-core/src/test/java/org/apache/helix/manager/zk/TestZkReconnect.java

package org.apache.helix.manager.zk;

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.helix.HelixException;
import org.apache.helix.HelixManager;
import org.apache.helix.HelixManagerFactory;
import org.apache.helix.InstanceType;
import org.apache.helix.SystemPropertyKeys;
import org.apache.helix.TestHelper;
import org.apache.helix.ZkTestHelper;
import org.apache.helix.zookeeper.api.client.HelixZkClient;
import org.apache.helix.store.zk.ZkHelixPropertyStore;
import org.apache.helix.tools.ClusterSetup;
import org.apache.helix.zookeeper.zkclient.ZkServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.Test;

public class TestZkReconnect {
  private static final Logger LOG = LoggerFactory.getLogger(TestZkReconnect.class);

  @Test
  public void testHelixManagerStateListenerCallback() throws Exception {
    final int zkPort = TestHelper.getRandomPort();
    final String zkAddr = String.format("localhost:%d", zkPort);
    final ZkServer zkServer = TestHelper.startZkServer(zkAddr);

    String className = TestHelper.getTestClassName();
    String methodName = TestHelper.getTestMethodName();
    final String clusterName = className + "_" + methodName;

    // Init onDisconnectedFlag to check if callback is triggered
    final AtomicReference<Boolean> onDisconnectedFlag = new AtomicReference<>(false);
    final AtomicReference<Boolean> onConnectedFlag = new AtomicReference<>(false);

    // Setup cluster
    LOG.info("Setup clusters");
    ClusterSetup clusterSetup = new ClusterSetup(zkAddr);
    clusterSetup.addCluster(clusterName, true);
    // For fast test, set short timeout
    System.setProperty(SystemPropertyKeys.ZK_CONNECTION_TIMEOUT, "2000");

    // Registers and starts controller, register listener for disconnect handling
    LOG.info("Starts controller");
    final ZKHelixManager controller =
        (ZKHelixManager) HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.CONTROLLER, zkAddr,
            new HelixManagerStateListener() {
              @Override
              public void onConnected(HelixManager helixManager) throws Exception {
                Assert.assertEquals(helixManager.getClusterName(), clusterName);
                onConnectedFlag.getAndSet(true);
              }

              @Override
              public void onDisconnected(HelixManager helixManager, Throwable error) throws Exception {
                Assert.assertEquals(helixManager.getClusterName(), clusterName);
                onDisconnectedFlag.getAndSet(true);
              }
            });

    Assert.assertEquals(controller.getMetadataStoreConnectionString(), zkAddr);

    try {
      controller.connect();
      // check onConnected() is triggered
      Assert.assertTrue(onConnectedFlag.getAndSet(false));

      // 1. shutdown zkServer and check if handler trigger callback
      zkServer.shutdown();
      // Simulate a retry in ZkClient that will not succeed
      ZkTestHelper.injectExpire(controller._zkclient);
      Assert.assertFalse(controller._zkclient.waitUntilConnected(5000, TimeUnit.MILLISECONDS));
      // While retrying, onDisconnectedFlag = false
      Assert.assertFalse(onDisconnectedFlag.get());

      // 2. restart zkServer and check if handler will recover connection
      zkServer.start();
      Assert.assertTrue(controller._zkclient
          .waitUntilConnected(HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, TimeUnit.MILLISECONDS));
      Assert.assertTrue(controller.isConnected());

      // New propertyStore should be in good state
      ZkHelixPropertyStore propertyStore = controller.getHelixPropertyStore();
      propertyStore.get("/", null, 0);

      TestHelper.verify(new TestHelper.Verifier() {
        @Override
        public boolean verify() throws Exception {
          return onConnectedFlag.getAndSet(false);
        }
      }, 1000);

      // Inject expire to test handler
      // onDisconnectedFlag should be set within onDisconnected handler
      controller.handleSessionEstablishmentError(new Exception("For testing"));
      TestHelper.verify(new TestHelper.Verifier() {
        @Override
        public boolean verify() throws Exception {
          return onDisconnectedFlag.get();
        }
      }, 1000);
      Assert.assertFalse(onConnectedFlag.get());
      Assert.assertFalse(controller.isConnected());

      // Verify ZK is down
      try {
        controller.getHelixPropertyStore();
      } catch (HelixException e) {
        // Expected exception
        System.out.println(e.getMessage());
      }
    } finally {
      controller.disconnect();
      zkServer.shutdown();
      System.clearProperty(SystemPropertyKeys.ZK_CONNECTION_TIMEOUT);
    }
  }
}