[DRILL-7191 / DRILL-7026]: RM state blob persistence in Zookeeper and Integration of Distributed queue configuration with Planner

exec/java-exec/src/main/java/org/apache/drill/exec/ExecConstants.java

@@ -691,6 +691,7 @@ private ExecConstants() {
       new OptionDescription("Indicates how long a query can wait in queue before the query fails. Range: 0-9223372036854775807"));
 
   // New Smart RM boot time configs
+  public static final String RM_WAIT_THREAD_INTERVAL = "exec.rm.wait_thread_interval";
   public static final String RM_QUERY_TAGS_KEY = "exec.rm.queryTags";
   public static final StringValidator RM_QUERY_TAGS_VALIDATOR = new StringValidator(RM_QUERY_TAGS_KEY,
     new OptionDescription("Allows user to set coma separated list of tags for all the queries submitted over a session"));

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

@@ -17,16 +17,17 @@
  */
 package org.apache.drill.exec.coord;
 
-import java.util.Collection;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
-
 import org.apache.drill.exec.coord.store.TransientStore;
 import org.apache.drill.exec.coord.store.TransientStoreConfig;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint.State;
 import org.apache.drill.exec.work.foreman.DrillbitStatusListener;
 
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
 /**
  * Pluggable interface built to manage cluster coordination. Allows Drillbit or DrillClient to register its capabilities
  * as well as understand other node's existence and capabilities.
@@ -60,6 +61,10 @@ public abstract class ClusterCoordinator implements AutoCloseable {
    */
   public abstract Collection<DrillbitEndpoint> getAvailableEndpoints();
 
+  public Map<String, DrillbitEndpoint> getAvailableEndpointsUUID() {
+    throw new UnsupportedOperationException("Only supported by Zookeeper Cluster Coordinator outside YARN");
+  }
+
   /**
    * Get a collection of ONLINE drillbit endpoints by excluding the drillbits
    * that are in QUIESCENT state (drillbits that are shutting down). Primarily used by the planner
@@ -70,6 +75,10 @@ public abstract class ClusterCoordinator implements AutoCloseable {
 
   public abstract Collection<DrillbitEndpoint> getOnlineEndPoints();
 
+  public Map<DrillbitEndpoint, String> getOnlineEndpointsUUID() {
+    throw new UnsupportedOperationException("Only supported by Zookeeper Cluster Coordinator outside YARN");
+  }
+
   public abstract RegistrationHandle update(RegistrationHandle handle, State state);
 
   public interface RegistrationHandle {
@@ -79,6 +88,8 @@ public interface RegistrationHandle {
      */
     public abstract DrillbitEndpoint getEndPoint();
 
+    public abstract String getId();
+
     public abstract void setEndPoint(DrillbitEndpoint endpoint);
   }
 

exec/java-exec/src/main/java/org/apache/drill/exec/coord/local/LocalClusterCoordinator.java

@@ -17,15 +17,6 @@
  */
 package org.apache.drill.exec.coord.local;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Map;
-import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
-
 import org.apache.drill.exec.coord.ClusterCoordinator;
 import org.apache.drill.exec.coord.DistributedSemaphore;
 import org.apache.drill.exec.coord.store.CachingTransientStoreFactory;
@@ -34,9 +25,17 @@
 import org.apache.drill.exec.coord.store.TransientStoreFactory;
 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.collect.Maps;
 
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+
 public class LocalClusterCoordinator extends ClusterCoordinator {
   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(LocalClusterCoordinator.class);
 
@@ -103,7 +102,15 @@ public RegistrationHandle update(RegistrationHandle handle, State state) {
 
   @Override
   public Collection<DrillbitEndpoint> getAvailableEndpoints() {
-    return endpoints.values();
+    return getAvailableEndpointsUUID().values();
+  }
+
+  public Map<String, DrillbitEndpoint> getAvailableEndpointsUUID() {
+    Map<String, DrillbitEndpoint> availableEndpointsUUID = new HashMap<>();
+    for (Map.Entry<RegistrationHandle, DrillbitEndpoint> entry : endpoints.entrySet()) {
+      availableEndpointsUUID.put(entry.getKey().getId(), entry.getValue());
+    }
+    return availableEndpointsUUID;
   }
 
   /**
@@ -115,19 +122,27 @@ public Collection<DrillbitEndpoint> getAvailableEndpoints() {
    */
   @Override
   public Collection<DrillbitEndpoint> getOnlineEndPoints() {
-    Collection<DrillbitEndpoint> runningEndPoints = new ArrayList<>();
-    for (DrillbitEndpoint endpoint: endpoints.values()){
-      if(isDrillbitInState(endpoint, State.ONLINE)) {
-        runningEndPoints.add(endpoint);
+    return getOnlineEndpointsUUID().keySet();
+  }
+
+  public Map<DrillbitEndpoint, String> getOnlineEndpointsUUID() {
+    Map<DrillbitEndpoint, String> onlineEndpointsUUID = new HashMap<>();
+    for (Map.Entry<RegistrationHandle, DrillbitEndpoint> entry : endpoints.entrySet()) {
+      if(isDrillbitInState(entry.getValue(), State.ONLINE)) {
+        onlineEndpointsUUID.put(entry.getValue(), entry.getKey().getId());
       }
     }
-    return runningEndPoints;
+    return onlineEndpointsUUID;
   }
 
   private class Handle implements RegistrationHandle {
     private final UUID id = UUID.randomUUID();
     private DrillbitEndpoint drillbitEndpoint;
 
+    private Handle(DrillbitEndpoint data) {
+      drillbitEndpoint = data;
+    }
+
     /**
      * Get the drillbit endpoint associated with the registration handle
      * @return drillbit endpoint
@@ -140,8 +155,9 @@ public void setEndPoint(DrillbitEndpoint endpoint) {
       this.drillbitEndpoint = endpoint;
     }
 
-    private Handle(DrillbitEndpoint data) {
-      drillbitEndpoint = data;
+    @Override
+    public String getId() {
+      return id.toString();
     }
 
     @Override

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

@@ -17,26 +17,11 @@
  */
 package org.apache.drill.exec.coord.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.ArrayList;
-import java.util.Set;
-import java.util.HashSet;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-import org.apache.curator.framework.imps.DefaultACLProvider;
-import org.apache.drill.shaded.guava.com.google.common.base.Throwables;
-import org.apache.commons.collections.keyvalue.MultiKey;
 import org.apache.curator.RetryPolicy;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.framework.api.ACLProvider;
+import org.apache.curator.framework.imps.DefaultACLProvider;
 import org.apache.curator.framework.state.ConnectionState;
 import org.apache.curator.framework.state.ConnectionStateListener;
 import org.apache.curator.retry.RetryNTimes;
@@ -57,7 +42,21 @@
 import org.apache.drill.exec.coord.store.TransientStoreFactory;
 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.Function;
+import org.apache.drill.shaded.guava.com.google.common.base.Throwables;
+
+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;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * Manages cluster coordination utilizing zookeeper. *
@@ -74,8 +73,8 @@ public class ZKClusterCoordinator extends ClusterCoordinator {
   private ServiceCache<DrillbitEndpoint> serviceCache;
   private DrillbitEndpoint endpoint;
 
-  // endpointsMap maps Multikey( comprises of endoint address and port) to Drillbit endpoints
-  private ConcurrentHashMap<MultiKey, DrillbitEndpoint> endpointsMap = new ConcurrentHashMap<MultiKey,DrillbitEndpoint>();
+  // endpointsMap maps String UUID to Drillbit endpoints
+  private ConcurrentHashMap<String, DrillbitEndpoint> endpointsMap = new ConcurrentHashMap<>();
   private static final Pattern ZK_COMPLEX_STRING = Pattern.compile("(^.*?)/(.*)/([^/]*)$");
 
   public ZKClusterCoordinator(DrillConfig config, String connect) {
@@ -237,7 +236,12 @@ public RegistrationHandle update(RegistrationHandle handle, State state) {
 
   @Override
   public Collection<DrillbitEndpoint> getAvailableEndpoints() {
-    return this.endpoints;
+    return getAvailableEndpointsUUID().values();
+  }
+
+  @Override
+  public Map<String, DrillbitEndpoint> getAvailableEndpointsUUID() {
+    return this.endpointsMap;
   }
 
   /*
@@ -249,14 +253,19 @@ public Collection<DrillbitEndpoint> getAvailableEndpoints() {
    */
   @Override
   public Collection<DrillbitEndpoint> getOnlineEndPoints() {
-    Collection<DrillbitEndpoint> runningEndPoints = new ArrayList<>();
-    for (DrillbitEndpoint endpoint: endpoints){
-      if(isDrillbitInState(endpoint, State.ONLINE)) {
-        runningEndPoints.add(endpoint);
+    return getOnlineEndpointsUUID().keySet();
+  }
+
+  @Override
+  public Map<DrillbitEndpoint, String> getOnlineEndpointsUUID() {
+    Map<DrillbitEndpoint, String> onlineEndpointsUUID = new HashMap<>();
+    for (Map.Entry<String, DrillbitEndpoint> endpointEntry : endpointsMap.entrySet()) {
+      if (isDrillbitInState(endpointEntry.getValue(), State.ONLINE)) {
+        onlineEndpointsUUID.put(endpointEntry.getValue(), endpointEntry.getKey());
       }
     }
-    logger.debug("Online endpoints in ZK are" + runningEndPoints.toString());
-    return runningEndPoints;
+    logger.debug("Online endpoints in ZK are" + onlineEndpointsUUID.keySet().toString());
+    return onlineEndpointsUUID;
   }
 
   @Override
@@ -273,14 +282,11 @@ public <V> TransientStore<V> getOrCreateTransientStore(final TransientStoreConfi
   private synchronized void updateEndpoints() {
     try {
       // All active bits in the Zookeeper
-      Collection<DrillbitEndpoint> newDrillbitSet =
-      transform(discovery.queryForInstances(serviceName),
-        new Function<ServiceInstance<DrillbitEndpoint>, DrillbitEndpoint>() {
-          @Override
-          public DrillbitEndpoint apply(ServiceInstance<DrillbitEndpoint> input) {
-            return input.getPayload();
-          }
-        });
+      final Map<String, DrillbitEndpoint> activeEndpointsUUID = 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));
 
       // set of newly dead bits : original bits - new set of active bits.
       Set<DrillbitEndpoint> unregisteredBits = new HashSet<>();
@@ -290,29 +296,32 @@ public DrillbitEndpoint apply(ServiceInstance<DrillbitEndpoint> input) {
 
       // 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 ( DrillbitEndpoint endpoint : newDrillbitSet) {
-        String endpointAddress = endpoint.getAddress();
-        int endpointPort = endpoint.getUserPort();
-        if (! endpointsMap.containsKey(new MultiKey(endpointAddress, endpointPort))) {
-          registeredBits.add(endpoint);
-        }
-        endpointsMap.put(new MultiKey(endpointAddress, endpointPort),endpoint);
+      for (Map.Entry<String, DrillbitEndpoint> endpointToUUID : activeEndpointsUUID.entrySet()) {
+        endpointsMap.put(endpointToUUID.getKey(), endpointToUUID.getValue());
       }
+
       // Remove all the endpoints that are newly dead
-      for ( MultiKey key: endpointsMap.keySet()) {
-        if(!newDrillbitSet.contains(endpointsMap.get(key))) {
-          unregisteredBits.add(endpointsMap.get(key));
-          endpointsMap.remove(key);
+      for ( String bitUUID: endpointsMap.keySet()) {
+        if (!activeEndpointsUUID.containsKey(bitUUID)) {
+          final DrillbitEndpoint unregisteredBit = endpointsMap.get(bitUUID);
+          unregisteredBits.add(unregisteredBit);
+
+          if (UUIDtoEndpoints.containsKey(unregisteredBit)) {
+            logger.info("Drillbit registered again with different UUID. [Details: Address: {}, UserPort: {}," +
+              " PreviousUUID: {}, CurrentUUID: {}", unregisteredBit.getAddress(), unregisteredBit.getUserPort(),
+              bitUUID, UUIDtoEndpoints.get(unregisteredBit));
+          }
+          endpointsMap.remove(bitUUID);
         }
       }
       endpoints = endpointsMap.values();
       if (logger.isDebugEnabled()) {
         StringBuilder builder = new StringBuilder();
         builder.append("Active drillbit set changed.  Now includes ");
-        builder.append(newDrillbitSet.size());
+        builder.append(activeEndpointsUUID.size());
         builder.append(" total bits. New active drillbits:\n");
         builder.append("Address | User Port | Control Port | Data Port | Version | State\n");
-        for (DrillbitEndpoint bit: newDrillbitSet) {
+        for (DrillbitEndpoint bit: activeEndpointsUUID.values()) {
           builder.append(bit.getAddress()).append(" | ");
           builder.append(bit.getUserPort()).append(" | ");
           builder.append(bit.getControlPort()).append(" | ");

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

@@ -26,19 +26,24 @@ public class ZKRegistrationHandle implements RegistrationHandle {
   public final String id;
   public DrillbitEndpoint endpoint;
 
+  public ZKRegistrationHandle(String id, DrillbitEndpoint endpoint) {
+    super();
+    this.id = id;
+    this.endpoint = endpoint;
+  }
+
   public DrillbitEndpoint getEndPoint() {
     return endpoint;
   }
 
+  public String getId() {
+    return id;
+  }
+
   @Override
   public void setEndPoint(DrillbitEndpoint endpoint) {
     this.endpoint = endpoint;
   }
 
-  public ZKRegistrationHandle(String id, DrillbitEndpoint endpoint) {
-    super();
-    this.id = id;
-    this.endpoint = endpoint;
-  }
 
 }