All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.I0Itec.zkclient.InMemoryConnection Maven / Gradle / Ivy

There is a newer version: 2024.03.7
Show newest version
/**
 * Copyright 2010 the original author or authors.
 *
 * 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 org.I0Itec.zkclient;

import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.I0Itec.zkclient.exception.ZkException;
import org.I0Itec.zkclient.exception.ZkInterruptedException;
import org.I0Itec.zkclient.exception.ZkNoNodeException;
import org.I0Itec.zkclient.util.ZkPathUtil;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.KeeperException.Code;
import org.apache.zookeeper.Op;
import org.apache.zookeeper.OpResult;
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;
import org.apache.zookeeper.ZooKeeper.States;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Id;
import org.apache.zookeeper.data.Stat;
import org.apache.zookeeper.proto.CheckVersionRequest;
import org.apache.zookeeper.proto.CreateRequest;
import org.apache.zookeeper.proto.DeleteRequest;
import org.apache.zookeeper.proto.SetDataRequest;

/**
 * Emulating a ZooKeeper server with few hash tables. Basically a mock class used for testing. Please avoid using this
 * as your ZK in production :)
 * 
 * Note that the addAuth is even more mocked than usual Since we have no authentication provider (i.e. Kerberos) around
 * we simply take the auth byte[] and convert it to string to get the Id scheme remains the same
 */
public class InMemoryConnection implements IZkConnection {

    public static class DataAndVersion {
        private byte[] _data;
        private int _version;
        private List _acl;

        public DataAndVersion(byte[] data, int version, List acl) {
            _data = data;
            _version = version;
            _acl = acl;
        }

        public DataAndVersion(byte[] data, int version) {
            this(data, version, null);
        }

        public byte[] getData() {
            return _data;
        }

        public int getVersion() {
            return _version;
        }

        public List getAcl() {
            return _acl;
        }
    }

    private Lock _lock = new ReentrantLock(true);
    private Map _data = new HashMap();
    private Map _creationTime = new HashMap();
    private List _ids = new ArrayList();
    private final AtomicInteger sequence = new AtomicInteger(0);

    private Set _dataWatches = new HashSet();
    private Set _nodeWatches = new HashSet();
    private EventThread _eventThread;

    private class EventThread extends Thread {

        private Watcher _watcher;
        private BlockingQueue _blockingQueue = new LinkedBlockingDeque();

        public EventThread(Watcher watcher) {
            _watcher = watcher;
        }

        @Override
        public void run() {
            try {
                while (true) {
                    _watcher.process(_blockingQueue.take());
                }
            } catch (InterruptedException e) {
                // stop event thread
            }
        }

        public void send(WatchedEvent event) {
            _blockingQueue.add(event);
        }
    }

    public InMemoryConnection() {
        try {
            create("/", null, CreateMode.PERSISTENT);
        } catch (KeeperException e) {
            throw ZkException.create(e);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new ZkInterruptedException(e);
        }
    }

    @Override
    public void close() throws InterruptedException {
        _lock.lockInterruptibly();
        try {
            if (_eventThread != null) {
                _eventThread.interrupt();
                _eventThread.join();
                _eventThread = null;
            }
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public void connect(Watcher watcher) {
        _lock.lock();
        try {
            if (_eventThread != null) {
                throw new IllegalStateException("Already connected.");
            }
            _eventThread = new EventThread(watcher);
            _eventThread.start();
            _eventThread.send(new WatchedEvent(null, KeeperState.SyncConnected, null));
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public String create(String path, byte[] data, List acl, CreateMode mode) throws KeeperException, InterruptedException {
        _lock.lock();
        try {

            if (mode.isSequential()) {
                final int newSequence = sequence.getAndIncrement();
                path = path + ZkPathUtil.leadingZeros(newSequence, 10);
            }

            if (exists(path, false)) {
                throw new KeeperException.NodeExistsException();
            }
            String parentPath = getParentPath(path);
            checkACL(parentPath, ZooDefs.Perms.CREATE);

            _data.put(path, new DataAndVersion(data, 0, acl));
            _creationTime.put(path, System.currentTimeMillis());
            checkWatch(_nodeWatches, path, EventType.NodeCreated);
            // we also need to send a child change event for the parent
            if (parentPath != null) {
                checkWatch(_nodeWatches, parentPath, EventType.NodeChildrenChanged);
            }
            return path;
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public String create(String path, byte[] data, CreateMode mode) throws KeeperException, InterruptedException {
        return create(path, data, null, mode);
    }

    private String getParentPath(String path) {
        int lastIndexOf = path.lastIndexOf("/");
        if (lastIndexOf == -1 || lastIndexOf == 0) {
            return null;
        }
        return path.substring(0, lastIndexOf);
    }

    @Override
    public void delete(String path) throws InterruptedException, KeeperException {
        _lock.lock();
        try {
            if (!exists(path, false)) {
                throw new KeeperException.NoNodeException();
            }
            String parentPath = getParentPath(path);
            checkACL(parentPath, ZooDefs.Perms.DELETE);
            _data.remove(path);
            _creationTime.remove(path);
            checkWatch(_nodeWatches, path, EventType.NodeDeleted);
            if (parentPath != null) {
                checkWatch(_nodeWatches, parentPath, EventType.NodeChildrenChanged);
            }
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public boolean exists(String path, boolean watch) throws KeeperException, InterruptedException {
        _lock.lock();
        try {
            if (watch) {
                installWatch(_nodeWatches, path);
            }
            return _data.containsKey(path);
        } finally {
            _lock.unlock();
        }
    }

    private void installWatch(Set watches, String path) {
        watches.add(path);
    }

    @Override
    public List getChildren(String path, boolean watch) throws KeeperException, InterruptedException {
        if (!exists(path, false)) {
            throw KeeperException.create(Code.NONODE, path);
        }
        if (exists(path, false) && watch) {
            installWatch(_nodeWatches, path);
        }

        checkACL(path, ZooDefs.Perms.READ);
        ArrayList children = new ArrayList();
        String[] directoryStack = path.split("/");
        Set keySet = _data.keySet();

        for (String string : keySet) {
            if (string.startsWith(path)) {
                String[] stack = string.split("/");
                // is one folder level below the one we loockig for and starts
                // with path...
                if (stack.length == directoryStack.length + 1) {
                    children.add(stack[stack.length - 1]);
                }
            }

        }
        return children;
    }

    @Override
    public States getZookeeperState() {
        _lock.lock();
        try {
            if (_eventThread == null) {
                return States.CLOSED;
            }
            return States.CONNECTED;
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public byte[] readData(String path, Stat stat, boolean watch) throws KeeperException, InterruptedException {
        if (watch) {
            installWatch(_dataWatches, path);
        }
        _lock.lock();
        try {
            DataAndVersion dataAndVersion = _data.get(path);
            if (dataAndVersion == null) {
                throw new ZkNoNodeException(new KeeperException.NoNodeException());
            }
            checkACL(path, ZooDefs.Perms.READ);
            byte[] bs = dataAndVersion.getData();
            if (stat != null)
                stat.setVersion(dataAndVersion.getVersion());
            return bs;
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public void writeData(String path, byte[] data, int expectedVersion) throws KeeperException, InterruptedException {
        writeDataReturnStat(path, data, expectedVersion);
    }

    @Override
    public Stat writeDataReturnStat(String path, byte[] data, int expectedVersion) throws KeeperException, InterruptedException {
        int newVersion = -1;
        _lock.lock();
        try {
            checkWatch(_dataWatches, path, EventType.NodeDataChanged);
            if (!exists(path, false)) {
                throw new KeeperException.NoNodeException();
            }
            checkACL(path, ZooDefs.Perms.WRITE);
            newVersion = _data.get(path).getVersion() + 1;
            _data.put(path, new DataAndVersion(data, newVersion));
            String parentPath = getParentPath(path);
            if (parentPath != null) {
                checkWatch(_nodeWatches, parentPath, EventType.NodeChildrenChanged);
            }
        } finally {
            _lock.unlock();
        }
        Stat stat = new Stat();
        stat.setVersion(newVersion);
        return stat;
    }

    private void checkWatch(Set watches, String path, EventType eventType) {
        if (watches.contains(path)) {
            watches.remove(path);
            _eventThread.send(new WatchedEvent(eventType, KeeperState.SyncConnected, path));
        }
    }

    @Override
    public long getCreateTime(String path) {
        Long time = _creationTime.get(path);
        if (time == null) {
            return -1;
        }
        return time;
    }

    @Override
    public String getServers() {
        return "mem";
    }

    public List multi(Iterable ops) throws KeeperException, InterruptedException {
        List opResults = new ArrayList();
        for (Op op : ops) {
            if (Op.Check.class.isAssignableFrom(op.getClass())) {
                CheckVersionRequest check = (CheckVersionRequest) op.toRequestRecord();
                exists(check.getPath(), false);
                opResults.add(new OpResult.CheckResult());
            } else if (Op.Create.class.isAssignableFrom(op.getClass())) {
                CreateRequest create = (CreateRequest) op.toRequestRecord();
                String path = create(create.getPath(), create.getData(), CreateMode.fromFlag(create.getFlags()));
                opResults.add(new OpResult.CreateResult(path));
            } else if (Op.Delete.class.isAssignableFrom(op.getClass())) {
                DeleteRequest delete = (DeleteRequest) op.toRequestRecord();
                delete(delete.getPath());
                opResults.add(new OpResult.DeleteResult());
            } else if (Op.SetData.class.isAssignableFrom(op.getClass())) {
                SetDataRequest setData = (SetDataRequest) op.toRequestRecord();
                writeData(setData.getPath(), setData.getData(), setData.getVersion());
                opResults.add(new OpResult.SetDataResult(null));
            }
        }
        return opResults;
    }

    @Override
    public void addAuthInfo(String scheme, byte[] auth) {
        _ids.add(new Id(scheme, new String(auth)));
    }

    @Override
    public void setAcl(String path, List acl, int version) throws KeeperException, InterruptedException {
        if (!exists(path, false)) {
            throw new KeeperException.NoNodeException();
        }

        DataAndVersion dataAndVersion = _data.get(path);
        if(version != dataAndVersion._version) {
            throw new KeeperException.BadVersionException();
        }

        checkACL(path, ZooDefs.Perms.ADMIN);

        _lock.lock();
        try {
            _data.put(path, new DataAndVersion(dataAndVersion.getData(), dataAndVersion.getVersion() + 1, acl));
        } finally {
            _lock.unlock();
        }
    }

    @Override
    public Map.Entry, Stat> getAcl(String path) throws KeeperException, InterruptedException {
        if (!exists(path, false)) {
            throw new KeeperException.NoNodeException();
        }

        DataAndVersion dataAndVersion = _data.get(path);

        Stat stat = new Stat();
        stat.setVersion(dataAndVersion.getVersion());
        stat.setCtime(_creationTime.get(path));

        return new AbstractMap.SimpleEntry, Stat>(dataAndVersion.getAcl(), stat);
    }

    /***
     * 
     * @param path
     *            - path of znode we are accessing
     * @param perm
     *            - Privileges required for the action
     * @throws KeeperException.NoAuthException
     */
    private void checkACL(String path, int perm) throws KeeperException.NoAuthException {
        DataAndVersion node = _data.get(path);
        if (node == null) {
            return;
        }
        List acl = node.getAcl();
        if (acl == null || acl.size() == 0) {
            return;
        }
        for (Id authId : _ids) {
            if (authId.getScheme().equals("super")) {
                return;
            }
        }
        for (ACL a : acl) {
            Id id = a.getId();
            if ((a.getPerms() & perm) != 0) {
                if (id.getScheme().equals("world") && id.getId().equals("anyone")) {
                    return;
                }
                for (Id authId : _ids) {
                    if (authId.getScheme().equals(id.getScheme()) && authId.getId().equals(id.getId())) {
                        return;
                    }
                }
            }
        }
        throw new KeeperException.NoAuthException();
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy