Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you 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.apache.twill.discovery; import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; import com.google.common.collect.HashMultimap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; import com.google.common.collect.Multimap; import com.google.common.hash.Hashing; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.twill.common.Cancellable; import org.apache.twill.common.Threads; import org.apache.twill.internal.Constants; import org.apache.twill.zookeeper.NodeChildren; import org.apache.twill.zookeeper.NodeData; import org.apache.twill.zookeeper.OperationFuture; import org.apache.twill.zookeeper.ZKClient; import org.apache.twill.zookeeper.ZKClients; import org.apache.twill.zookeeper.ZKOperations; import org.apache.zookeeper.CreateMode; 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.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; /** * Zookeeper implementation of {@link DiscoveryService} and {@link DiscoveryServiceClient}. * <p> * Discoverable services are registered within Zookeeper under the namespace 'discoverable' by default. * If you would like to change the namespace under which the services are registered then you can pass * in the namespace during construction of {@link ZKDiscoveryService}. * </p> * * <p> * Following is a simple example of how {@link ZKDiscoveryService} can be used for registering services * and also for discovering the registered services. * </p> * * <blockquote> * <pre> * {@code * * DiscoveryService service = new ZKDiscoveryService(zkClient); * service.register(new Discoverable() { * @Override * public String getName() { * return 'service-name'; * } * * @Override * public InetSocketAddress getSocketAddress() { * return new InetSocketAddress(hostname, port); * } * }); * ... * ... * ServiceDiscovered services = service.discovery("service-name"); * ... * } * </pre> * </blockquote> */ public class ZKDiscoveryService implements DiscoveryService, DiscoveryServiceClient { private static final Logger LOG = LoggerFactory.getLogger(ZKDiscoveryService.class); private static final long RETRY_MILLIS = 1000; // In memory map for recreating ephemeral nodes after session expires. // It map from discoverable to the corresponding Cancellable private final Multimap<Discoverable, DiscoveryCancellable> discoverables; private final Lock lock; private final LoadingCache<String, ServiceDiscovered> services; private final ZKClient zkClient; private final ScheduledExecutorService retryExecutor; /** * Constructs ZKDiscoveryService using the provided zookeeper client for storing service registry. * @param zkClient The {@link ZKClient} for interacting with zookeeper. */ public ZKDiscoveryService(ZKClient zkClient) { this(zkClient, Constants.DISCOVERY_PATH_PREFIX); } /** * Constructs ZKDiscoveryService using the provided zookeeper client for storing service registry under namespace. * @param zkClient of zookeeper quorum * @param namespace under which the service registered would be stored in zookeeper. * If namespace is {@code null}, no namespace will be used. */ public ZKDiscoveryService(ZKClient zkClient, String namespace) { this.discoverables = HashMultimap.create(); this.lock = new ReentrantLock(); this.retryExecutor = Executors .newSingleThreadScheduledExecutor(Threads.createDaemonThreadFactory("zk-discovery-retry")); this.zkClient = namespace == null ? zkClient : ZKClients.namespace(zkClient, namespace); this.services = CacheBuilder.newBuilder().build(createServiceLoader()); this.zkClient.addConnectionWatcher(createConnectionWatcher()); } /** * Registers a {@link Discoverable} in zookeeper. * <p> * Registering a {@link Discoverable} will create a node <base>/<service-name> * in zookeeper as a ephemeral node. If the node already exists (timeout associated with emphemeral node creation), * then a runtime exception is thrown to make sure that a service with an intent to register is not started without * registering. * When a runtime exception is thrown, expectation is that the process being started will fail and would be started * again by the monitoring service. * </p> * @param discoverable Information of the service provider that could be discovered. * @return An instance of {@link Cancellable} */ @Override public Cancellable register(final Discoverable discoverable) { final Discoverable wrapper = new DiscoverableWrapper(discoverable); final SettableFuture<String> future = SettableFuture.create(); final DiscoveryCancellable cancellable = new DiscoveryCancellable(wrapper); // Create the zk ephemeral node. Futures.addCallback(doRegister(wrapper), new FutureCallback<String>() { @Override public void onSuccess(String result) { // Set the sequence node path to cancellable for future cancellation. cancellable.setPath(result); lock.lock(); try { discoverables.put(wrapper, cancellable); } finally { lock.unlock(); } LOG.debug("Service registered: {} {}", wrapper, result); future.set(result); } @Override public void onFailure(Throwable t) { if (t instanceof KeeperException.NodeExistsException) { handleRegisterFailure(discoverable, future, this, t); } else { LOG.warn("Failed to register: {}", wrapper, t); future.setException(t); } } }, Threads.SAME_THREAD_EXECUTOR); Futures.getUnchecked(future); return cancellable; } @Override public ServiceDiscovered discover(String service) { return services.getUnchecked(service); } /** * Handle registration failure. * * @param discoverable The discoverable to register. * @param completion A settable future to set when registration is completed / failed. * @param creationCallback A future callback for path creation. * @param failureCause The original cause of failure. */ private void handleRegisterFailure(final Discoverable discoverable, final SettableFuture<String> completion, final FutureCallback<String> creationCallback, final Throwable failureCause) { final String path = getNodePath(discoverable); Futures.addCallback(zkClient.exists(path), new FutureCallback<Stat>() { @Override public void onSuccess(Stat result) { if (result == null) { // If the node is gone, simply retry. LOG.info("Node {} is gone. Retry registration for {}.", path, discoverable); retryRegister(discoverable, creationCallback); return; } long ephemeralOwner = result.getEphemeralOwner(); if (ephemeralOwner == 0) { // it is not an ephemeral node, something wrong. LOG.error( "Node {} already exists and is not an ephemeral node. Discoverable registration failed: {}.", path, discoverable); completion.setException(failureCause); return; } Long sessionId = zkClient.getSessionId(); if (sessionId == null || ephemeralOwner != sessionId) { // This zkClient is not valid or doesn't own the ephemeral node, simply keep retrying. LOG.info("Owner of {} is different. Retry registration for {}.", path, discoverable); retryRegister(discoverable, creationCallback); } else { // This client owned the node, treat the registration as completed. // This could happen if same client tries to register twice (due to mistake or failure race condition). completion.set(path); } } @Override public void onFailure(Throwable t) { // If exists call failed, simply retry creation. LOG.warn("Error when getting stats on {}. Retry registration for {}.", path, discoverable); retryRegister(discoverable, creationCallback); } }, Threads.SAME_THREAD_EXECUTOR); } private OperationFuture<String> doRegister(Discoverable discoverable) { byte[] discoverableBytes = DiscoverableAdapter.encode(discoverable); return zkClient.create(getNodePath(discoverable), discoverableBytes, CreateMode.EPHEMERAL, true); } private void retryRegister(final Discoverable discoverable, final FutureCallback<String> creationCallback) { retryExecutor.schedule(new Runnable() { @Override public void run() { Futures.addCallback(doRegister(discoverable), creationCallback, Threads.SAME_THREAD_EXECUTOR); } }, RETRY_MILLIS, TimeUnit.MILLISECONDS); } /** * Generate unique node path for a given {@link Discoverable}. * @param discoverable An instance of {@link Discoverable}. * @return A node name based on the discoverable. */ private String getNodePath(Discoverable discoverable) { InetSocketAddress socketAddress = discoverable.getSocketAddress(); String node = Hashing.md5().newHasher().putBytes(socketAddress.getAddress().getAddress()) .putInt(socketAddress.getPort()).hash().toString(); return String.format("/%s/%s", discoverable.getName(), node); } private Watcher createConnectionWatcher() { return new Watcher() { // Watcher is invoked from single event thread, hence safe to use normal mutable variable. private boolean expired; @Override public void process(WatchedEvent event) { if (event.getState() == Event.KeeperState.Expired) { LOG.warn("ZK Session expired: {}", zkClient.getConnectString()); expired = true; } else if (event.getState() == Event.KeeperState.SyncConnected && expired) { LOG.info("Reconnected after expiration: {}", zkClient.getConnectString()); expired = false; // Re-register all services lock.lock(); try { for (final Map.Entry<Discoverable, DiscoveryCancellable> entry : discoverables.entries()) { LOG.info("Re-registering service: {}", entry.getKey()); // Must be non-blocking in here. Futures.addCallback(doRegister(entry.getKey()), new FutureCallback<String>() { @Override public void onSuccess(String result) { // Updates the cancellable to the newly created sequential node. entry.getValue().setPath(result); LOG.debug("Service re-registered: {} {}", entry.getKey(), result); } @Override public void onFailure(Throwable t) { // When failed to create the node, there would be no retry and simply make the cancellable do nothing. entry.getValue().setPath(null); LOG.error("Failed to re-register service: {}", entry.getKey(), t); } }, Threads.SAME_THREAD_EXECUTOR); } } finally { lock.unlock(); } } } }; } /** * Creates a CacheLoader for creating live Iterable for watching instances changes for a given service. */ private CacheLoader<String, ServiceDiscovered> createServiceLoader() { return new CacheLoader<String, ServiceDiscovered>() { @Override public ServiceDiscovered load(String service) throws Exception { final DefaultServiceDiscovered serviceDiscovered = new DefaultServiceDiscovered(service); final String serviceBase = "/" + service; // Watch for children changes in /service ZKOperations.watchChildren(zkClient, serviceBase, new ZKOperations.ChildrenCallback() { @Override public void updated(NodeChildren nodeChildren) { // Fetch data of all children nodes in parallel. List<String> children = nodeChildren.getChildren(); List<OperationFuture<NodeData>> dataFutures = Lists .newArrayListWithCapacity(children.size()); for (String child : children) { dataFutures.add(zkClient.getData(serviceBase + "/" + child)); } // Update the service map when all fetching are done. final ListenableFuture<List<NodeData>> fetchFuture = Futures.successfulAsList(dataFutures); fetchFuture.addListener(new Runnable() { @Override public void run() { ImmutableSet.Builder<Discoverable> builder = ImmutableSet.builder(); for (NodeData nodeData : Futures.getUnchecked(fetchFuture)) { // For successful fetch, decode the content. if (nodeData != null) { Discoverable discoverable = DiscoverableAdapter.decode(nodeData.getData()); if (discoverable != null) { builder.add(discoverable); } } } serviceDiscovered.setDiscoverables(builder.build()); } }, Threads.SAME_THREAD_EXECUTOR); } }); return serviceDiscovered; } }; } /** * Inner class for cancelling (un-register) discovery service. */ private final class DiscoveryCancellable implements Cancellable { private final Discoverable discoverable; private final AtomicBoolean cancelled; private volatile String path; DiscoveryCancellable(Discoverable discoverable) { this.discoverable = discoverable; this.cancelled = new AtomicBoolean(); } /** * Set the zk node path representing the ephemeral sequence node of this registered discoverable. * Called from ZK event thread when creating of the node completed, either from normal registration or * re-registration due to session expiration. * * @param path The path to ephemeral sequence node. */ void setPath(String path) { this.path = path; if (cancelled.get() && path != null) { // Simply delete the path if it's already cancelled // It's for the case when session expire happened and re-registration completed after this has been cancelled. // Not bother with the result as if there is error, nothing much we could do. zkClient.delete(path); } } @Override public void cancel() { if (!cancelled.compareAndSet(false, true)) { return; } // Take a snapshot of the volatile path. String path = this.path; // If it is null, meaning cancel() is called before the ephemeral node is created, hence // setPath() will be called in future (through zk callback when creation is completed) // so that deletion will be done in setPath(). if (path == null) { return; } // Remove this Cancellable from the map so that upon session expiration won't try to register. lock.lock(); try { discoverables.remove(discoverable, this); } finally { lock.unlock(); } // Delete the path. It's ok if the path not exists // (e.g. what session expired and before node has been re-created) Futures.getUnchecked( ZKOperations.ignoreError(zkClient.delete(path), KeeperException.NoNodeException.class, path)); LOG.debug("Service unregistered: {} {}", discoverable, path); } } }