Java tutorial
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.io.BufferedReader; import java.io.InputStreamReader; import java.io.PrintWriter; import java.net.Socket; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.TreeMap; import java.util.TreeSet; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; import org.I0Itec.zkclient.IZkChildListener; import org.I0Itec.zkclient.IZkDataListener; import org.I0Itec.zkclient.IZkStateListener; import org.I0Itec.zkclient.ZkConnection; import org.apache.helix.PropertyKey.Builder; import org.apache.helix.manager.zk.ZKHelixDataAccessor; import org.apache.helix.manager.zk.ZkBaseDataAccessor; import org.apache.helix.manager.zk.ZkClient; import org.apache.helix.model.ExternalView; import org.apache.log4j.Logger; import org.apache.zookeeper.WatchedEvent; import org.apache.zookeeper.Watcher; import org.apache.zookeeper.Watcher.Event.EventType; import org.apache.zookeeper.Watcher.Event.KeeperState; import org.apache.zookeeper.ZooKeeper; import org.apache.zookeeper.ZooKeeper.States; import org.testng.Assert; public class ZkTestHelper { private static Logger LOG = Logger.getLogger(ZkTestHelper.class); static { // Logger.getRootLogger().setLevel(Level.DEBUG); } public static void disconnectSession(final ZkClient zkClient) throws Exception { IZkStateListener listener = new IZkStateListener() { @Override public void handleStateChanged(KeeperState state) throws Exception { // System.err.println("disconnectSession handleStateChanged. state: " + state); } @Override public void handleNewSession() throws Exception { // make sure zkclient is connected again zkClient.waitUntilConnected(); ZkConnection connection = ((ZkConnection) zkClient.getConnection()); ZooKeeper curZookeeper = connection.getZookeeper(); LOG.info("handleNewSession. sessionId: " + Long.toHexString(curZookeeper.getSessionId())); } }; zkClient.subscribeStateChanges(listener); ZkConnection connection = ((ZkConnection) zkClient.getConnection()); ZooKeeper curZookeeper = connection.getZookeeper(); LOG.info("Before expiry. sessionId: " + Long.toHexString(curZookeeper.getSessionId())); Watcher watcher = new Watcher() { @Override public void process(WatchedEvent event) { LOG.info("Process watchEvent: " + event); } }; final ZooKeeper dupZookeeper = new ZooKeeper(connection.getServers(), curZookeeper.getSessionTimeout(), watcher, curZookeeper.getSessionId(), curZookeeper.getSessionPasswd()); // wait until connected, then close while (dupZookeeper.getState() != States.CONNECTED) { Thread.sleep(10); } dupZookeeper.close(); connection = (ZkConnection) zkClient.getConnection(); curZookeeper = connection.getZookeeper(); zkClient.unsubscribeStateChanges(listener); // System.err.println("zk: " + oldZookeeper); LOG.info("After expiry. sessionId: " + Long.toHexString(curZookeeper.getSessionId())); } /** * Simulate a zk state change by calling {@link ZkClient#process(WatchedEvent)} directly */ public static void simulateZkStateDisconnected(ZkClient client) { WatchedEvent event = new WatchedEvent(EventType.None, KeeperState.Disconnected, null); client.process(event); } /** * Get zk connection session id * @param client * @return */ public static String getSessionId(ZkClient client) { ZkConnection connection = ((ZkConnection) client.getConnection()); ZooKeeper curZookeeper = connection.getZookeeper(); return Long.toHexString(curZookeeper.getSessionId()); } /** * Expire current zk session and wait for {@link IZkStateListener#handleNewSession()} invoked * @param zkClient * @throws Exception */ public static void expireSession(final ZkClient zkClient) throws Exception { final CountDownLatch waitNewSession = new CountDownLatch(1); IZkStateListener listener = new IZkStateListener() { @Override public void handleStateChanged(KeeperState state) throws Exception { LOG.info("IZkStateListener#handleStateChanged, state: " + state); } @Override public void handleNewSession() throws Exception { // make sure zkclient is connected again zkClient.waitUntilConnected(); ZkConnection connection = ((ZkConnection) zkClient.getConnection()); ZooKeeper curZookeeper = connection.getZookeeper(); LOG.info("handleNewSession. sessionId: " + Long.toHexString(curZookeeper.getSessionId())); waitNewSession.countDown(); } }; zkClient.subscribeStateChanges(listener); ZkConnection connection = ((ZkConnection) zkClient.getConnection()); ZooKeeper curZookeeper = connection.getZookeeper(); String oldSessionId = Long.toHexString(curZookeeper.getSessionId()); LOG.info("Before session expiry. sessionId: " + oldSessionId + ", zk: " + curZookeeper); Watcher watcher = new Watcher() { @Override public void process(WatchedEvent event) { LOG.info("Watcher#process, event: " + event); } }; final ZooKeeper dupZookeeper = new ZooKeeper(connection.getServers(), curZookeeper.getSessionTimeout(), watcher, curZookeeper.getSessionId(), curZookeeper.getSessionPasswd()); // wait until connected, then close while (dupZookeeper.getState() != States.CONNECTED) { Thread.sleep(10); } Assert.assertEquals(dupZookeeper.getState(), States.CONNECTED, "Fail to connect to zk using current session info"); dupZookeeper.close(); // make sure session expiry really happens waitNewSession.await(); zkClient.unsubscribeStateChanges(listener); connection = (ZkConnection) zkClient.getConnection(); curZookeeper = connection.getZookeeper(); String newSessionId = Long.toHexString(curZookeeper.getSessionId()); LOG.info("After session expiry. sessionId: " + newSessionId + ", zk: " + curZookeeper); Assert.assertNotSame(newSessionId, oldSessionId, "Fail to expire current session, zk: " + curZookeeper); } /** * expire zk session asynchronously * @param zkClient * @throws Exception */ public static void asyncExpireSession(final ZkClient zkClient) throws Exception { ZkConnection connection = ((ZkConnection) zkClient.getConnection()); ZooKeeper curZookeeper = connection.getZookeeper(); LOG.info("Before expiry. sessionId: " + Long.toHexString(curZookeeper.getSessionId())); Watcher watcher = new Watcher() { @Override public void process(WatchedEvent event) { LOG.info("Process watchEvent: " + event); } }; final ZooKeeper dupZookeeper = new ZooKeeper(connection.getServers(), curZookeeper.getSessionTimeout(), watcher, curZookeeper.getSessionId(), curZookeeper.getSessionPasswd()); // wait until connected, then close while (dupZookeeper.getState() != States.CONNECTED) { Thread.sleep(10); } dupZookeeper.close(); connection = (ZkConnection) zkClient.getConnection(); curZookeeper = connection.getZookeeper(); // System.err.println("zk: " + oldZookeeper); LOG.info("After expiry. sessionId: " + Long.toHexString(curZookeeper.getSessionId())); } /* * stateMap: partition->instance->state */ public static boolean verifyState(ZkClient zkclient, String clusterName, String resourceName, Map<String, Map<String, String>> expectStateMap, String op) { boolean result = true; ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(zkclient); ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor); Builder keyBuilder = accessor.keyBuilder(); ExternalView extView = accessor.getProperty(keyBuilder.externalView(resourceName)); Map<String, Map<String, String>> actualStateMap = extView.getRecord().getMapFields(); for (String partition : actualStateMap.keySet()) { for (String expectPartiton : expectStateMap.keySet()) { if (!partition.matches(expectPartiton)) { continue; } Map<String, String> actualInstanceStateMap = actualStateMap.get(partition); Map<String, String> expectInstanceStateMap = expectStateMap.get(expectPartiton); for (String instance : actualInstanceStateMap.keySet()) { for (String expectInstance : expectStateMap.get(expectPartiton).keySet()) { if (!instance.matches(expectInstance)) { continue; } String actualState = actualInstanceStateMap.get(instance); String expectState = expectInstanceStateMap.get(expectInstance); boolean equals = expectState.equals(actualState); if (op.equals("==") && !equals || op.equals("!=") && equals) { System.out.println(partition + "/" + instance + " state mismatch. actual state: " + actualState + ", but expect: " + expectState + ", op: " + op); result = false; } } } } } return result; } /** * return the number of listeners on given zk-path * @param zkAddr * @param path * @return * @throws Exception */ public static int numberOfListeners(String zkAddr, String path) throws Exception { Map<String, Set<String>> listenerMap = getListenersByZkPath(zkAddr); if (listenerMap.containsKey(path)) { return listenerMap.get(path).size(); } return 0; } /** * return a map from zk-path to a set of zk-session-id that put watches on the zk-path * @param zkAddr * @return * @throws Exception */ public static Map<String, Set<String>> getListenersByZkPath(String zkAddr) throws Exception { String splits[] = zkAddr.split(":"); Map<String, Set<String>> listenerMap = new TreeMap<String, Set<String>>(); Socket sock = null; int retry = 5; while (retry > 0) { try { sock = new Socket(splits[0], Integer.parseInt(splits[1])); PrintWriter out = new PrintWriter(sock.getOutputStream(), true); BufferedReader in = new BufferedReader(new InputStreamReader(sock.getInputStream())); out.println("wchp"); listenerMap.clear(); String lastPath = null; String line = in.readLine(); while (line != null) { line = line.trim(); if (line.startsWith("/")) { lastPath = line; if (!listenerMap.containsKey(lastPath)) { listenerMap.put(lastPath, new TreeSet<String>()); } } else if (line.startsWith("0x")) { if (lastPath != null && listenerMap.containsKey(lastPath)) { listenerMap.get(lastPath).add(line); } else { LOG.error("Not path associated with listener sessionId: " + line + ", lastPath: " + lastPath); } } else { // LOG.error("unrecognized line: " + line); } line = in.readLine(); } break; } catch (Exception e) { // sometimes in test, we see connection-reset exceptions when in.readLine() // so add this retry logic retry--; } finally { if (sock != null) sock.close(); } } return listenerMap; } /** * return a map from session-id to a set of zk-path that the session has watches on * @return */ public static Map<String, Set<String>> getListenersBySession(String zkAddr) throws Exception { Map<String, Set<String>> listenerMapByInstance = getListenersByZkPath(zkAddr); // convert to index by sessionId Map<String, Set<String>> listenerMapBySession = new TreeMap<String, Set<String>>(); for (String path : listenerMapByInstance.keySet()) { for (String sessionId : listenerMapByInstance.get(path)) { if (!listenerMapBySession.containsKey(sessionId)) { listenerMapBySession.put(sessionId, new TreeSet<String>()); } listenerMapBySession.get(sessionId).add(path); } } return listenerMapBySession; } static java.lang.reflect.Field getField(Class<?> clazz, String fieldName) throws NoSuchFieldException { try { return clazz.getDeclaredField(fieldName); } catch (NoSuchFieldException e) { Class<?> superClass = clazz.getSuperclass(); if (superClass == null) { throw e; } else { return getField(superClass, fieldName); } } } public static Map<String, List<String>> getZkWatch(ZkClient client) throws Exception { Map<String, List<String>> lists = new HashMap<String, List<String>>(); ZkConnection connection = ((ZkConnection) client.getConnection()); ZooKeeper zk = connection.getZookeeper(); java.lang.reflect.Field field = getField(zk.getClass(), "watchManager"); field.setAccessible(true); Object watchManager = field.get(zk); java.lang.reflect.Field field2 = getField(watchManager.getClass(), "dataWatches"); field2.setAccessible(true); HashMap<String, Set<Watcher>> dataWatches = (HashMap<String, Set<Watcher>>) field2.get(watchManager); field2 = getField(watchManager.getClass(), "existWatches"); field2.setAccessible(true); HashMap<String, Set<Watcher>> existWatches = (HashMap<String, Set<Watcher>>) field2.get(watchManager); field2 = getField(watchManager.getClass(), "childWatches"); field2.setAccessible(true); HashMap<String, Set<Watcher>> childWatches = (HashMap<String, Set<Watcher>>) field2.get(watchManager); lists.put("dataWatches", new ArrayList<String>(dataWatches.keySet())); lists.put("existWatches", new ArrayList<String>(existWatches.keySet())); lists.put("childWatches", new ArrayList<String>(childWatches.keySet())); return lists; } public static Map<String, Set<IZkDataListener>> getZkDataListener(ZkClient client) throws Exception { java.lang.reflect.Field field = getField(client.getClass(), "_dataListener"); field.setAccessible(true); Map<String, Set<IZkDataListener>> dataListener = (Map<String, Set<IZkDataListener>>) field.get(client); return dataListener; } public static Map<String, Set<IZkChildListener>> getZkChildListener(ZkClient client) throws Exception { java.lang.reflect.Field field = getField(client.getClass(), "_childListener"); field.setAccessible(true); Map<String, Set<IZkChildListener>> childListener = (Map<String, Set<IZkChildListener>>) field.get(client); return childListener; } public static boolean tryWaitZkEventsCleaned(ZkClient zkclient) throws Exception { java.lang.reflect.Field field = getField(zkclient.getClass(), "_eventThread"); field.setAccessible(true); Object eventThread = field.get(zkclient); // System.out.println("field: " + eventThread); java.lang.reflect.Field field2 = getField(eventThread.getClass(), "_events"); field2.setAccessible(true); BlockingQueue queue = (BlockingQueue) field2.get(eventThread); // System.out.println("field2: " + queue + ", " + queue.size()); if (queue == null) { LOG.error("fail to get event-queue from zkclient. skip waiting"); return false; } for (int i = 0; i < 20; i++) { if (queue.size() == 0) { return true; } Thread.sleep(100); System.out.println("pending zk-events in queue: " + queue); } return false; } }