org.kiji.schema.zookeeper.UsersTracker.java Source code

Java tutorial

Introduction

Here is the source code for org.kiji.schema.zookeeper.UsersTracker.java

Source

/**
 * (c) Copyright 2014 WibiData, Inc.
 *
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership.
 *
 * 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.kiji.schema.zookeeper;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.net.URLDecoder;
import java.util.Collections;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import com.google.common.base.Charsets;
import com.google.common.base.Objects;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCache.StartMode;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Tracks users of a table or instance. Allows retrieving the current users, and allows callbacks
 * to be registered which fire on changes to the set of registered users.
 */
public class UsersTracker implements Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(UsersTracker.class);

    private final File mUsersDir;

    private final PathChildrenCache mCache;

    private final Set<UsersUpdateHandler> mHandlers = Collections
            .newSetFromMap(Maps.<UsersUpdateHandler, Boolean>newConcurrentMap());

    /**
     * Initializes a new users tracker. Package private, users should use the
     * {@code newTableUsersTracker()} or {@code newInstanceUsersTracker()} factory methods in
     * {@link org.kiji.schema.zookeeper.ZooKeeperUtils} to construct a new {@link UsersTracker}.
     *
     * @param zkClient connection to ZooKeeper.
     * @param usersDir path to the directory containing user nodes.
     */
    UsersTracker(CuratorFramework zkClient, File usersDir) {
        mUsersDir = usersDir;
        mCache = new PathChildrenCache(zkClient, usersDir.getPath(), false);
    }

    /**
     * Register a new update handler on this users tracker.
     *
     * @param handler to register.
     * @return this.
     */
    public UsersTracker registerUpdateHandler(UsersUpdateHandler handler) {
        mHandlers.add(handler);
        return this;
    }

    /**
     * Unregister an existing update handler from this users tracker.
     *
     * @param handler to unregister.
     * @return this.
     */
    public UsersTracker unregisterUpdateHandler(UsersUpdateHandler handler) {
        mHandlers.remove(handler);
        return this;
    }

    /**
     * Update the registered update handlers with the current users.
     *
     * @throws IOException on unrecoverable error.
     */
    public void updateRegisteredHandlers() throws IOException {
        final Multimap<String, String> users = getUsers();
        for (UsersUpdateHandler handler : mHandlers) {
            handler.update(users);
        }
    }

    /**
     * Starts the tracker.  Will block until the users list is initialized. When the block is
     * initialized all handlers will be invoked in another thread.
     *
     * @throws IOException on unrecoverable ZooKeeper error.
     * @return this.
     */
    public UsersTracker start() throws IOException {
        LOG.debug("Starting table user tracker for path {}.", mUsersDir);
        final CountDownLatch initializationLatch = new CountDownLatch(1);

        mCache.getListenable().addListener(new PathChildrenCacheListener() {
            @Override
            public void childEvent(CuratorFramework client, PathChildrenCacheEvent event) throws Exception {
                LOG.debug("Users tracker event received for path {}: {}.", mUsersDir, event);
                switch (event.getType()) {
                case CHILD_ADDED:
                case CHILD_UPDATED:
                case CHILD_REMOVED: {
                    updateRegisteredHandlers();
                    break;
                }
                case INITIALIZED: {
                    initializationLatch.countDown();
                    updateRegisteredHandlers();
                    break;
                }
                default:
                    break; // Connection state changes are already logged at the client level.
                }
            }
        });

        try {
            mCache.start(StartMode.POST_INITIALIZED_EVENT);
            if (!initializationLatch.await(5, TimeUnit.SECONDS)) {
                throw new IOException("Unable to start ZooKeeper path cache before timeout.");
            }
        } catch (Exception e) {
            ZooKeeperUtils.wrapAndRethrow(e);
        }
        return this;
    }

    /**
     * Retrieve a multimap of the current users to their respective value (either table layout version
     * if tracking table users, or system version if tracking instance users).
     *
     * @return a multimap of the current table users to their respective value.
     * @throws IOException on failure.
     */
    public Multimap<String, String> getUsers() throws IOException {
        ImmutableMultimap.Builder<String, String> users = ImmutableSetMultimap.builder();
        for (ChildData child : mCache.getCurrentData()) {
            String nodeName = Iterables.getLast(Splitter.on('/').split(child.getPath()));
            String[] split = nodeName.split(ZooKeeperUtils.ZK_NODE_NAME_SEPARATOR);
            if (split.length != 2) {
                LOG.error("Ignoring invalid ZooKeeper table user node: {}.", nodeName);
                continue;
            }
            final String userID = URLDecoder.decode(split[0], Charsets.UTF_8.displayName());
            final String layoutID = URLDecoder.decode(split[1], Charsets.UTF_8.displayName());

            users.put(userID, layoutID);
        }
        return users.build();
    }

    /** {@inheritDoc} */
    @Override
    public void close() throws IOException {
        LOG.debug("Closing UsersTracker {}.", this);
        mCache.close();
        mHandlers.clear();
    }

    /** {@inheritDoc} */
    @Override
    public String toString() {
        return Objects.toStringHelper(UsersTracker.class).add("directory", mUsersDir).toString();
    }
}