/*
* Copyright (C) 2012 Facebook, Inc.
*
* Licensed 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 com.facebook.zookeeper.cmd;
import com.facebook.util.StreamImporter;
import com.facebook.zookeeper.ZkUtil;
import com.facebook.zookeeper.path.ZkGenericPath;
import com.facebook.zookeeper.ZooKeeperIface;
import com.facebook.zookeeper.convenience.ZkQuickConnectionManager;
import com.facebook.zookeeper.convenience.ZkScript;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Options;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.data.Stat;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
public class ZNodePruner extends ZkScript {
public static final String DEFAULT_KEYWORD = "freeze";
private volatile boolean verbose = false;
public ZNodePruner(ZkQuickConnectionManager zkQuickConnectionManager) {
super(zkQuickConnectionManager);
}
public ZNodePruner() {
this(new ZkQuickConnectionManager());
}
public void setVerbose(boolean verbose) {
this.verbose = verbose;
}
public void prunePersistent(String pathStr, String keyword, File saveTemplate)
throws IOException, InterruptedException, KeeperException {
InputStream in = new FileInputStream(saveTemplate);
try {
prunePersistent(pathStr, keyword, in);
} finally {
in.close();
}
}
public void prunePersistent(String pathStr, String keyword, InputStream in)
throws IOException, InterruptedException, KeeperException {
prunePersistent(pathStr, keyword, StreamImporter.importLines(in));
}
public void prunePersistent(
String pathStr, String keyword, List<String> toKeep
) throws InterruptedException, KeeperException {
ZkGenericPath path = new ZkGenericPath(pathStr);
Set<String> keepSet = expandToKeepSet(toKeep);
internalPrunePersistent(path, keyword, keepSet);
}
private void internalPrunePersistent(
ZkGenericPath path, String keyword, Set<String> keepSet
) throws InterruptedException, KeeperException {
if (isEphemeral(path.toString())) {
// We should only be scanning persistent nodes (although we may delete
// ephemeral nodes in order to remove a parent persistent node)
return;
}
if (keepSet.contains(path.toString())) {
try {
ZooKeeperIface zk = getZk();
byte[] data = zk.getData(path.toString(), false, null);
if (data != null &&
keyword.equals(ZkUtil.bytesToString(data))
) {
zk.setData(path.toString(), new byte[0], -1);
}
List<String> children = zk.getChildren(path.toString(), null);
for (String child : children) {
internalPrunePersistent(path.appendChild(child), keyword, keepSet);
}
} catch (KeeperException.NoNodeException e) {
// If the node disappears while scanning it, then just ignore
}
} else {
deleteSubtree(path, keyword);
}
}
private void deleteSubtree(ZkGenericPath path, String keyword)
throws InterruptedException, KeeperException {
ZooKeeperIface zk = getZk();
try {
if (!isEphemeral(path.toString())) {
// Only set freeze on non-ephemerals since ephemerals cant have children
zk.setData(
path.toString(),
ZkUtil.stringToBytes(keyword),
-1
);
}
while (true) {
List<String> children = zk.getChildren(path.toString(), null);
for (String child : children) {
deleteSubtree(path.appendChild(child), keyword);
}
try {
zk.delete(path.toString(), -1);
break;
} catch (KeeperException.NotEmptyException e) {
// Repeat since children were re-added
}
}
} catch (KeeperException.NoNodeException e) {
// Ignore since we are trying to delete it
}
if (verbose) {
System.out.println("Deleted ZNode: " + path);
}
}
private boolean isEphemeral(String pathStr)
throws InterruptedException, KeeperException {
ZooKeeperIface zk = getZk();
Stat stat = new Stat();
zk.getData(pathStr, null, stat);
return stat.getEphemeralOwner() != 0;
}
private Set<String> expandToKeepSet(List<String> toKeep) {
// Add all listed ZNodes as well as all of their ancestors
Set<String> keepSet = new HashSet<String>();
for (String pathStr : toKeep) {
ZkGenericPath path = ZkGenericPath.parse("/", pathStr);
Iterator<ZkGenericPath> lineageIter = path.lineageIterator();
while(lineageIter.hasNext()) {
keepSet.add(lineageIter.next().toString());
}
}
return keepSet;
}
@Override
protected String getName() {
return ZNodePruner.class.getName();
}
@Override
protected Options getSpecificOptions() {
Options options = new Options();
options.addOption(
"z",
"zkpath",
true,
"ZooKeeper path to prune (includes all descendants) [Required]"
);
options.addOption(
"k",
"keyword",
true,
"Keyword to write into zNodes that will be removed. This may be used " +
"to help reduce activity on or underneath this node [Default: freeze]"
);
options.addOption(
"f",
"file-template",
true,
"Path to file containing new-line delimited list of ZNodes to save. " +
"The list only applies to persistent ZNodes. If this parameter is " +
"not specified, expects the values to be provided via standard input."
);
options.addOption(
"v",
"verbose",
false,
"Print verbose messages [Default: off]"
);
return options;
}
@Override
protected boolean verifySpecificOptions(CommandLine cmd) {
setVerbose(cmd.hasOption("verbose"));
if (!cmd.hasOption("zkpath")) {
System.err.println("Error: You must specify a ZooKeeper path.\n");
return false;
}
if (cmd.hasOption("file-template")) {
File template = new File(cmd.getOptionValue("file-template"));
if (!template.exists()) {
System.err.println("Error: invalid file-template path.\n");
return false;
}
}
return true;
}
@Override
protected void runScript(CommandLine cmd) throws Exception {
String root = cmd.getOptionValue("zkpath");
String keyword = cmd.getOptionValue("keyword", ZNodePruner.DEFAULT_KEYWORD);
if (cmd.hasOption("file-template")) {
File saveTemplate = new File(cmd.getOptionValue("file-template"));
prunePersistent(root, keyword, saveTemplate);
} else {
prunePersistent(root, keyword, System.in);
}
}
public static void main(String[] args) throws Exception {
ZkScript script = new ZNodePruner();
script.runMain(args);
System.out.println("DONE");
}
}