Java tutorial
/* * Copyright 2014 Cask Data, Inc. * * 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 co.cask.cdap.common.zookeeper.coordination; import co.cask.cdap.api.common.Bytes; import co.cask.cdap.common.zookeeper.ZKExtOperations; import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSortedSet; import com.google.common.collect.Maps; import com.google.common.collect.Multimap; import com.google.common.collect.Sets; import com.google.common.collect.TreeMultimap; import com.google.common.util.concurrent.AbstractService; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import org.apache.twill.common.Cancellable; import org.apache.twill.common.Threads; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; import org.apache.twill.discovery.ServiceDiscovered; import org.apache.twill.zookeeper.NodeChildren; import org.apache.twill.zookeeper.NodeData; import org.apache.twill.zookeeper.ZKClient; 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.util.Map; import java.util.Set; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; /** * Allocate resources to registered handler. It is expected to have single instance of this class * running per ZK namespace (determined by the ZKClient passed to constructor). User can use * {@link org.apache.twill.internal.zookeeper.LeaderElection} class to help archiving requirement. * * ZK structure. * * <pre> * {@code * /requirements * /[resource_name] - Contains ResourceRequirement encoded in json * /assignments * /[resource_name] - Contains ResourceAssignment encoded in json * } * </pre> * */ public final class ResourceCoordinator extends AbstractService { private static final Logger LOG = LoggerFactory.getLogger(ResourceCoordinator.class); private final ZKClient zkClient; private final DiscoveryServiceClient discoveryService; private final AssignmentStrategy assignmentStrategy; private final Map<String, ResourceRequirement> requirements; private final Map<String, ResourceAssignment> assignments; private final Map<String, CancellableServiceDiscovered> serviceDiscovered; private final DiscoverableChangeListener discoverableListener; // A single thread executor to process resource requests and perform allocation. private ExecutorService executor; public ResourceCoordinator(ZKClient zkClient, DiscoveryServiceClient discoveryService, AssignmentStrategy assignmentStrategy) { this.zkClient = zkClient; this.discoveryService = discoveryService; this.assignmentStrategy = assignmentStrategy; this.requirements = Maps.newHashMap(); this.assignments = Maps.newHashMap(); this.serviceDiscovered = Maps.newHashMap(); this.discoverableListener = new DiscoverableChangeListener(); } @Override protected void doStart() { executor = Executors.newSingleThreadExecutor(Threads.createDaemonThreadFactory("resource-coordinator")); beginWatch(wrapWatcher(new ResourceWatcher())); notifyStarted(); } @Override protected void doStop() { try { executor.execute(createShutdownTask(null)); } finally { executor.shutdown(); } } /** * Signals failure in this service and terminates itself. * * @param cause Reason for failure. */ private void doNotifyFailed(Throwable cause) { try { executor.execute(createShutdownTask(cause)); } finally { executor.shutdown(); } } private Runnable createShutdownTask(final Throwable failureCause) { return new Runnable() { @Override public void run() { for (Cancellable cancellable : serviceDiscovered.values()) { try { cancellable.cancel(); } catch (Throwable t) { LOG.warn("Exception when cancelling service discovery listener.", t); } } if (failureCause == null) { notifyStopped(); } else { notifyFailed(failureCause); } } }; } /** * Start watching for changes in resources requirements. */ private void beginWatch(final Watcher watcher) { Futures.addCallback(zkClient.exists(CoordinationConstants.REQUIREMENTS_PATH, watcher), wrapCallback(new FutureCallback<Stat>() { @Override public void onSuccess(Stat result) { if (result != null) { fetchAndProcessAllResources(watcher); } // If the node doesn't exists yet, that's ok, the watcher would handle it once it's created. } @Override public void onFailure(Throwable t) { // Something very wrong to have exists call failed. LOG.error("Failed to call exists on ZK node {}{}", zkClient.getConnectString(), CoordinationConstants.REQUIREMENTS_PATH, t); doNotifyFailed(t); } }), executor); } /** * Fetches all {@link ResourceRequirement} and perform assignment for the one that changed. Also, it will * remove assignments for the resource requirements that are removed. */ private void fetchAndProcessAllResources(final Watcher watcher) { Futures.addCallback(zkClient.getChildren(CoordinationConstants.REQUIREMENTS_PATH, watcher), wrapCallback(new FutureCallback<NodeChildren>() { @Override public void onSuccess(NodeChildren result) { Set<String> children = ImmutableSet.copyOf(result.getChildren()); // Handle new resources for (String child : children) { String path = CoordinationConstants.REQUIREMENTS_PATH + "/" + child; Watcher requirementWatcher = wrapWatcher(new ResourceRequirementWatcher(path)); fetchAndProcessRequirement(path, requirementWatcher); } // Handle removed resources for (String removed : ImmutableSet .copyOf(Sets.difference(requirements.keySet(), children))) { ResourceRequirement requirement = requirements.remove(removed); LOG.info("Requirement deleted {}", requirement); // Delete the assignment node. removeAssignment(removed); } } @Override public void onFailure(Throwable t) { // If the resource path node doesn't exists, resort to watch for exists. if (t instanceof KeeperException.NoNodeException) { beginWatch(watcher); } // Otherwise, it's a unexpected failure. LOG.error("Failed to getChildren on ZK node {}{}", zkClient.getConnectString(), CoordinationConstants.REQUIREMENTS_PATH, t); doNotifyFailed(t); } }), executor); } /** * Gets the data from a resource node, decode it to {@link ResourceRequirement} and performs resource assignment * if the requirement changed. */ private void fetchAndProcessRequirement(final String path, Watcher watcher) { Futures.addCallback(zkClient.getData(path, watcher), wrapCallback(new FutureCallback<NodeData>() { @Override public void onSuccess(NodeData result) { byte[] nodeData = result.getData(); if (nodeData == null) { LOG.warn("Ignore empty data in ZK node {}{}", zkClient.getConnectString(), path); return; } try { ResourceRequirement requirement = CoordinationConstants.RESOURCE_REQUIREMENT_CODEC .decode(nodeData); LOG.info("Get requirement {}", requirement); // See if the requirement changed. ResourceRequirement oldRequirement = requirements.get(requirement.getName()); if (requirement.equals(oldRequirement)) { LOG.info("Requirement for {} is not changed. No assignment is needed. {} = {}", requirement.getName(), oldRequirement, requirement); return; } // Requirement change, perform assignment, optional subscribe to service discovery if not yet did. requirements.put(requirement.getName(), requirement); CancellableServiceDiscovered discovered = serviceDiscovered.get(requirement.getName()); if (discovered == null) { discovered = new CancellableServiceDiscovered( discoveryService.discover(requirement.getName()), discoverableListener, executor); serviceDiscovered.put(requirement.getName(), discovered); // If it is the first time it subscribes, no need to trigger assignment logic here as the first call // to listener.onChange would do so. } else { performAssignment(requirement, discovered.serviceDiscovered); } } catch (Exception e) { LOG.warn("Failed to process requirement ZK node {}{}: {}", zkClient.getConnectString(), path, Bytes.toStringBinary(nodeData), e); } } @Override public void onFailure(Throwable t) { // Just log LOG.error("Failed to getData on ZK node {}{}", zkClient.getConnectString(), path, t); } }), executor); } /** * Performs resource assignment based on the resource requirement. * This method should only be called from the single thread executor owned by this class. * * @param requirement The resource requirement that needs to be fulfilled. * @param serviceDiscovered The set of handlers available. */ private void performAssignment(ResourceRequirement requirement, ServiceDiscovered serviceDiscovered) { ResourceAssignment oldAssignment = assignments.get(requirement.getName()); if (oldAssignment == null) { // Try to fetch it from ZK. // This is the penalty to pay for the first time this coordinator instance sees a given requirement. fetchAndPerformAssignment(requirement, serviceDiscovered); return; } Set<Discoverable> handlers = ImmutableSortedSet.copyOf(DiscoverableComparator.COMPARATOR, serviceDiscovered); LOG.info("Perform assign for requirement {}. Number of handlers is {}", requirement, handlers.size()); // Build a map from partition name to number of replicas Map<String, Integer> partitions = Maps.newHashMap(); for (ResourceRequirement.Partition partition : requirement.getPartitions()) { partitions.put(partition.getName(), partition.getReplicas()); } // Copy all valid partition replicas assignments and drops the invalid one. Multimap<Discoverable, PartitionReplica> assignmentMap = TreeMultimap .create(DiscoverableComparator.COMPARATOR, PartitionReplica.COMPARATOR); for (Map.Entry<Discoverable, PartitionReplica> entry : oldAssignment.getAssignments().entries()) { Integer replicas = partitions.get(entry.getValue().getName()); // Check if the partition replica is still valid and if the handler being assigned still exists. if (replicas != null && entry.getValue().getReplicaId() < replicas && handlers.contains(entry.getKey())) { assignmentMap.put(entry.getKey(), entry.getValue()); } } ResourceAssigner<Discoverable> assigner = DefaultResourceAssigner.create(assignmentMap); // Call the strategy for assignment only if there are some handlers and the requirement is not empty. // Otherwise the assignment will be just an empty assignment if (!handlers.isEmpty() && !partitions.isEmpty()) { assignmentStrategy.assign(requirement, handlers, assigner); } // Save the new assignment saveAssignment(new ResourceAssignment(requirement.getName(), assigner.get())); } /** * Fetch the {@link ResourceAssignment} from ZK and then perform resource assignment logic. This is done with best * effort to let the {@link AssignmentStrategy} has access to existing assignment. If failed to get existing * {@link ResourceAssignment} or if it's simply not exists, assignment will still be triggered as if there is no * existing assignment. * * @param requirement The resource requirement that needs to be fulfilled. * @param serviceDiscovered The set of handlers available. */ private void fetchAndPerformAssignment(final ResourceRequirement requirement, final ServiceDiscovered serviceDiscovered) { final String name = requirement.getName(); String zkPath = CoordinationConstants.ASSIGNMENTS_PATH + "/" + name; Futures.addCallback(zkClient.getData(zkPath), new FutureCallback<NodeData>() { @Override public void onSuccess(NodeData result) { if (assignments.get(name) != null) { // Assignment should has been performed while this one is fetching. So, ignore this. return; } byte[] data = result.getData(); ResourceAssignment resourceAssignment = new ResourceAssignment(name); try { if (data != null) { resourceAssignment = CoordinationConstants.RESOURCE_ASSIGNMENT_CODEC.decode(data); } } catch (Throwable t) { LOG.warn( "Failed to decode resource assignment. Perform assignment as if no assignment existed.", t); } assignments.put(name, resourceAssignment); performAssignment(requirement, serviceDiscovered); } @Override public void onFailure(Throwable t) { if (!(t instanceof KeeperException.NoNodeException)) { // If failure is not because node doesn't exists, log a warning LOG.warn("Failed to fetch current assignment. Perform assignment as if no assignment existed.", t); } assignments.put(name, new ResourceAssignment(name)); performAssignment(requirement, serviceDiscovered); } }, executor); } /** * Save a {@link ResourceAssignment} to local cache as well as ZK ZK. * @param assignment The assignment to be persisted. */ private void saveAssignment(ResourceAssignment assignment) { assignments.put(assignment.getName(), assignment); try { final byte[] data = CoordinationConstants.RESOURCE_ASSIGNMENT_CODEC.encode(assignment); String zkPath = CoordinationConstants.ASSIGNMENTS_PATH + "/" + assignment.getName(); Futures.addCallback(ZKExtOperations.setOrCreate(zkClient, zkPath, data, assignment, CoordinationConstants.MAX_ZK_FAILURE_RETRY), new FutureCallback<ResourceAssignment>() { @Override public void onSuccess(ResourceAssignment result) { // Done. Just log for debugging. LOG.info("Resource assignment updated for {}. {}", result.getName(), Bytes.toString(data)); } @Override public void onFailure(Throwable t) { LOG.error("Failed to save assignment {}", Bytes.toStringBinary(data), t); doNotifyFailed(t); } }, executor); } catch (Exception e) { // Something very wrong LOG.error("Failed to save assignment: {}", assignment.getName(), e); } } /** * Removes the {@link ResourceAssignment} with the given name from local cache as well as ZK. * @param name Name of the resource. */ private void removeAssignment(String name) { assignments.remove(name); String zkPath = CoordinationConstants.ASSIGNMENTS_PATH + "/" + name; // Simply delete the assignment node. No need to care about the result. // Even if failed to delete the node and leaves stale assignment, next time when an assignment action is // triggered, it'll correct it. zkClient.delete(zkPath); } /** * Returns true if this service is up and running, hence should process any events that it received. */ private boolean shouldProcess() { State state = state(); return state == State.STARTING || state == State.RUNNING; } /** * Wraps a given callback so that it only get triggered if {@link #shouldProcess()} returns true. */ private <T> FutureCallback<T> wrapCallback(final FutureCallback<T> callback) { return new FutureCallback<T>() { @Override public void onSuccess(T result) { if (shouldProcess()) { callback.onSuccess(result); } } @Override public void onFailure(Throwable t) { if (shouldProcess()) { callback.onFailure(t); } } }; } /** * Wraps a given Watcher so that it only get triggered if {@link #shouldProcess()} returns true. */ private Watcher wrapWatcher(final Watcher watcher) { return new Watcher() { @Override public void process(WatchedEvent event) { if (!shouldProcess()) { return; } watcher.process(event); } }; } /** * Watcher to handle children nodes changes in the resource requirement. Child node will be added / removed * when requirement is added or removed. */ private final class ResourceWatcher implements Watcher { @Override public void process(WatchedEvent event) { if (!shouldProcess()) { return; } switch (event.getType()) { case NodeCreated: case NodeChildrenChanged: fetchAndProcessAllResources(this); break; case NodeDeleted: beginWatch(this); break; default: // No-op } } } /** * Watcher to get updates in resource requirement node. */ private final class ResourceRequirementWatcher implements Watcher { private final String path; private ResourceRequirementWatcher(String path) { this.path = path; } @Override public void process(WatchedEvent event) { if (!shouldProcess()) { return; } // Only interested in data change event. Other type of events is handled by the watcher on parent node. if (event.getType() == Event.EventType.NodeDataChanged) { fetchAndProcessRequirement(path, this); } } } private final class DiscoverableChangeListener implements ServiceDiscovered.ChangeListener { @Override public void onChange(ServiceDiscovered serviceDiscovered) { ResourceRequirement requirement = requirements.get(serviceDiscovered.getName()); if (requirement != null) { performAssignment(requirement, serviceDiscovered); } } } private static final class CancellableServiceDiscovered implements Cancellable { private final Cancellable cancellable; private final ServiceDiscovered serviceDiscovered; private CancellableServiceDiscovered(ServiceDiscovered serviceDiscovered, ServiceDiscovered.ChangeListener listener, Executor executor) { this.cancellable = serviceDiscovered.watchChanges(listener, executor); this.serviceDiscovered = serviceDiscovered; } @Override public void cancel() { cancellable.cancel(); } } }