Added stubs for QueryResourceManager exit and wait/cleanup thread

Update MemoryCalculator to use DrillNode instead of DrillbitEndpoint
Changes to support localbit resource registration to cluster state blob using DrillbitStatusListener
Support ThrottledResourceManager via ResourceManagerBuilder
Add some E2E tests and RMStateBlobs tests along with some bug fixes
Fix TestRMConfigLoad tests to handle case where ZKQueues are explicitly enabled

Author: sohami

drill-yarn/src/main/java/org/apache/drill/yarn/zk/ZKClusterCoordinator.java

@@ -17,17 +17,6 @@
  */
 package org.apache.drill.yarn.zk;
 
-import static org.apache.drill.shaded.guava.com.google.common.collect.Collections2.transform;
-
-import java.io.IOException;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.Set;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-
-import org.apache.drill.shaded.guava.com.google.common.base.Throwables;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.curator.RetryPolicy;
@@ -55,8 +44,16 @@
 import org.apache.drill.exec.coord.zk.ZkTransientStoreFactory;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint.State;
+import org.apache.drill.shaded.guava.com.google.common.base.Throwables;
 
-import org.apache.drill.shaded.guava.com.google.common.base.Function;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 /**
  * Manages cluster coordination utilizing zookeeper.
@@ -87,8 +84,7 @@ public class ZKClusterCoordinator extends ClusterCoordinator {
 
   private CuratorFramework curator;
   private ServiceDiscovery<DrillbitEndpoint> discovery;
-  private volatile Collection<DrillbitEndpoint> endpoints = Collections
-      .emptyList();
+  private volatile Map<String, DrillbitEndpoint> endpointsMap = Collections.emptyMap();
   private final String serviceName;
   private final CountDownLatch initialConnection = new CountDownLatch(1);
   private final TransientStoreFactory factory;
@@ -214,7 +210,7 @@ public void unregister(RegistrationHandle handle) {
 
   @Override
   public Collection<DrillbitEndpoint> getAvailableEndpoints() {
-    return this.endpoints;
+    return this.endpointsMap.values();
   }
 
   @Override
@@ -233,35 +229,33 @@ public <V> TransientStore<V> getOrCreateTransientStore(
 
   private synchronized void updateEndpoints() {
     try {
-      Collection<DrillbitEndpoint> newDrillbitSet = transform(
-          discovery.queryForInstances(serviceName),
-          new Function<ServiceInstance<DrillbitEndpoint>, DrillbitEndpoint>() {
-            @Override
-            public DrillbitEndpoint apply(
-                ServiceInstance<DrillbitEndpoint> input) {
-              return input.getPayload();
-            }
-          });
+      // All active bits in the Zookeeper
+      final Map<String, DrillbitEndpoint> UUIDtoEndpoints = discovery.queryForInstances(serviceName).stream()
+        .collect(Collectors.toConcurrentMap(ServiceInstance::getId, ServiceInstance::getPayload));
 
       // set of newly dead bits : original bits - new set of active bits.
-      Set<DrillbitEndpoint> unregisteredBits = new HashSet<>(endpoints);
-      unregisteredBits.removeAll(newDrillbitSet);
+      Map<String, DrillbitEndpoint> unregisteredBits = new HashMap<>(endpointsMap);
+      for (Map.Entry<String, DrillbitEndpoint> newEndpoint : UUIDtoEndpoints.entrySet()) {
+        unregisteredBits.remove(newEndpoint.getKey());
+      }
 
       // Set of newly live bits : new set of active bits - original bits.
-      Set<DrillbitEndpoint> registeredBits = new HashSet<>(newDrillbitSet);
-      registeredBits.removeAll(endpoints);
+      Map<String, DrillbitEndpoint> registeredBits = new HashMap<>(UUIDtoEndpoints);
+      for (Map.Entry<String, DrillbitEndpoint> newEndpoint : endpointsMap.entrySet()) {
+        registeredBits.remove(newEndpoint.getKey());
+      }
 
-      endpoints = newDrillbitSet;
+      endpointsMap = UUIDtoEndpoints;
 
       if (logger.isDebugEnabled()) {
         StringBuilder builder = new StringBuilder();
         builder.append("Active drillbit set changed.  Now includes ");
-        builder.append(newDrillbitSet.size());
+        builder.append(UUIDtoEndpoints.size());
         builder.append(" total bits.");
-        if (!newDrillbitSet.isEmpty()) {
+        if (!UUIDtoEndpoints.isEmpty()) {
           builder.append(" New active drillbits: \n");
         }
-        for (DrillbitEndpoint bit : newDrillbitSet) {
+        for (DrillbitEndpoint bit : UUIDtoEndpoints.values()) {
           builder.append('\t');
           builder.append(bit.getAddress());
           builder.append(':');
@@ -277,11 +271,13 @@ public DrillbitEndpoint apply(
 
       // Notify the drillbit listener for newly unregistered bits.
       if (!(unregisteredBits.isEmpty())) {
-        drillbitUnregistered(unregisteredBits);
+        drillbitUnregistered(unregisteredBits.entrySet().stream().collect(
+          Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey)));
       }
       // Notify the drillbit listener for newly registered bits.
       if (!(registeredBits.isEmpty())) {
-        drillbitRegistered(registeredBits);
+        drillbitRegistered(registeredBits.entrySet().stream().collect(
+          Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey)));
       }
 
     } catch (Exception e) {

drill-yarn/src/main/java/org/apache/drill/yarn/zk/ZKRegistry.java

@@ -17,24 +17,24 @@
  */
 package org.apache.drill.yarn.zk;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
-import org.apache.drill.shaded.guava.com.google.common.annotations.VisibleForTesting;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.exec.work.foreman.DrillbitStatusListener;
+import org.apache.drill.shaded.guava.com.google.common.annotations.VisibleForTesting;
 import org.apache.drill.yarn.appMaster.AMWrapperException;
 import org.apache.drill.yarn.appMaster.EventContext;
 import org.apache.drill.yarn.appMaster.Pollable;
 import org.apache.drill.yarn.appMaster.RegistryHandler;
 import org.apache.drill.yarn.appMaster.Task;
 import org.apache.drill.yarn.appMaster.TaskLifecycleListener;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
 /**
  * AM-specific implementation of a Drillbit registry backed by ZooKeeper.
  * Listens to ZK events for registering a Drillbit and deregistering. Alerts the
@@ -297,11 +297,12 @@ public AckEvent(Task task, DrillbitEndpoint endpoint) {
    * Callback from ZK to indicate that one or more drillbits have become
    * registered. We handle registrations in a critical section, then alert the
    * cluster controller outside the critical section.
+   * @param registeredDrillbitsUUID
    */
 
   @Override
-  public void drillbitRegistered(Set<DrillbitEndpoint> registeredDrillbits) {
-    List<AckEvent> updates = registerDrillbits(registeredDrillbits);
+  public void drillbitRegistered(Map<DrillbitEndpoint, String> registeredDrillbitsUUID) {
+    List<AckEvent> updates = registerDrillbits(registeredDrillbitsUUID.keySet());
     for (AckEvent event : updates) {
       if (event.task == null) {
         registryHandler.reserveHost(event.endpoint.getAddress());
@@ -363,12 +364,12 @@ private AckEvent drillbitRegistered(DrillbitEndpoint dbe) {
    * Callback from ZK to indicate that one or more drillbits have become
    * deregistered from ZK. We handle the deregistrations in a critical section,
    * but updates to the cluster controller outside of a critical section.
+   * @param unregisteredDrillbitsUUID
    */
 
   @Override
-  public void drillbitUnregistered(
-      Set<DrillbitEndpoint> unregisteredDrillbits) {
-    List<AckEvent> updates = unregisterDrillbits(unregisteredDrillbits);
+  public void drillbitUnregistered(Map<DrillbitEndpoint, String> unregisteredDrillbitsUUID) {
+    List<AckEvent> updates = unregisterDrillbits(unregisteredDrillbitsUUID.keySet());
     for (AckEvent event : updates) {
       registryHandler.completionAck(event.task, ENDPOINT_PROPERTY);
     }

drill-yarn/src/test/java/org/apache/drill/yarn/zk/TestZkRegistry.java

@@ -17,15 +17,6 @@
  */
 package org.apache.drill.yarn.zk;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.retry.RetryNTimes;
@@ -41,6 +32,15 @@
 import org.apache.drill.yarn.zk.ZKRegistry.DrillbitTracker;
 import org.junit.Test;
 
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
 /**
  * Tests for the AM version of the cluster coordinator. The AM version has no
  * dependencies on the DoY config system or other systems, making it easy to
@@ -108,14 +108,13 @@ private class TestDrillbitStatusListener implements DrillbitStatusListener {
     protected Set<DrillbitEndpoint> removed;
 
     @Override
-    public void drillbitUnregistered(
-        Set<DrillbitEndpoint> unregisteredDrillbits) {
-      removed = unregisteredDrillbits;
+    public void drillbitUnregistered(Map<DrillbitEndpoint, String> unregisteredDrillbitsUUID) {
+      removed = unregisteredDrillbitsUUID.keySet();
     }
 
     @Override
-    public void drillbitRegistered(Set<DrillbitEndpoint> registeredDrillbits) {
-      added = registeredDrillbits;
+    public void drillbitRegistered(Map<DrillbitEndpoint, String> registeredDrillbitsUUID) {
+      added = registeredDrillbitsUUID.keySet();
     }
 
     public void clear() {

exec/java-exec/src/main/java/org/apache/drill/exec/coord/ClusterCoordinator.java

@@ -25,7 +25,6 @@
 
 import java.util.Collection;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -107,17 +106,17 @@ public interface RegistrationHandle {
 
   /**
    * Actions to take when there are a set of new de-active drillbits.
-   * @param unregisteredBits
+   * @param unregisteredBitsUUID
    */
-  protected void drillbitUnregistered(Set<DrillbitEndpoint> unregisteredBits) {
+  protected void drillbitUnregistered(Map<DrillbitEndpoint, String> unregisteredBitsUUID) {
     for (DrillbitStatusListener listener : listeners.keySet()) {
-      listener.drillbitUnregistered(unregisteredBits);
+      listener.drillbitUnregistered(unregisteredBitsUUID);
     }
   }
 
-  protected void drillbitRegistered(Set<DrillbitEndpoint> registeredBits) {
+  protected void drillbitRegistered(Map<DrillbitEndpoint, String> registeredBitsUUID) {
     for (DrillbitStatusListener listener : listeners.keySet()) {
-      listener.drillbitRegistered(registeredBits);
+      listener.drillbitRegistered(registeredBitsUUID);
     }
   }
 

exec/java-exec/src/main/java/org/apache/drill/exec/coord/zk/ZKClusterCoordinator.java

@@ -31,6 +31,7 @@
 import org.apache.curator.x.discovery.ServiceInstance;
 import org.apache.curator.x.discovery.details.ServiceCacheListener;
 import org.apache.drill.common.AutoCloseables;
+import org.apache.drill.common.DrillNode;
 import org.apache.drill.common.config.DrillConfig;
 import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.coord.ClusterCoordinator;
@@ -46,11 +47,8 @@
 
 import java.io.IOException;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -66,15 +64,13 @@ public class ZKClusterCoordinator extends ClusterCoordinator {
 
   private CuratorFramework curator;
   private ServiceDiscovery<DrillbitEndpoint> discovery;
-  private volatile Collection<DrillbitEndpoint> endpoints = Collections.emptyList();
   private final String serviceName;
   private final CountDownLatch initialConnection = new CountDownLatch(1);
   private final TransientStoreFactory factory;
   private ServiceCache<DrillbitEndpoint> serviceCache;
-  private DrillbitEndpoint endpoint;
 
   // endpointsMap maps String UUID to Drillbit endpoints
-  private ConcurrentHashMap<String, DrillbitEndpoint> endpointsMap = new ConcurrentHashMap<>();
+  private  Map<String, DrillbitEndpoint> endpointsMap = new ConcurrentHashMap<>();
   private static final Pattern ZK_COMPLEX_STRING = Pattern.compile("(^.*?)/(.*)/([^/]*)$");
 
   public ZKClusterCoordinator(DrillConfig config, String connect) {
@@ -220,7 +216,8 @@ public void unregister(RegistrationHandle handle) {
    */
   public RegistrationHandle update(RegistrationHandle handle, State state) {
     ZKRegistrationHandle h = (ZKRegistrationHandle) handle;
-      try {
+    final DrillbitEndpoint endpoint;
+    try {
         endpoint = h.endpoint.toBuilder().setState(state).build();
         ServiceInstance<DrillbitEndpoint> serviceInstance = ServiceInstance.<DrillbitEndpoint>builder()
                 .name(serviceName)
@@ -231,6 +228,7 @@ public RegistrationHandle update(RegistrationHandle handle, State state) {
         Throwables.throwIfUnchecked(e);
         throw new RuntimeException(e);
       }
+      handle.setEndPoint(endpoint);
       return handle;
   }
 
@@ -282,46 +280,49 @@ public <V> TransientStore<V> getOrCreateTransientStore(final TransientStoreConfi
   private synchronized void updateEndpoints() {
     try {
       // All active bits in the Zookeeper
-      final Map<String, DrillbitEndpoint> activeEndpointsUUID = discovery.queryForInstances(serviceName).stream()
+      final Map<String, DrillbitEndpoint> UUIDtoEndpoints = discovery.queryForInstances(serviceName).stream()
         .collect(Collectors.toMap(ServiceInstance::getId, ServiceInstance::getPayload));
 
-      final Map<DrillbitEndpoint, String> UUIDtoEndpoints = activeEndpointsUUID.entrySet().stream()
-        .collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
+      final Map<DrillNode, String> activeEndpointsUUID = UUIDtoEndpoints.entrySet().stream()
+        .collect(Collectors.toMap(x -> DrillNode.create(x.getValue()), Map.Entry::getKey));
 
       // set of newly dead bits : original bits - new set of active bits.
-      Set<DrillbitEndpoint> unregisteredBits = new HashSet<>();
+      final Map<DrillbitEndpoint, String> unregisteredBits = new HashMap<>();
       // Set of newly live bits : new set of active bits - original bits.
-      Set<DrillbitEndpoint> registeredBits = new HashSet<>();
+      final Map<DrillbitEndpoint, String> registeredBits = new HashMap<>();
 
 
       // Updates the endpoints map if there is a change in state of the endpoint or with the addition
       // of new drillbit endpoints. Registered endpoints is set to newly live drillbit endpoints.
-      for (Map.Entry<String, DrillbitEndpoint> endpointToUUID : activeEndpointsUUID.entrySet()) {
-        endpointsMap.put(endpointToUUID.getKey(), endpointToUUID.getValue());
+      for (Map.Entry<String, DrillbitEndpoint> endpoint : UUIDtoEndpoints.entrySet()) {
+        // check if this bit is newly added bit
+        if (!endpointsMap.containsKey(endpoint.getKey())) {
+          registeredBits.put(endpoint.getValue(), endpoint.getKey());
+        }
+        endpointsMap.put(endpoint.getKey(), endpoint.getValue());
       }
 
       // Remove all the endpoints that are newly dead
       for ( String bitUUID: endpointsMap.keySet()) {
-        if (!activeEndpointsUUID.containsKey(bitUUID)) {
+        if (!UUIDtoEndpoints.containsKey(bitUUID)) {
           final DrillbitEndpoint unregisteredBit = endpointsMap.get(bitUUID);
-          unregisteredBits.add(unregisteredBit);
-
-          if (UUIDtoEndpoints.containsKey(unregisteredBit)) {
+          unregisteredBits.put(unregisteredBit, bitUUID);
+          final DrillNode unregisteredNode = DrillNode.create(unregisteredBit);
+          if (activeEndpointsUUID.containsKey(unregisteredNode)) {
             logger.info("Drillbit registered again with different UUID. [Details: Address: {}, UserPort: {}," +
               " PreviousUUID: {}, CurrentUUID: {}", unregisteredBit.getAddress(), unregisteredBit.getUserPort(),
-              bitUUID, UUIDtoEndpoints.get(unregisteredBit));
+              bitUUID, activeEndpointsUUID.get(unregisteredNode));
           }
           endpointsMap.remove(bitUUID);
         }
       }
-      endpoints = endpointsMap.values();
       if (logger.isDebugEnabled()) {
         StringBuilder builder = new StringBuilder();
         builder.append("Active drillbit set changed.  Now includes ");
-        builder.append(activeEndpointsUUID.size());
+        builder.append(UUIDtoEndpoints.size());
         builder.append(" total bits. New active drillbits:\n");
         builder.append("Address | User Port | Control Port | Data Port | Version | State\n");
-        for (DrillbitEndpoint bit: activeEndpointsUUID.values()) {
+        for (DrillbitEndpoint bit: UUIDtoEndpoints.values()) {
           builder.append(bit.getAddress()).append(" | ");
           builder.append(bit.getUserPort()).append(" | ");
           builder.append(bit.getControlPort()).append(" | ");