ZookeeperClient.java

package io.fathom.cloud.zookeeper;

import java.io.IOException;
import java.util.List;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.KeeperException.Code;
import org.apache.zookeeper.KeeperException.NoNodeException;
import org.apache.zookeeper.KeeperException.NodeExistsException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.CharMatcher;
import com.google.common.util.concurrent.SettableFuture;

public class ZookeeperClient {
    private static final Logger log = LoggerFactory.getLogger(ZookeeperClient.class);

    ZooKeeper zk;

    private final String connectString;

    public ZookeeperClient(String connectString) {
        log.info("Building zookeeper client to {}", connectString);
        this.connectString = connectString;
    }

    protected synchronized ZooKeeper getZk() throws IOException {
        ZooKeeper zk = this.zk;
        if (zk == null) {
            int sessionTimeout = 20000;
            Watcher watcher = new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    log.debug("Got ZK event {}", event);
                }
            };
            try {
                zk = new ZooKeeper(connectString, sessionTimeout, watcher);
            } catch (IOException e) {
                throw new IOException("Unable to connect to zookeeper", e);
            }
            this.zk = zk;
        }
        return zk;
    }

    protected synchronized ZooKeeper closeZk() {
        ZooKeeper zk = this.zk;
        if (zk != null) {
            this.zk = null;

            try {
                zk.close();
            } catch (Exception e) {
                log.warn("Error closing zk", e);
            }
        }
        return zk;
    }

    public Stat exists(String path, boolean watch) throws KeeperException, IOException {
        try {
            return getZk().exists(path, watch);
        } catch (KeeperException e) {
            if (processException(e)) {
                return exists(path, watch);
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted while communicating with zookeeper", e);
        }
    }

    public void delete(String path, int version) throws KeeperException, IOException {
        try {
            getZk().delete(path, version);
        } catch (KeeperException e) {
            if (processException(e)) {
                delete(path, version);
                return;
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted while communicating with zookeeper", e);
        }
    }

    public List<String> getChildren(String path, boolean watch) throws KeeperException, IOException {
        try {
            return getZk().getChildren(path, watch);
        } catch (KeeperException e) {
            if (processException(e)) {
                return getChildren(path, watch);
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted while communicating with zookeeper", e);
        }
    }

    public byte[] getData(String path, SettableFuture<Object> watch, Stat stat) throws IOException, KeeperException {
        try {
            Watcher watcher = null;

            if (watch != null) {
                watcher = new ListenableWatcher(watch);
            }

            return getZk().getData(path, watcher, stat);
        } catch (KeeperException e) {
            if (processException(e)) {
                return getData(path, watch, stat);
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted while communicating with zookeeper", e);
        }
    }

    public Stat setData(String path, byte[] data, int version) throws KeeperException, IOException {
        try {
            return getZk().setData(path, data, version);
        } catch (KeeperException e) {
            if (processException(e)) {
                return setData(path, data, version);
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted while communicating with zookeeper", e);
        }
    }

    private boolean processException(KeeperException e) {
        Code code = e.code();
        switch (code) {
        case SESSIONEXPIRED:
            log.warn("Closing ZK session after SESSIONEXPIRED");
            closeZk();
            return true;
        case CONNECTIONLOSS:
            log.warn("Closing ZK session after CONNECTIONLOSS");
            closeZk();
            return true;

        default:
            return false;
        }
    }

    public String create(final String path, byte[] data, List<ACL> acl, CreateMode createMode) throws IOException,
            KeeperException {
        try {
            return getZk().create(path, data, acl, createMode);
        } catch (KeeperException e) {
            if (processException(e)) {
                return create(path, data, acl, createMode);
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted while communicating with zookeeper", e);
        }
    }

    public void mkdirs(String path) throws KeeperException, IOException {
        try {
            log.debug("ZK create on {}", path);

            getZk().create(path, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        } catch (NoNodeException e) {
            mkdirs(getParent(path));
            mkdirs(path);
        } catch (NodeExistsException e) {
            return;
        } catch (KeeperException e) {
            if (processException(e)) {
                mkdirs(path);
                return;
            }
            throw e;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException("Error communicating with zookeeper", e);
        }
    }

    public static String getParent(String zkPath) {
        zkPath = CharMatcher.is('/').trimTrailingFrom(zkPath);

        int lastSlash = zkPath.lastIndexOf('/');
        if (lastSlash == -1) {
            throw new IllegalArgumentException();
        }
        return zkPath.substring(0, lastSlash);
    }

    public void createOrUpdate(String path, byte[] data, boolean mkdirs) throws KeeperException, IOException {
        try {
            setData(path, data, -1);
        } catch (NoNodeException e) {
            create(path, data, mkdirs);
        }
    }

    public String create(final String path, byte[] data, boolean mkdirs) throws IOException, KeeperException {
        try {
            return create(path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        } catch (NoNodeException e) {
            if (mkdirs) {
                mkdirs(getParent(path));
                return create(path, data, false);
            } else {
                throw e;
            }
        }
    }

    public static String escape(String s) {
        StringBuilder sb = new StringBuilder(s.length());

        int length = s.length();
        for (int i = 0; i < length; i++) {
            char c = s.charAt(i);

            if ((c >= '0' && c <= '9') || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')) {
                sb.append(c);
            } else {
                switch (c) {
                case '.':
                case ':':
                case '-':
                    sb.append(c);
                    break;

                default: {
                    String v = Integer.toHexString(c);
                    switch (v.length()) {
                    case 1:
                        sb.append("_0" + v);
                        break;

                    case 2:
                        sb.append("_" + v);
                        break;

                    case 3:
                        sb.append("__0" + v);
                        break;

                    case 4:
                        sb.append("__" + v);
                        break;

                    default:
                        throw new UnsupportedOperationException();
                    }
                }
                    break;
                }
            }
        }

        return sb.toString();
    }

}