/**
*
* 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.
*
*/
package org.apache.bookkeeper.zookeeper;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.bookkeeper.stats.NullStatsLogger;
import org.apache.bookkeeper.test.ZooKeeperUtil;
import org.apache.zookeeper.AsyncCallback.Create2Callback;
import org.apache.zookeeper.AsyncCallback.Children2Callback;
import org.apache.zookeeper.AsyncCallback.DataCallback;
import org.apache.zookeeper.AsyncCallback.StatCallback;
import org.apache.zookeeper.AsyncCallback.StringCallback;
import org.apache.zookeeper.AsyncCallback.VoidCallback;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
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.ZooDefs.Ids;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.junit.Assert;
import junit.framework.TestCase;
/**
* Test the wrapper of {@link org.apache.zookeeper.ZooKeeper} client.
*/
public class TestZooKeeperClient extends TestCase {
static final Logger logger = LoggerFactory.getLogger(TestZooKeeperClient.class);
// ZooKeeper related variables
protected ZooKeeperUtil zkUtil = new ZooKeeperUtil();
@Before
@Override
public void setUp() throws Exception {
logger.info("Setting up test {}.", getName());
zkUtil.startServer();
}
@After
@Override
public void tearDown() throws Exception {
zkUtil.killServer();
logger.info("Teared down test {}.", getName());
}
private void expireZooKeeperSession(ZooKeeper zk, int timeout)
throws IOException, InterruptedException, KeeperException {
final CountDownLatch latch = new CountDownLatch(1);
ZooKeeper newZk = new ZooKeeper(zkUtil.getZooKeeperConnectString(), timeout,
new Watcher() {
@Override
public void process(WatchedEvent event) {
if (event.getType() == EventType.None &&
event.getState() == KeeperState.SyncConnected) {
latch.countDown();
}
}
}, zk.getSessionId(), zk.getSessionPasswd());
if (!latch.await(timeout, TimeUnit.MILLISECONDS)) {
throw KeeperException.create(KeeperException.Code.CONNECTIONLOSS);
}
newZk.close();
}
/**
* Shutdown Zk Server when client received an expire event.
* So the client issue recreation client task but it would not succeed
* until we start the zookeeper server again.
*/
class ShutdownZkServerClient extends ZooKeeperClient {
ShutdownZkServerClient(String connectString, int sessionTimeoutMs,
ZooKeeperWatcherBase watcher, RetryPolicy operationRetryPolicy)
throws IOException {
super(connectString, sessionTimeoutMs, watcher,
new BoundExponentialBackoffRetryPolicy(sessionTimeoutMs, sessionTimeoutMs, Integer.MAX_VALUE),
operationRetryPolicy,
NullStatsLogger.INSTANCE, 1, 0);
}
@Override
public void process(WatchedEvent event) {
if (event.getType() == EventType.None &&
event.getState() == KeeperState.Expired) {
try {
zkUtil.stopServer();
} catch (Exception e) {
logger.error("Failed to stop zookeeper server : ", e);
}
}
super.process(event);
}
}
@Test(timeout=12000)
public void testReconnectAfterExipred() throws Exception {
final CountDownLatch expireLatch = new CountDownLatch(1);
Watcher testWatcher = new Watcher() {
@Override
public void process(WatchedEvent event) {
if (event.getType() == EventType.None &&
event.getState() == KeeperState.Expired) {
expireLatch.countDown();
}
}
};
final int timeout = 2000;
ZooKeeperWatcherBase watcherManager =
new ZooKeeperWatcherBase(timeout).addChildWatcher(testWatcher);
List<Watcher> watchers = new ArrayList<Watcher>(1);
watchers.add(testWatcher);
ZooKeeperClient client = new ShutdownZkServerClient(
zkUtil.getZooKeeperConnectString(), timeout, watcherManager,
new BoundExponentialBackoffRetryPolicy(timeout, timeout, 0)
);
client.waitForConnection();
Assert.assertTrue("Client failed to connect an alive ZooKeeper.",
client.getState().isConnected());
logger.info("Expire zookeeper client");
expireZooKeeperSession(client, timeout);
// wait until session expire
Assert.assertTrue("Client registered watcher should receive expire event.",
expireLatch.await(2 * timeout, TimeUnit.MILLISECONDS));
Assert.assertFalse("Client doesn't receive expire event from ZooKeeper.",
client.getState().isConnected());
try {
client.exists("/tmp", false);
Assert.fail("Should fail due to connection loss.");
} catch (KeeperException.ConnectionLossException cle) {
// expected
} catch (KeeperException.SessionExpiredException cle) {
// expected
}
zkUtil.restartServer();
// wait for a reconnect cycle
Thread.sleep(2*timeout);
Assert.assertTrue("Client failed to connect zookeeper even it was back.",
client.getState().isConnected());
try {
client.exists("/tmp", false);
} catch (KeeperException.ConnectionLossException cle) {
Assert.fail("Should not throw ConnectionLossException");
} catch (KeeperException.SessionExpiredException cle) {
Assert.fail("Should not throw SessionExpiredException");
}
}
@Test(timeout=60000)
public void testRetrySyncOperations() throws Exception {
final int timeout = 2000;
ZooKeeperClient client = ZooKeeperClient.createConnectedZooKeeperClient(
zkUtil.getZooKeeperConnectString(), timeout, new HashSet<Watcher>(),
new BoundExponentialBackoffRetryPolicy(timeout, timeout, Integer.MAX_VALUE)
);
Assert.assertTrue("Client failed to connect an alive ZooKeeper.",
client.getState().isConnected());
String path = "/a";
byte[] data = "test".getBytes();
expireZooKeeperSession(client, timeout);
logger.info("Create znode " + path);
client.create(path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
logger.info("Created znode " + path);
expireZooKeeperSession(client, timeout);
logger.info("Exists znode " + path);
Stat stat = client.exists(path, false);
Assert.assertNotNull("znode doesn't existed", stat);
expireZooKeeperSession(client, timeout);
logger.info("Get data from znode " + path);
Stat newStat = new Stat();
client.getData(path, false, newStat);
Assert.assertEquals(stat, newStat);
expireZooKeeperSession(client, timeout);
logger.info("Create children under znode " + path);
client.create(path + "/children", data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
expireZooKeeperSession(client, timeout);
logger.info("Create children under znode " + path);
client.create(path + "/children2", data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, new Stat());
expireZooKeeperSession(client, timeout);
List<String> children = client.getChildren(path, false, newStat);
Assert.assertEquals(2, children.size());
Assert.assertTrue(children.contains("children"));
Assert.assertTrue(children.contains("children2"));
logger.info("Get children under znode " + path);
expireZooKeeperSession(client, timeout);
client.delete(path + "/children", -1);
logger.info("Delete children from znode " + path);
}
@Test(timeout=60000)
public void testRetryOnCreatingEphemeralZnode() throws Exception {
final int timeout = 2000;
ZooKeeperClient client = ZooKeeperClient.createConnectedZooKeeperClient(
zkUtil.getZooKeeperConnectString(), timeout, new HashSet<Watcher>(),
new BoundExponentialBackoffRetryPolicy(timeout, timeout, Integer.MAX_VALUE)
);
Assert.assertTrue("Client failed to connect an alive ZooKeeper.",
client.getState().isConnected());
String path = "/a";
byte[] data = "test".getBytes();
logger.info("Create znode " + path);
client.create(path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
logger.info("Created znode " + path);
expireZooKeeperSession(client, timeout);
logger.info("Create znode w/ new session : " + path);
client.create(path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
logger.info("Created znode w/ new session : " + path);
}
@Test(timeout=60000)
public void testRetryAsyncOperations() throws Exception {
final int timeout = 2000;
ZooKeeperClient client = ZooKeeperClient.createConnectedZooKeeperClient(
zkUtil.getZooKeeperConnectString(), timeout, new HashSet<Watcher>(),
new BoundExponentialBackoffRetryPolicy(timeout, timeout, Integer.MAX_VALUE)
);
Assert.assertTrue("Client failed to connect an alive ZooKeeper.",
client.getState().isConnected());
String path = "/a";
byte[] data = "test".getBytes();
expireZooKeeperSession(client, timeout);
logger.info("Create znode " + path);
final CountDownLatch createLatch = new CountDownLatch(1);
client.create(path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT,
new StringCallback() {
@Override
public void processResult(int rc, String path, Object ctx, String name) {
if (KeeperException.Code.OK.intValue() == rc) {
createLatch.countDown();
}
}
}, null);
createLatch.await();
logger.info("Created znode " + path);
expireZooKeeperSession(client, timeout);
logger.info("Create znode " + path);
final CountDownLatch create2Latch = new CountDownLatch(1);
client.create(path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT,
new Create2Callback() {
@Override
public void processResult(int rc, String path, Object ctx, String name, Stat stat) {
if (KeeperException.Code.NODEEXISTS.intValue() == rc) {
create2Latch.countDown();
}
}
}, null);
create2Latch.await();
logger.info("Created znode " + path);
expireZooKeeperSession(client, timeout);
logger.info("Exists znode " + path);
final CountDownLatch existsLatch = new CountDownLatch(1);
client.exists(path, false, new StatCallback() {
@Override
public void processResult(int rc, String path, Object ctx, Stat stat) {
if (KeeperException.Code.OK.intValue() == rc) {
existsLatch.countDown();
}
}
}, null);
existsLatch.await();
expireZooKeeperSession(client, timeout);
final CountDownLatch getLatch = new CountDownLatch(1);
logger.info("Get data from znode " + path);
client.getData(path, false, new DataCallback() {
@Override
public void processResult(int rc, String path, Object ctx, byte[] data, Stat stat) {
if (KeeperException.Code.OK.intValue() == rc) {
getLatch.countDown();
}
}
}, null);
getLatch.await();
expireZooKeeperSession(client, timeout);
logger.info("Create children under znode " + path);
final CountDownLatch createChildLatch = new CountDownLatch(1);
client.create(path + "/children", data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT,
new StringCallback() {
@Override
public void processResult(int rc, String path, Object ctx, String name) {
if (KeeperException.Code.OK.intValue() == rc) {
createChildLatch.countDown();
}
}
}, null);
createChildLatch.await();
expireZooKeeperSession(client, timeout);
final CountDownLatch getChildLatch = new CountDownLatch(1);
final AtomicReference<List<String>> children =
new AtomicReference<List<String>>();
client.getChildren(path, false, new Children2Callback() {
@Override
public void processResult(int rc, String path, Object ctx, List<String> childList, Stat stat) {
if (KeeperException.Code.OK.intValue() == rc) {
children.set(childList);
getChildLatch.countDown();
}
}
}, null);
getChildLatch.await();
Assert.assertNotNull(children.get());
Assert.assertEquals(1, children.get().size());
Assert.assertEquals("children", children.get().get(0));
logger.info("Get children under znode " + path);
expireZooKeeperSession(client, timeout);
final CountDownLatch deleteChildLatch = new CountDownLatch(1);
client.delete(path + "/children", -1, new VoidCallback() {
@Override
public void processResult(int rc, String path, Object ctx) {
if (KeeperException.Code.OK.intValue() == rc) {
deleteChildLatch.countDown();
}
}
}, null);
deleteChildLatch.await();
logger.info("Delete children from znode " + path);
}
}