com.continuuity.weave.internal.kafka.client.KafkaBrokerCache.java Source code

Java tutorial

Introduction

Here is the source code for com.continuuity.weave.internal.kafka.client.KafkaBrokerCache.java

Source

/*
 * Copyright 2012-2013 Continuuity,Inc. All Rights Reserved.
 *
 * 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.continuuity.weave.internal.kafka.client;

import com.continuuity.weave.common.Threads;
import com.continuuity.weave.zookeeper.NodeChildren;
import com.continuuity.weave.zookeeper.NodeData;
import com.continuuity.weave.zookeeper.ZKClient;
import com.google.common.base.Charsets;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.AbstractIdleService;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.SortedMap;

/**
 * A Service to cache kafka broker information by subscribing to ZooKeeper.
 */
final class KafkaBrokerCache extends AbstractIdleService {

    private static final Logger LOG = LoggerFactory.getLogger(KafkaBrokerCache.class);

    private static final String BROKERS_PATH = "/brokers";

    private final ZKClient zkClient;
    private final Map<String, InetSocketAddress> brokers;
    // topicBrokers is from topic->partition size->brokerId
    private final Map<String, SortedMap<Integer, Set<String>>> topicBrokers;
    private final Runnable invokeGetBrokers = new Runnable() {
        @Override
        public void run() {
            getBrokers();
        }
    };
    private final Runnable invokeGetTopics = new Runnable() {
        @Override
        public void run() {
            getTopics();
        }
    };

    KafkaBrokerCache(ZKClient zkClient) {
        this.zkClient = zkClient;
        this.brokers = Maps.newConcurrentMap();
        this.topicBrokers = Maps.newConcurrentMap();
    }

    @Override
    protected void startUp() throws Exception {
        getBrokers();
        getTopics();
    }

    @Override
    protected void shutDown() throws Exception {
        // No-op
    }

    public int getPartitionSize(String topic) {
        SortedMap<Integer, Set<String>> partitionBrokers = topicBrokers.get(topic);
        if (partitionBrokers == null || partitionBrokers.isEmpty()) {
            return 1;
        }
        return partitionBrokers.lastKey();
    }

    public TopicBroker getBrokerAddress(String topic, int partition) {
        SortedMap<Integer, Set<String>> partitionBrokers = topicBrokers.get(topic);
        if (partitionBrokers == null || partitionBrokers.isEmpty()) {
            return pickRandomBroker(topic);
        }

        // If the requested partition is greater than supported partition size, randomly pick one
        if (partition >= partitionBrokers.lastKey()) {
            return pickRandomBroker(topic);
        }

        // Randomly pick a partition size and randomly pick a broker from it
        Random random = new Random();
        partitionBrokers = partitionBrokers.tailMap(partition + 1);
        List<Integer> sizes = Lists.newArrayList(partitionBrokers.keySet());
        Integer partitionSize = pickRandomItem(sizes, random);
        List<String> ids = Lists.newArrayList(partitionBrokers.get(partitionSize));
        InetSocketAddress address = brokers.get(ids.get(new Random().nextInt(ids.size())));
        return address == null ? pickRandomBroker(topic) : new TopicBroker(topic, address, partitionSize);
    }

    private TopicBroker pickRandomBroker(String topic) {
        Map.Entry<String, InetSocketAddress> entry = Iterables.getFirst(brokers.entrySet(), null);
        if (entry == null) {
            return null;
        }
        InetSocketAddress address = entry.getValue();
        return new TopicBroker(topic, address, 0);
    }

    private <T> T pickRandomItem(List<T> list, Random random) {
        return list.get(random.nextInt(list.size()));
    }

    private void getBrokers() {
        final String idsPath = BROKERS_PATH + "/ids";

        Futures.addCallback(zkClient.getChildren(idsPath, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                getBrokers();
            }
        }), new ExistsOnFailureFutureCallback<NodeChildren>(idsPath, invokeGetBrokers) {
            @Override
            public void onSuccess(NodeChildren result) {
                Set<String> children = ImmutableSet.copyOf(result.getChildren());
                for (String child : children) {
                    getBrokenData(idsPath + "/" + child, child);
                }
                // Remove all removed brokers
                removeDiff(children, brokers);
            }
        });
    }

    private void getTopics() {
        final String topicsPath = BROKERS_PATH + "/topics";
        Futures.addCallback(zkClient.getChildren(topicsPath, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                getTopics();
            }
        }), new ExistsOnFailureFutureCallback<NodeChildren>(topicsPath, invokeGetTopics) {
            @Override
            public void onSuccess(NodeChildren result) {
                Set<String> children = ImmutableSet.copyOf(result.getChildren());

                // Process new children
                for (String topic : ImmutableSet.copyOf(Sets.difference(children, topicBrokers.keySet()))) {
                    getTopic(topicsPath + "/" + topic, topic);
                }

                // Remove old children
                removeDiff(children, topicBrokers);
            }
        });
    }

    private void getBrokenData(String path, final String brokerId) {
        Futures.addCallback(zkClient.getData(path), new FutureCallback<NodeData>() {
            @Override
            public void onSuccess(NodeData result) {
                String data = new String(result.getData(), Charsets.UTF_8);
                String hostPort = data.substring(data.indexOf(':') + 1);
                int idx = hostPort.indexOf(':');
                brokers.put(brokerId, new InetSocketAddress(hostPort.substring(0, idx),
                        Integer.parseInt(hostPort.substring(idx + 1))));
            }

            @Override
            public void onFailure(Throwable t) {
                // No-op, the watch on the parent node will handle it.
            }
        });
    }

    private void getTopic(final String path, final String topic) {
        Futures.addCallback(zkClient.getChildren(path, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                // Other event type changes are either could be ignored or handled by parent watcher
                if (event.getType() == Event.EventType.NodeChildrenChanged) {
                    getTopic(path, topic);
                }
            }
        }), new FutureCallback<NodeChildren>() {
            @Override
            public void onSuccess(NodeChildren result) {
                List<String> children = result.getChildren();
                final List<ListenableFuture<BrokerPartition>> futures = Lists
                        .newArrayListWithCapacity(children.size());

                // Fetch data from each broken node
                for (final String brokerId : children) {
                    Futures.transform(zkClient.getData(path + "/" + brokerId),
                            new Function<NodeData, BrokerPartition>() {
                                @Override
                                public BrokerPartition apply(NodeData input) {
                                    return new BrokerPartition(brokerId,
                                            Integer.parseInt(new String(input.getData(), Charsets.UTF_8)));
                                }
                            });
                }

                // When all fetching is done, build the partition size->broker map for this topic
                Futures.successfulAsList(futures).addListener(new Runnable() {
                    @Override
                    public void run() {
                        Map<Integer, Set<String>> partitionBrokers = Maps.newHashMap();
                        for (ListenableFuture<BrokerPartition> future : futures) {
                            try {
                                BrokerPartition info = future.get();
                                Set<String> brokerSet = partitionBrokers.get(info.getPartitionSize());
                                if (brokerSet == null) {
                                    brokerSet = Sets.newHashSet();
                                    partitionBrokers.put(info.getPartitionSize(), brokerSet);
                                }
                                brokerSet.add(info.getBrokerId());
                            } catch (Exception e) {
                                // Exception is ignored, as it will be handled by parent watcher
                            }
                        }
                        topicBrokers.put(topic, ImmutableSortedMap.copyOf(partitionBrokers));
                    }
                }, Threads.SAME_THREAD_EXECUTOR);
            }

            @Override
            public void onFailure(Throwable t) {
                // No-op. Failure would be handled by parent watcher already (e.g. node not exists -> children change in parent)
            }
        });
    }

    private <K, V> void removeDiff(Set<K> keys, Map<K, V> map) {
        for (K key : ImmutableSet.copyOf(Sets.difference(map.keySet(), keys))) {
            map.remove(key);
        }
    }

    private abstract class ExistsOnFailureFutureCallback<V> implements FutureCallback<V> {

        private final String path;
        private final Runnable action;

        protected ExistsOnFailureFutureCallback(String path, Runnable action) {
            this.path = path;
            this.action = action;
        }

        @Override
        public final void onFailure(Throwable t) {
            if (!isNotExists(t)) {
                LOG.error("Fail to watch for kafka brokers: " + path, t);
                return;
            }

            waitExists(path);
        }

        private boolean isNotExists(Throwable t) {
            return ((t instanceof KeeperException) && ((KeeperException) t).code() == KeeperException.Code.NONODE);
        }

        private void waitExists(String path) {
            LOG.info("Path " + path + " not exists. Watch for creation.");

            // If the node doesn't exists, use the "exists" call to watch for node creation.
            Futures.addCallback(zkClient.exists(path, new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    if (event.getType() == Event.EventType.NodeCreated
                            || event.getType() == Event.EventType.NodeDeleted) {
                        action.run();
                    }
                }
            }), new FutureCallback<Stat>() {
                @Override
                public void onSuccess(Stat result) {
                    // If path exists, get children again, otherwise wait for watch to get triggered
                    if (result != null) {
                        action.run();
                    }
                }

                @Override
                public void onFailure(Throwable t) {
                    action.run();
                }
            });
        }
    }

    private static final class BrokerPartition {
        private final String brokerId;
        private final int partitionSize;

        private BrokerPartition(String brokerId, int partitionSize) {
            this.brokerId = brokerId;
            this.partitionSize = partitionSize;
        }

        public String getBrokerId() {
            return brokerId;
        }

        public int getPartitionSize() {
            return partitionSize;
        }
    }
}