helix-core/src/main/java/org/apache/helix/tools/ClusterVerifiers/ZkHelixClusterVerifier.java - helix - Git at Google

 package org.apache.helix.tools.ClusterVerifiers;

 /*
  * 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.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;

 import org.I0Itec.zkclient.IZkChildListener;
 import org.I0Itec.zkclient.IZkDataListener;
 import org.apache.helix.HelixDataAccessor;
 import org.apache.helix.PropertyKey;
 import org.apache.helix.api.listeners.PreFetch;
 import org.apache.helix.manager.zk.ZKHelixDataAccessor;
 import org.apache.helix.manager.zk.ZNRecordSerializer;
 import org.apache.helix.manager.zk.ZkBaseDataAccessor;
 import org.apache.helix.manager.zk.ZkClient;
 import org.apache.helix.manager.zk.client.DedicatedZkClientFactory;
 import org.apache.helix.manager.zk.client.HelixZkClient;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public abstract class ZkHelixClusterVerifier
     implements IZkChildListener, IZkDataListener, HelixClusterVerifier {
   private static Logger LOG = LoggerFactory.getLogger(ZkHelixClusterVerifier.class);
   protected static int DEFAULT_TIMEOUT = 300 * 1000;
   protected static int DEFAULT_PERIOD = 500;

   protected final HelixZkClient _zkClient;
   // true if ZkHelixClusterVerifier was instantiated with a HelixZkClient, false otherwise
   // This is used for close() to determine how ZkHelixClusterVerifier should close the underlying
   // ZkClient
   private final boolean _usesExternalZkClient;
   protected final String _clusterName;
   protected final HelixDataAccessor _accessor;
   protected final PropertyKey.Builder _keyBuilder;
   private CountDownLatch _countdown;

   private ExecutorService _verifyTaskThreadPool =
       Executors.newSingleThreadExecutor(r -> new Thread(r, "ZkHelixClusterVerifier-verify_thread"));

   protected static class ClusterVerifyTrigger {
     final PropertyKey _triggerKey;
     final boolean _triggerOnDataChange;
     final boolean _triggerOnChildChange;
     final boolean _triggerOnChildDataChange;

     public ClusterVerifyTrigger(PropertyKey triggerKey, boolean triggerOnDataChange,
         boolean triggerOnChildChange, boolean triggerOnChildDataChange) {
       _triggerKey = triggerKey;
       _triggerOnDataChange = triggerOnDataChange;
       _triggerOnChildChange = triggerOnChildChange;
       _triggerOnChildDataChange = triggerOnChildDataChange;
     }

     public boolean isTriggerOnDataChange() {
       return _triggerOnDataChange;
     }

     public PropertyKey getTriggerKey() {
       return _triggerKey;
     }

     public boolean isTriggerOnChildChange() {
       return _triggerOnChildChange;
     }

     public boolean isTriggerOnChildDataChange() {
       return _triggerOnChildDataChange;
     }
   }

   public ZkHelixClusterVerifier(HelixZkClient zkClient, String clusterName) {
     if (zkClient == null || clusterName == null) {
       throw new IllegalArgumentException("requires zkClient|clusterName");
     }
     _zkClient = zkClient;
     _usesExternalZkClient = true;
     _clusterName = clusterName;
     _accessor = new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
     _keyBuilder = _accessor.keyBuilder();
   }

   public ZkHelixClusterVerifier(String zkAddr, String clusterName) {
     if (zkAddr == null || clusterName == null) {
       throw new IllegalArgumentException("requires zkAddr|clusterName");
     }
     _zkClient = DedicatedZkClientFactory.getInstance()
         .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkAddr));
     _usesExternalZkClient = false;
     _zkClient.setZkSerializer(new ZNRecordSerializer());
     _clusterName = clusterName;
     _accessor = new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
     _keyBuilder = _accessor.keyBuilder();
   }

   /**
    * Verify the cluster.
    * The method will be blocked at most {@code timeout}.
    * Return true if the verify succeed, otherwise return false.
    * @param timeout in milliseconds
    * @return true if succeed, false if not.
    */
   public boolean verify(long timeout) {
     return verifyByZkCallback(timeout);
   }

   /**
    * Verify the cluster.
    * The method will be blocked at most 30 seconds.
    * Return true if the verify succeed, otherwise return false.
    * @return true if succeed, false if not.
    */
   public boolean verify() {
     return verify(DEFAULT_TIMEOUT);
   }

   /**
    * Verify the cluster by relying on zookeeper callback and verify.
    * The method will be blocked at most {@code timeout}.
    * Return true if the verify succeed, otherwise return false.
    * @param timeout in milliseconds
    * @return true if succeed, false if not.
    */
   public abstract boolean verifyByZkCallback(long timeout);

   /**
    * Verify the cluster by relying on zookeeper callback and verify.
    * The method will be blocked at most 30 seconds.
    * Return true if the verify succeed, otherwise return false.
    * @return true if succeed, false if not.
    */
   public boolean verifyByZkCallback() {
     return verifyByZkCallback(DEFAULT_TIMEOUT);
   }

   /**
    * Verify the cluster by periodically polling the cluster status and verify.
    * The method will be blocked at most {@code timeout}.
    * Return true if the verify succeed, otherwise return false.
    * @param timeout
    * @param period polling interval
    * @return
    */
   public boolean verifyByPolling(long timeout, long period) {
     try {
       long start = System.currentTimeMillis();
       boolean success;
       do {
         success = verifyState();
         if (success) {
           return true;
         }
         TimeUnit.MILLISECONDS.sleep(period);
       } while ((System.currentTimeMillis() - start) <= timeout);
     } catch (Exception e) {
       LOG.error("Exception in verifier", e);
     }
     return false;
   }

   /**
    * Verify the cluster by periodically polling the cluster status and verify.
    * The method will be blocked at most 30 seconds.
    * Return true if the verify succeed, otherwise return false.
    * @return true if succeed, false if not.
    */
   public boolean verifyByPolling() {
     return verifyByPolling(DEFAULT_TIMEOUT, DEFAULT_PERIOD);
   }

   public void close() {
     if (_zkClient != null && !_usesExternalZkClient) {
       _zkClient.close();
     }
   }

   protected boolean verifyByCallback(long timeout, List<ClusterVerifyTrigger> triggers) {
     _countdown = new CountDownLatch(1);

     for (ClusterVerifyTrigger trigger : triggers) {
       subscribeTrigger(trigger);
     }

     boolean success = false;
     try {
       success = verifyState();
       if (!success) {
         success = _countdown.await(timeout, TimeUnit.MILLISECONDS);
         if (!success) {
           // make a final try if timeout
           success = verifyState();
         }
       }
     } catch (Exception e) {
       LOG.error("Exception in verifier", e);
     }

     // clean up
     _zkClient.unsubscribeAll();
     _verifyTaskThreadPool.shutdownNow();

     return success;
   }

   private void subscribeTrigger(ClusterVerifyTrigger trigger) {
     String path = trigger.getTriggerKey().getPath();
     if (trigger.isTriggerOnDataChange()) {
       _zkClient.subscribeDataChanges(path, this);
     }

     if (trigger.isTriggerOnChildChange()) {
       _zkClient.subscribeChildChanges(path, this);
     }

     if (trigger.isTriggerOnChildDataChange()) {
       List<String> childs = _zkClient.getChildren(path);
       for (String child : childs) {
         String childPath = String.format("%s/%s", path, child);
         _zkClient.subscribeDataChanges(childPath, this);
       }
     }
   }

   /**
    * The method actually performs the required verifications.
    * @return
    * @throws Exception
    */
   protected abstract boolean verifyState() throws Exception;

   class VerifyStateCallbackTask implements Runnable {
     @Override
     public void run() {
       try {
         boolean success = verifyState();
         if (success) {
           _countdown.countDown();
         }
       } catch (Exception ex) {
         LOG.info("verifyState() throws exception: " + ex);
       }
     }
   }

   @Override
   @PreFetch(enabled = false)
   public void handleDataChange(String dataPath, Object data) throws Exception {
     if (!_verifyTaskThreadPool.isShutdown()) {
       _verifyTaskThreadPool.submit(new VerifyStateCallbackTask());
     }
   }

   @Override
   public void handleDataDeleted(String dataPath) throws Exception {
     _zkClient.unsubscribeDataChanges(dataPath, this);
     if (!_verifyTaskThreadPool.isShutdown()) {
       _verifyTaskThreadPool.submit(new VerifyStateCallbackTask());
     }
   }

   @Override
   public void handleChildChange(String parentPath, List<String> currentChilds) throws Exception {
     for (String child : currentChilds) {
       String childPath = String.format("%s/%s", parentPath, child);
       _zkClient.subscribeDataChanges(childPath, this);
     }
     if (!_verifyTaskThreadPool.isShutdown()) {
       _verifyTaskThreadPool.submit(new VerifyStateCallbackTask());
     }
   }

   public HelixZkClient getHelixZkClient() {
     return _zkClient;
   }

   @Deprecated
   public ZkClient getZkClient() {
     return (ZkClient) getHelixZkClient();
   }

   public String getClusterName() {
     return _clusterName;
   }
 }
	package org.apache.helix.tools.ClusterVerifiers;

	/*
	* 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.List;
	import java.util.concurrent.CountDownLatch;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Executors;
	import java.util.concurrent.TimeUnit;

	import org.I0Itec.zkclient.IZkChildListener;
	import org.I0Itec.zkclient.IZkDataListener;
	import org.apache.helix.HelixDataAccessor;
	import org.apache.helix.PropertyKey;
	import org.apache.helix.api.listeners.PreFetch;
	import org.apache.helix.manager.zk.ZKHelixDataAccessor;
	import org.apache.helix.manager.zk.ZNRecordSerializer;
	import org.apache.helix.manager.zk.ZkBaseDataAccessor;
	import org.apache.helix.manager.zk.ZkClient;
	import org.apache.helix.manager.zk.client.DedicatedZkClientFactory;
	import org.apache.helix.manager.zk.client.HelixZkClient;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public abstract class ZkHelixClusterVerifier
	implements IZkChildListener, IZkDataListener, HelixClusterVerifier {
	private static Logger LOG = LoggerFactory.getLogger(ZkHelixClusterVerifier.class);
	protected static int DEFAULT_TIMEOUT = 300 * 1000;
	protected static int DEFAULT_PERIOD = 500;

	protected final HelixZkClient _zkClient;
	// true if ZkHelixClusterVerifier was instantiated with a HelixZkClient, false otherwise
	// This is used for close() to determine how ZkHelixClusterVerifier should close the underlying
	// ZkClient
	private final boolean _usesExternalZkClient;
	protected final String _clusterName;
	protected final HelixDataAccessor _accessor;
	protected final PropertyKey.Builder _keyBuilder;
	private CountDownLatch _countdown;

	private ExecutorService _verifyTaskThreadPool =
	Executors.newSingleThreadExecutor(r -> new Thread(r, "ZkHelixClusterVerifier-verify_thread"));

	protected static class ClusterVerifyTrigger {
	final PropertyKey _triggerKey;
	final boolean _triggerOnDataChange;
	final boolean _triggerOnChildChange;
	final boolean _triggerOnChildDataChange;

	public ClusterVerifyTrigger(PropertyKey triggerKey, boolean triggerOnDataChange,
	boolean triggerOnChildChange, boolean triggerOnChildDataChange) {
	_triggerKey = triggerKey;
	_triggerOnDataChange = triggerOnDataChange;
	_triggerOnChildChange = triggerOnChildChange;
	_triggerOnChildDataChange = triggerOnChildDataChange;
	}

	public boolean isTriggerOnDataChange() {
	return _triggerOnDataChange;
	}

	public PropertyKey getTriggerKey() {
	return _triggerKey;
	}

	public boolean isTriggerOnChildChange() {
	return _triggerOnChildChange;
	}

	public boolean isTriggerOnChildDataChange() {
	return _triggerOnChildDataChange;
	}
	}

	public ZkHelixClusterVerifier(HelixZkClient zkClient, String clusterName) {
	if (zkClient == null \|\| clusterName == null) {
	throw new IllegalArgumentException("requires zkClient\|clusterName");
	}
	_zkClient = zkClient;
	_usesExternalZkClient = true;
	_clusterName = clusterName;
	_accessor = new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
	_keyBuilder = _accessor.keyBuilder();
	}

	public ZkHelixClusterVerifier(String zkAddr, String clusterName) {
	if (zkAddr == null \|\| clusterName == null) {
	throw new IllegalArgumentException("requires zkAddr\|clusterName");
	}
	_zkClient = DedicatedZkClientFactory.getInstance()
	.buildZkClient(new HelixZkClient.ZkConnectionConfig(zkAddr));
	_usesExternalZkClient = false;
	_zkClient.setZkSerializer(new ZNRecordSerializer());
	_clusterName = clusterName;
	_accessor = new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
	_keyBuilder = _accessor.keyBuilder();
	}

	/**
	* Verify the cluster.
	* The method will be blocked at most {@code timeout}.
	* Return true if the verify succeed, otherwise return false.
	* @param timeout in milliseconds
	* @return true if succeed, false if not.
	*/
	public boolean verify(long timeout) {
	return verifyByZkCallback(timeout);
	}

	/**
	* Verify the cluster.
	* The method will be blocked at most 30 seconds.
	* Return true if the verify succeed, otherwise return false.
	* @return true if succeed, false if not.
	*/
	public boolean verify() {
	return verify(DEFAULT_TIMEOUT);
	}

	/**
	* Verify the cluster by relying on zookeeper callback and verify.
	* The method will be blocked at most {@code timeout}.
	* Return true if the verify succeed, otherwise return false.
	* @param timeout in milliseconds
	* @return true if succeed, false if not.
	*/
	public abstract boolean verifyByZkCallback(long timeout);

	/**
	* Verify the cluster by relying on zookeeper callback and verify.
	* The method will be blocked at most 30 seconds.
	* Return true if the verify succeed, otherwise return false.
	* @return true if succeed, false if not.
	*/
	public boolean verifyByZkCallback() {
	return verifyByZkCallback(DEFAULT_TIMEOUT);
	}

	/**
	* Verify the cluster by periodically polling the cluster status and verify.
	* The method will be blocked at most {@code timeout}.
	* Return true if the verify succeed, otherwise return false.
	* @param timeout
	* @param period polling interval
	* @return
	*/
	public boolean verifyByPolling(long timeout, long period) {
	try {
	long start = System.currentTimeMillis();
	boolean success;
	do {
	success = verifyState();
	if (success) {
	return true;
	}
	TimeUnit.MILLISECONDS.sleep(period);
	} while ((System.currentTimeMillis() - start) <= timeout);
	} catch (Exception e) {
	LOG.error("Exception in verifier", e);
	}
	return false;
	}

	/**
	* Verify the cluster by periodically polling the cluster status and verify.
	* The method will be blocked at most 30 seconds.
	* Return true if the verify succeed, otherwise return false.
	* @return true if succeed, false if not.
	*/
	public boolean verifyByPolling() {
	return verifyByPolling(DEFAULT_TIMEOUT, DEFAULT_PERIOD);
	}

	public void close() {
	if (_zkClient != null && !_usesExternalZkClient) {
	_zkClient.close();
	}
	}

	protected boolean verifyByCallback(long timeout, List<ClusterVerifyTrigger> triggers) {
	_countdown = new CountDownLatch(1);

	for (ClusterVerifyTrigger trigger : triggers) {
	subscribeTrigger(trigger);
	}

	boolean success = false;
	try {
	success = verifyState();
	if (!success) {
	success = _countdown.await(timeout, TimeUnit.MILLISECONDS);
	if (!success) {
	// make a final try if timeout
	success = verifyState();
	}
	}
	} catch (Exception e) {
	LOG.error("Exception in verifier", e);
	}

	// clean up
	_zkClient.unsubscribeAll();
	_verifyTaskThreadPool.shutdownNow();

	return success;
	}

	private void subscribeTrigger(ClusterVerifyTrigger trigger) {
	String path = trigger.getTriggerKey().getPath();
	if (trigger.isTriggerOnDataChange()) {
	_zkClient.subscribeDataChanges(path, this);
	}

	if (trigger.isTriggerOnChildChange()) {
	_zkClient.subscribeChildChanges(path, this);
	}

	if (trigger.isTriggerOnChildDataChange()) {
	List<String> childs = _zkClient.getChildren(path);
	for (String child : childs) {
	String childPath = String.format("%s/%s", path, child);
	_zkClient.subscribeDataChanges(childPath, this);
	}
	}
	}

	/**
	* The method actually performs the required verifications.
	* @return
	* @throws Exception
	*/
	protected abstract boolean verifyState() throws Exception;

	class VerifyStateCallbackTask implements Runnable {
	@Override
	public void run() {
	try {
	boolean success = verifyState();
	if (success) {
	_countdown.countDown();
	}
	} catch (Exception ex) {
	LOG.info("verifyState() throws exception: " + ex);
	}
	}
	}

	@Override
	@PreFetch(enabled = false)
	public void handleDataChange(String dataPath, Object data) throws Exception {
	if (!_verifyTaskThreadPool.isShutdown()) {
	_verifyTaskThreadPool.submit(new VerifyStateCallbackTask());
	}
	}

	@Override
	public void handleDataDeleted(String dataPath) throws Exception {
	_zkClient.unsubscribeDataChanges(dataPath, this);
	if (!_verifyTaskThreadPool.isShutdown()) {
	_verifyTaskThreadPool.submit(new VerifyStateCallbackTask());
	}
	}

	@Override
	public void handleChildChange(String parentPath, List<String> currentChilds) throws Exception {
	for (String child : currentChilds) {
	String childPath = String.format("%s/%s", parentPath, child);
	_zkClient.subscribeDataChanges(childPath, this);
	}
	if (!_verifyTaskThreadPool.isShutdown()) {
	_verifyTaskThreadPool.submit(new VerifyStateCallbackTask());
	}
	}

	public HelixZkClient getHelixZkClient() {
	return _zkClient;
	}

	@Deprecated
	public ZkClient getZkClient() {
	return (ZkClient) getHelixZkClient();
	}

	public String getClusterName() {
	return _clusterName;
	}
	}