/** * ZKBlockingQueue.java * * Copyright 2016 the original author or authors. * * We 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. */ package org.apache.niolex.queue; import java.io.Closeable; import java.io.IOException; import java.io.Serializable; import java.util.Arrays; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import org.apache.niolex.commons.bean.BeanUtil; import org.apache.niolex.commons.collection.CollectionUtil; import org.apache.niolex.zookeeper.core.ZKConnector; import org.apache.niolex.zookeeper.core.ZKException; import org.apache.zookeeper.KeeperException; import org.apache.zookeeper.WatchedEvent; import org.apache.zookeeper.Watcher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * The distributed blocking queue implementation powered by Zookeeper.<br>If you call the constructor * {@link #ZKBlockingQueue(String, int, String)} to create a new instance, you need to call the {@link #close()} * method on this instance to close the inner ZKConnector when you do not want to use this lock any more. * <br>If you call constructor {@link #ZKBlockingQueue(ZKConnector, String)}, the ZKConnector instance will be managed * by you. Even if you call {@link #close()}, we will not close the ZKConnector instance. * * @author <a href="mailto:xiejiyun@foxmail.com">Xie, Jiyun</a> * @version 1.0.0 * @since 2016-4-19 */ public class ZKBlockingQueue<E extends Serializable> extends DistributedBlockingQueue<E> implements Closeable { protected static final Logger LOG = LoggerFactory.getLogger(ZKBlockingQueue.class); protected static final String PREFIX = "/zkbq-"; private static final String[] STUB = new String[0]; // Used to store current parsed items. private final LinkedBlockingQueue<String> itemQueue = new LinkedBlockingQueue<String>(); private final ZKConnector zkc; private final String basePath; private boolean closeZKC = false; /** * The Constructor to create a {@link ZKConnector} inside it. * * @param clusterAddress the zookeeper cluster servers address list * @param sessionTimeout the zookeeper session timeout in microseconds * @param basePath the queue base path * @throws IOException in cases of network failure * @throws IllegalArgumentException if sessionTimeout is too small */ public ZKBlockingQueue(String clusterAddress, int sessionTimeout, String basePath) throws IOException { this(new ZKConnector(clusterAddress, sessionTimeout), basePath); closeZKC = true; } /** * The ZKBlockingQueue Constructor. * * @param zkc the zookeeper connector * @param basePath the queue base path */ public ZKBlockingQueue(ZKConnector zkc, String basePath) { super(); this.zkc = zkc; if (basePath.endsWith("/")) { basePath = basePath.substring(0, basePath.length() - 1); } this.basePath = basePath; zkc.makeSurePathExists(basePath); } /** * Add authenticate info for this client. * 添加client的权限认证信息 * * @param username the user name of this client * @param password the password of this client * @see org.apache.niolex.zookeeper.core.ZKConnector#addAuthInfo(java.lang.String, java.lang.String) */ public void addAuthInfo(String username, String password) { zkc.addAuthInfo(username, password); } /** * This is the override of super method. * @see java.util.concurrent.BlockingQueue#offer(java.lang.Object) */ @Override public boolean offer(E e) { byte[] data; try { data = BeanUtil.toBytes(e); } catch (Exception ex) { throw new IllegalArgumentException(ex); } zkc.createNode(basePath + PREFIX, data, false, true); return true; } /** * Get children of the base path, No sort. * * @return the children */ protected List<String> getChildren() { return zkc.getChildren(basePath); } /** * Get the real data from ZK and store them in the {@link #itemQueue}. We do sort here. * * @param children the children of the base path */ protected void parseDataFromZK(List<String> children) { // Sort itemArray. String[] data = children.toArray(STUB); Arrays.sort(data); for (String s : data) { itemQueue.offer(s); } } /** * This is the override of super method. * @see java.util.Queue#peek() */ @Override public E peek() { return getHead(false); } /** * This is the override of super method. * @see java.util.Queue#poll() */ @Override public E poll() { return getHead(true); } /** * Get the current queue head. Remove it from queue if the parameter <tt>removeItem</tt> is {@code true}. * * @param removeItem whether we need to remove the item from the queue or not * @return the head item or null if queue is empty */ public synchronized E getHead(boolean removeItem) { boolean queryZK = false; String child = null; byte[] arr = null; while (true) { // Get item from queue. if (removeItem) child = itemQueue.poll(); else child = itemQueue.peek(); // Get item from ZK only once. if (child == null) { if (!queryZK) { queryZK = true; parseDataFromZK(getChildren()); // After query, poll queue again. continue; } else { return null; } } // Make whole path. String path = basePath + "/" + child; // Get data from ZK. try { arr = zkc.getData(path); if (removeItem) zkc.deleteNode(path); // We got the data. break; } catch (ZKException zk) { if (zk.getCode() != ZKException.Code.NO_NODE) { throw zk; } else if (!removeItem) { // Item already removed from ZK, we need to remove it from itemQueue. itemQueue.remove(child); } } } // De-serialize data to object. if (arr == null) { return null; } else { try { @SuppressWarnings("unchecked") E o = (E) BeanUtil.toObject(arr); return o; } catch (Exception e) { throw new IllegalStateException("Failed to deserialize object for " + child, e); } } } /** * This is the override of super method. * @see org.apache.niolex.queue.DistributedBlockingQueue#watchQueue() */ @Override protected synchronized void watchQueue() throws InterruptedException { try { CountDownLatch latch = new CountDownLatch(1); List<String> children = zkc.zooKeeper().getChildren(basePath, new ChildrenChangeWather(latch)); if (!CollectionUtil.isEmpty(children)) { parseDataFromZK(children); return; } else { latch.await(); } } catch (KeeperException e) { throw ZKException.makeInstance("ZKBlockingQueue internal error.", e); } } /** * This is the override of super method. * @see org.apache.niolex.queue.DistributedBlockingQueue#watchQueue(long, java.util.concurrent.TimeUnit) */ @Override protected synchronized boolean watchQueue(long timeout, TimeUnit unit) throws InterruptedException { try { CountDownLatch latch = new CountDownLatch(1); List<String> children = zkc.zooKeeper().getChildren(basePath, new ChildrenChangeWather(latch)); if (!CollectionUtil.isEmpty(children)) { parseDataFromZK(children); return true; } else { return latch.await(timeout, unit); } } catch (KeeperException e) { throw ZKException.makeInstance("ZKBlockingQueue internal error.", e); } } /** * Check the children changes of the base path. * * @author <a href="mailto:xiejiyun@foxmail.com">Xie, Jiyun</a> * @version 1.0.0 * @since 2016-4-21 */ public class ChildrenChangeWather implements Watcher { // The latch used to wait for the children changes. private final CountDownLatch latch; /** * Constructor. * * @param latch the latch used to notify event */ public ChildrenChangeWather(CountDownLatch latch) { super(); this.latch = latch; } /** * This is the override of super method. * @see org.apache.zookeeper.Watcher#process(org.apache.zookeeper.WatchedEvent) */ @Override public void process(WatchedEvent event) { // Each watch can be triggered only once. // We need to make sure do this thing right. try { if (event.getType() == Watcher.Event.EventType.None) { // If event type is none, then something wrong with the zookeeper. while (!zkc.connected()) zkc.waitForConnectedTillDeath(); } } finally { latch.countDown(); } } } /** * This is the override of super method. * @see java.util.AbstractCollection#size() */ @Override public int size() { return getChildren().size(); } /** * Close the internal zookeeper connector if and only if it's created by this class. * * This is the override of super method. * @see java.io.Closeable#close() */ @Override public void close() { if (closeZKC) { zkc.close(); } } }