helix-core/src/test/java/org/apache/helix/TestZkBasis.java - helix - Git at Google

 package org.apache.helix;

 /*
  * 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.Collections;
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;

 import org.apache.helix.manager.zk.ZNRecordSerializer;
 import org.apache.helix.zookeeper.impl.client.ZkClient;
 import org.apache.helix.zookeeper.api.client.HelixZkClient;
 import org.apache.helix.zookeeper.zkclient.IZkChildListener;
 import org.apache.helix.zookeeper.zkclient.IZkDataListener;
 import org.testng.Assert;
 import org.testng.annotations.Test;

 /**
  * test zookeeper basis
  */
 public class TestZkBasis extends ZkUnitTestBase {
   class ZkListener implements IZkDataListener, IZkChildListener {
     String _parentPath = null;
     String _dataDeletePath = null;
     List<String> _currentChilds = Collections.emptyList(); // make sure it's set to null in
                                                            // #handleChildChange()

     CountDownLatch _childChangeCountDown = new CountDownLatch(1);
     CountDownLatch _dataDeleteCountDown = new CountDownLatch(1);

     @Override
     public void handleChildChange(String parentPath, List<String> currentChilds) {
       _parentPath = parentPath;
       _currentChilds = currentChilds;
       _childChangeCountDown.countDown();
     }

     @Override
     public void handleDataChange(String dataPath, Object data) {
       // To change body of implemented methods use File | Settings | File Templates.
     }

     @Override
     public void handleDataDeleted(String dataPath) {
       _dataDeletePath = dataPath;
       _dataDeleteCountDown.countDown();
     }
   }

   @Test
   public void testZkSessionExpiry() throws Exception {
     String className = TestHelper.getTestClassName();
     String methodName = TestHelper.getTestMethodName();
     String clusterName = className + "_" + methodName;
     System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

     ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
         HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());

     String path = String.format("/%s", clusterName);
     client.createEphemeral(path);
     String oldSessionId = ZkTestHelper.getSessionId(client);
     ZkTestHelper.expireSession(client);
     String newSessionId = ZkTestHelper.getSessionId(client);
     Assert.assertNotSame(newSessionId, oldSessionId);
     Assert.assertFalse(client.exists(path), "Ephemeral znode should be gone after session expiry");
     client.close();
     System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
   }

   @Test
   public void testCloseZkClient() {
     String className = TestHelper.getTestClassName();
     String methodName = TestHelper.getTestMethodName();
     String clusterName = className + "_" + methodName;
     System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

     ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
         HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
     String path = String.format("/%s", clusterName);
     client.createEphemeral(path);

     client.close();
     Assert.assertFalse(_gZkClient.exists(path),
         "Ephemeral node: " + path + " should be removed after ZkClient#close()");
     System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
   }

   @Test
   public void testCloseZkClientInZkClientEventThread() throws Exception {
     String className = TestHelper.getTestClassName();
     String methodName = TestHelper.getTestMethodName();
     String clusterName = className + "_" + methodName;
     System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

     final CountDownLatch waitCallback = new CountDownLatch(1);
     final ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
         HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
     String path = String.format("/%s", clusterName);
     client.createEphemeral(path);
     client.subscribeDataChanges(path, new IZkDataListener() {

       @Override
       public void handleDataDeleted(String dataPath) {
       }

       @Override
       public void handleDataChange(String dataPath, Object data) {
         client.close();
         waitCallback.countDown();
       }
     });

     client.writeData(path, new ZNRecord("test"));
     waitCallback.await();
     Assert.assertFalse(_gZkClient.exists(path), "Ephemeral node: " + path
         + " should be removed after ZkClient#close() in its own event-thread");

     System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));

   }

   /**
    * test zk watchers are renewed automatically after session expiry
    * zookeeper-client side keeps all registered watchers see ZooKeeper.WatchRegistration.register()
    * after session expiry, all watchers are renewed
    * if a path that has watches on it has been removed during session expiry,
    * the watchers on that path will still get callbacks after session renewal, especially:
    * a data-watch will get data-deleted callback
    * a child-watch will get a child-change callback with current-child-list = null
    * this can be used for cleanup watchers on the zookeeper-client side
    */
   @Test
   public void testWatchRenew() throws Exception {
     String className = TestHelper.getTestClassName();
     String methodName = TestHelper.getTestMethodName();
     String testName = className + "_" + methodName;

     final ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
         HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
     // make sure "/testName/test" doesn't exist
     final String path = "/" + testName + "/test";
     client.delete(path);

     ZkListener listener = new ZkListener();
     client.subscribeDataChanges(path, listener);
     client.subscribeChildChanges(path, listener);

     ZkTestHelper.expireSession(client);

     boolean succeed = listener._childChangeCountDown.await(10, TimeUnit.SECONDS);
     Assert.assertTrue(succeed,
         "fail to wait on child-change count-down in 10 seconds after session-expiry");
     Assert.assertEquals(listener._parentPath, path,
         "fail to get child-change callback after session-expiry");
     Assert.assertNull(listener._currentChilds,
         "fail to get child-change callback with currentChilds=null after session expiry");

     succeed = listener._dataDeleteCountDown.await(10, TimeUnit.SECONDS);
     Assert.assertTrue(succeed,
         "fail to wait on data-delete count-down in 10 seconds after session-expiry");
     Assert.assertEquals(listener._dataDeletePath, path,
         "fail to get data-delete callback after session-expiry");

     client.close();
   }

   /**
    * after calling zkclient#unsubscribeXXXListener()
    * an already registered watch will not be removed from ZooKeeper#watchManager#XXXWatches
    * immediately.
    * the watch will get removed on the following conditions:
    * 1) there is a set/delete on the listening path via the zkclient
    * 2) session expiry on the zkclient (i.e. the watch will not be renewed after session expiry)
    * @throws Exception
    */
   @Test
   public void testWatchRemove() throws Exception {
     String className = TestHelper.getTestClassName();
     String methodName = TestHelper.getTestMethodName();
     String testName = className + "_" + methodName;

     final ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
         HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
     // make sure "/testName/test" doesn't exist
     final String path = "/" + testName + "/test";
     client.createPersistent(path, true);

     ZkListener listener = new ZkListener();
     client.subscribeDataChanges(path, listener);
     client.subscribeChildChanges(path, listener);

     // listener should be in both ZkClient#_dataListener and ZkClient#_childListener set
     Map<String, Set<IZkDataListener>> dataListenerMap = ZkTestHelper.getZkDataListener(client);
     Assert.assertEquals(dataListenerMap.size(), 1, "ZkClient#_dataListener should have 1 listener");
     Set<IZkDataListener> dataListenerSet = dataListenerMap.get(path);
     Assert.assertNotNull(dataListenerSet,
         "ZkClient#_dataListener should have 1 listener on path: " + path);
     Assert.assertEquals(dataListenerSet.size(), 1,
         "ZkClient#_dataListener should have 1 listener on path: " + path);

     Map<String, Set<IZkChildListener>> childListenerMap = ZkTestHelper.getZkChildListener(client);
     Assert.assertEquals(childListenerMap.size(), 1,
         "ZkClient#_childListener should have 1 listener");
     Set<IZkChildListener> childListenerSet = childListenerMap.get(path);
     Assert.assertNotNull(childListenerSet,
         "ZkClient#_childListener should have 1 listener on path: " + path);
     Assert.assertEquals(childListenerSet.size(), 1,
         "ZkClient#_childListener should have 1 listener on path: " + path);

     // watch should be in ZooKeeper#watchManager#XXXWatches
     Map<String, List<String>> watchMap = ZkTestHelper.getZkWatch(client);
     // System.out.println("watchMap1: " + watchMap);
     List<String> dataWatch = watchMap.get("dataWatches");
     Assert.assertNotNull(dataWatch,
         "ZooKeeper#watchManager#dataWatches should have 1 data watch on path: " + path);
     Assert.assertEquals(dataWatch.size(), 1,
         "ZooKeeper#watchManager#dataWatches should have 1 data watch on path: " + path);
     Assert.assertEquals(dataWatch.get(0), path,
         "ZooKeeper#watchManager#dataWatches should have 1 data watch on path: " + path);

     List<String> childWatch = watchMap.get("childWatches");
     Assert.assertNotNull(childWatch,
         "ZooKeeper#watchManager#childWatches should have 1 child watch on path: " + path);
     Assert.assertEquals(childWatch.size(), 1,
         "ZooKeeper#watchManager#childWatches should have 1 child watch on path: " + path);
     Assert.assertEquals(childWatch.get(0), path,
         "ZooKeeper#watchManager#childWatches should have 1 child watch on path: " + path);

     client.unsubscribeDataChanges(path, listener);
     client.unsubscribeChildChanges(path, listener);
     // System.out.println("watchMap2: " + watchMap);
     ZkTestHelper.expireSession(client);

     // after session expiry, those watches should be removed
     watchMap = ZkTestHelper.getZkWatch(client);
     // System.out.println("watchMap3: " + watchMap);
     dataWatch = watchMap.get("dataWatches");
     Assert.assertTrue(dataWatch.isEmpty(), "ZooKeeper#watchManager#dataWatches should be empty");
     childWatch = watchMap.get("childWatches");
     Assert.assertTrue(childWatch.isEmpty(), "ZooKeeper#watchManager#childWatches should be empty");

     client.close();
     deleteCluster(testName);
   }
 }
	package org.apache.helix;

	/*
	* 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.Collections;
	import java.util.Date;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.concurrent.CountDownLatch;
	import java.util.concurrent.TimeUnit;

	import org.apache.helix.manager.zk.ZNRecordSerializer;
	import org.apache.helix.zookeeper.impl.client.ZkClient;
	import org.apache.helix.zookeeper.api.client.HelixZkClient;
	import org.apache.helix.zookeeper.zkclient.IZkChildListener;
	import org.apache.helix.zookeeper.zkclient.IZkDataListener;
	import org.testng.Assert;
	import org.testng.annotations.Test;

	/**
	* test zookeeper basis
	*/
	public class TestZkBasis extends ZkUnitTestBase {
	class ZkListener implements IZkDataListener, IZkChildListener {
	String _parentPath = null;
	String _dataDeletePath = null;
	List<String> _currentChilds = Collections.emptyList(); // make sure it's set to null in
	// #handleChildChange()

	CountDownLatch _childChangeCountDown = new CountDownLatch(1);
	CountDownLatch _dataDeleteCountDown = new CountDownLatch(1);

	@Override
	public void handleChildChange(String parentPath, List<String> currentChilds) {
	_parentPath = parentPath;
	_currentChilds = currentChilds;
	_childChangeCountDown.countDown();
	}

	@Override
	public void handleDataChange(String dataPath, Object data) {
	// To change body of implemented methods use File \| Settings \| File Templates.
	}

	@Override
	public void handleDataDeleted(String dataPath) {
	_dataDeletePath = dataPath;
	_dataDeleteCountDown.countDown();
	}
	}

	@Test
	public void testZkSessionExpiry() throws Exception {
	String className = TestHelper.getTestClassName();
	String methodName = TestHelper.getTestMethodName();
	String clusterName = className + "_" + methodName;
	System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

	ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
	HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());

	String path = String.format("/%s", clusterName);
	client.createEphemeral(path);
	String oldSessionId = ZkTestHelper.getSessionId(client);
	ZkTestHelper.expireSession(client);
	String newSessionId = ZkTestHelper.getSessionId(client);
	Assert.assertNotSame(newSessionId, oldSessionId);
	Assert.assertFalse(client.exists(path), "Ephemeral znode should be gone after session expiry");
	client.close();
	System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
	}

	@Test
	public void testCloseZkClient() {
	String className = TestHelper.getTestClassName();
	String methodName = TestHelper.getTestMethodName();
	String clusterName = className + "_" + methodName;
	System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

	ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
	HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
	String path = String.format("/%s", clusterName);
	client.createEphemeral(path);

	client.close();
	Assert.assertFalse(_gZkClient.exists(path),
	"Ephemeral node: " + path + " should be removed after ZkClient#close()");
	System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
	}

	@Test
	public void testCloseZkClientInZkClientEventThread() throws Exception {
	String className = TestHelper.getTestClassName();
	String methodName = TestHelper.getTestMethodName();
	String clusterName = className + "_" + methodName;
	System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

	final CountDownLatch waitCallback = new CountDownLatch(1);
	final ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
	HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
	String path = String.format("/%s", clusterName);
	client.createEphemeral(path);
	client.subscribeDataChanges(path, new IZkDataListener() {

	@Override
	public void handleDataDeleted(String dataPath) {
	}

	@Override
	public void handleDataChange(String dataPath, Object data) {
	client.close();
	waitCallback.countDown();
	}
	});

	client.writeData(path, new ZNRecord("test"));
	waitCallback.await();
	Assert.assertFalse(_gZkClient.exists(path), "Ephemeral node: " + path
	+ " should be removed after ZkClient#close() in its own event-thread");

	System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));

	}

	/**
	* test zk watchers are renewed automatically after session expiry
	* zookeeper-client side keeps all registered watchers see ZooKeeper.WatchRegistration.register()
	* after session expiry, all watchers are renewed
	* if a path that has watches on it has been removed during session expiry,
	* the watchers on that path will still get callbacks after session renewal, especially:
	* a data-watch will get data-deleted callback
	* a child-watch will get a child-change callback with current-child-list = null
	* this can be used for cleanup watchers on the zookeeper-client side
	*/
	@Test
	public void testWatchRenew() throws Exception {
	String className = TestHelper.getTestClassName();
	String methodName = TestHelper.getTestMethodName();
	String testName = className + "_" + methodName;

	final ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
	HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
	// make sure "/testName/test" doesn't exist
	final String path = "/" + testName + "/test";
	client.delete(path);

	ZkListener listener = new ZkListener();
	client.subscribeDataChanges(path, listener);
	client.subscribeChildChanges(path, listener);

	ZkTestHelper.expireSession(client);

	boolean succeed = listener._childChangeCountDown.await(10, TimeUnit.SECONDS);
	Assert.assertTrue(succeed,
	"fail to wait on child-change count-down in 10 seconds after session-expiry");
	Assert.assertEquals(listener._parentPath, path,
	"fail to get child-change callback after session-expiry");
	Assert.assertNull(listener._currentChilds,
	"fail to get child-change callback with currentChilds=null after session expiry");

	succeed = listener._dataDeleteCountDown.await(10, TimeUnit.SECONDS);
	Assert.assertTrue(succeed,
	"fail to wait on data-delete count-down in 10 seconds after session-expiry");
	Assert.assertEquals(listener._dataDeletePath, path,
	"fail to get data-delete callback after session-expiry");

	client.close();
	}

	/**
	* after calling zkclient#unsubscribeXXXListener()
	* an already registered watch will not be removed from ZooKeeper#watchManager#XXXWatches
	* immediately.
	* the watch will get removed on the following conditions:
	* 1) there is a set/delete on the listening path via the zkclient
	* 2) session expiry on the zkclient (i.e. the watch will not be renewed after session expiry)
	* @throws Exception
	*/
	@Test
	public void testWatchRemove() throws Exception {
	String className = TestHelper.getTestClassName();
	String methodName = TestHelper.getTestMethodName();
	String testName = className + "_" + methodName;

	final ZkClient client = new ZkClient(ZK_ADDR, HelixZkClient.DEFAULT_SESSION_TIMEOUT,
	HelixZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
	// make sure "/testName/test" doesn't exist
	final String path = "/" + testName + "/test";
	client.createPersistent(path, true);

	ZkListener listener = new ZkListener();
	client.subscribeDataChanges(path, listener);
	client.subscribeChildChanges(path, listener);

	// listener should be in both ZkClient#_dataListener and ZkClient#_childListener set
	Map<String, Set<IZkDataListener>> dataListenerMap = ZkTestHelper.getZkDataListener(client);
	Assert.assertEquals(dataListenerMap.size(), 1, "ZkClient#_dataListener should have 1 listener");
	Set<IZkDataListener> dataListenerSet = dataListenerMap.get(path);
	Assert.assertNotNull(dataListenerSet,
	"ZkClient#_dataListener should have 1 listener on path: " + path);
	Assert.assertEquals(dataListenerSet.size(), 1,
	"ZkClient#_dataListener should have 1 listener on path: " + path);

	Map<String, Set<IZkChildListener>> childListenerMap = ZkTestHelper.getZkChildListener(client);
	Assert.assertEquals(childListenerMap.size(), 1,
	"ZkClient#_childListener should have 1 listener");
	Set<IZkChildListener> childListenerSet = childListenerMap.get(path);
	Assert.assertNotNull(childListenerSet,
	"ZkClient#_childListener should have 1 listener on path: " + path);
	Assert.assertEquals(childListenerSet.size(), 1,
	"ZkClient#_childListener should have 1 listener on path: " + path);

	// watch should be in ZooKeeper#watchManager#XXXWatches
	Map<String, List<String>> watchMap = ZkTestHelper.getZkWatch(client);
	// System.out.println("watchMap1: " + watchMap);
	List<String> dataWatch = watchMap.get("dataWatches");
	Assert.assertNotNull(dataWatch,
	"ZooKeeper#watchManager#dataWatches should have 1 data watch on path: " + path);
	Assert.assertEquals(dataWatch.size(), 1,
	"ZooKeeper#watchManager#dataWatches should have 1 data watch on path: " + path);
	Assert.assertEquals(dataWatch.get(0), path,
	"ZooKeeper#watchManager#dataWatches should have 1 data watch on path: " + path);

	List<String> childWatch = watchMap.get("childWatches");
	Assert.assertNotNull(childWatch,
	"ZooKeeper#watchManager#childWatches should have 1 child watch on path: " + path);
	Assert.assertEquals(childWatch.size(), 1,
	"ZooKeeper#watchManager#childWatches should have 1 child watch on path: " + path);
	Assert.assertEquals(childWatch.get(0), path,
	"ZooKeeper#watchManager#childWatches should have 1 child watch on path: " + path);

	client.unsubscribeDataChanges(path, listener);
	client.unsubscribeChildChanges(path, listener);
	// System.out.println("watchMap2: " + watchMap);
	ZkTestHelper.expireSession(client);

	// after session expiry, those watches should be removed
	watchMap = ZkTestHelper.getZkWatch(client);
	// System.out.println("watchMap3: " + watchMap);
	dataWatch = watchMap.get("dataWatches");
	Assert.assertTrue(dataWatch.isEmpty(), "ZooKeeper#watchManager#dataWatches should be empty");
	childWatch = watchMap.get("childWatches");
	Assert.assertTrue(childWatch.isEmpty(), "ZooKeeper#watchManager#childWatches should be empty");

	client.close();
	deleteCluster(testName);
	}
	}