Merge branch '2.1'

Author: dlmarion

core/src/main/java/org/apache/accumulo/core/clientImpl/ThriftTransportKey.java

@@ -24,12 +24,14 @@
 
 import org.apache.accumulo.core.rpc.SaslConnectionParams;
 import org.apache.accumulo.core.rpc.SslConnectionParams;
+import org.apache.accumulo.core.rpc.clients.ThriftClientTypes;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.net.HostAndPort;
 
 @VisibleForTesting
 public class ThriftTransportKey {
+  private final ThriftClientTypes<?> type;
   private final HostAndPort server;
   private final long timeout;
   private final SslConnectionParams sslParams;
@@ -38,16 +40,18 @@ public class ThriftTransportKey {
   private final int hash;
 
   @VisibleForTesting
-  public ThriftTransportKey(HostAndPort server, long timeout, ClientContext context) {
-    this(server, timeout, context.getClientSslParams(), context.getSaslParams());
+  public ThriftTransportKey(ThriftClientTypes<?> type, HostAndPort server, long timeout,
+      ClientContext context) {
+    this(type, server, timeout, context.getClientSslParams(), context.getSaslParams());
   }
 
   /**
    * Visible only for testing
    */
-  ThriftTransportKey(HostAndPort server, long timeout, SslConnectionParams sslParams,
-      SaslConnectionParams saslParams) {
+  ThriftTransportKey(ThriftClientTypes<?> type, HostAndPort server, long timeout,
+      SslConnectionParams sslParams, SaslConnectionParams saslParams) {
     requireNonNull(server, "location is null");
+    this.type = type;
     this.server = server;
     this.timeout = timeout;
     this.sslParams = sslParams;
@@ -56,14 +60,21 @@ public ThriftTransportKey(HostAndPort server, long timeout, ClientContext contex
       // TSasl and TSSL transport factories don't play nicely together
       throw new IllegalArgumentException("Cannot use both SSL and SASL thrift transports");
     }
-    this.hash = Objects.hash(server, timeout, sslParams, saslParams);
+    this.hash = Objects.hash(type, server, timeout, sslParams, saslParams);
   }
 
-  HostAndPort getServer() {
+  @VisibleForTesting
+  public ThriftClientTypes<?> getType() {
+    return type;
+  }
+
+  @VisibleForTesting
+  public HostAndPort getServer() {
     return server;
   }
 
-  long getTimeout() {
+  @VisibleForTesting
+  public long getTimeout() {
     return timeout;
   }
 
@@ -81,7 +92,7 @@ public boolean equals(Object o) {
       return false;
     }
     ThriftTransportKey ttk = (ThriftTransportKey) o;
-    return server.equals(ttk.server) && timeout == ttk.timeout
+    return type.equals(ttk.type) && server.equals(ttk.server) && timeout == ttk.timeout
         && (!isSsl() || (ttk.isSsl() && sslParams.equals(ttk.sslParams)))
         && (!isSasl() || (ttk.isSasl() && saslParams.equals(ttk.saslParams)));
   }
@@ -99,7 +110,7 @@ public String toString() {
     } else if (isSasl()) {
       prefix = saslParams + ":";
     }
-    return prefix + server + " (" + timeout + ")";
+    return prefix + type + ":" + server + " (" + timeout + ")";
   }
 
   public SslConnectionParams getSslParams() {

core/src/main/java/org/apache/accumulo/core/clientImpl/ThriftTransportPool.java

@@ -41,6 +41,7 @@
 import java.util.function.Supplier;
 
 import org.apache.accumulo.core.rpc.ThriftUtil;
+import org.apache.accumulo.core.rpc.clients.ThriftClientTypes;
 import org.apache.accumulo.core.util.Pair;
 import org.apache.accumulo.core.util.threads.Threads;
 import org.apache.thrift.TConfiguration;
@@ -49,7 +50,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.net.HostAndPort;
 
@@ -109,71 +109,40 @@ static ThriftTransportPool startNew(LongSupplier maxAgeMillis) {
     return pool;
   }
 
-  public TTransport getTransport(HostAndPort location, long milliseconds, ClientContext context)
-      throws TTransportException {
-    ThriftTransportKey cacheKey = new ThriftTransportKey(location, milliseconds, context);
+  public TTransport getTransport(ThriftClientTypes<?> type, HostAndPort location, long milliseconds,
+      ClientContext context, boolean preferCached) throws TTransportException {
 
-    CachedConnection connection = connectionPool.reserveAny(cacheKey);
-
-    if (connection != null) {
-      log.trace("Using existing connection to {}", cacheKey.getServer());
-      return connection.transport;
-    } else {
-      return createNewTransport(cacheKey);
+    ThriftTransportKey cacheKey = new ThriftTransportKey(type, location, milliseconds, context);
+    if (preferCached) {
+      CachedConnection connection = connectionPool.reserveAny(cacheKey);
+      if (connection != null) {
+        log.trace("Using existing connection to {}", cacheKey.getServer());
+        return connection.transport;
+      }
     }
+    return createNewTransport(cacheKey);
   }
 
-  @VisibleForTesting
-  public Pair<String,TTransport> getAnyTransport(List<ThriftTransportKey> servers,
-      boolean preferCachedConnection) throws TTransportException {
-
-    servers = new ArrayList<>(servers);
-
-    if (preferCachedConnection) {
-      HashSet<ThriftTransportKey> serversSet = new HashSet<>(servers);
-
-      // randomly pick a server from the connection cache
-      serversSet.retainAll(connectionPool.getThriftTransportKeys());
-
-      if (!serversSet.isEmpty()) {
-        ArrayList<ThriftTransportKey> cachedServers = new ArrayList<>(serversSet);
-        Collections.shuffle(cachedServers, RANDOM.get());
-
-        for (ThriftTransportKey ttk : cachedServers) {
-          CachedConnection connection = connectionPool.reserveAny(ttk);
-          if (connection != null) {
-            final String serverAddr = ttk.getServer().toString();
-            log.trace("Using existing connection to {}", serverAddr);
-            return new Pair<>(serverAddr, connection.transport);
-          }
-
-        }
+  public Pair<String,TTransport> getAnyCachedTransport(ThriftClientTypes<?> type) {
+    final List<ThriftTransportKey> serversSet = new ArrayList<>();
+    for (ThriftTransportKey ttk : connectionPool.getThriftTransportKeys()) {
+      if (ttk.getType().equals(type)) {
+        serversSet.add(ttk);
       }
     }
-
-    int retryCount = 0;
-    while (!servers.isEmpty() && retryCount < 10) {
-
-      int index = RANDOM.get().nextInt(servers.size());
-      ThriftTransportKey ttk = servers.get(index);
-
-      if (preferCachedConnection) {
-        CachedConnection connection = connectionPool.reserveAnyIfPresent(ttk);
-        if (connection != null) {
-          return new Pair<>(ttk.getServer().toString(), connection.transport);
-        }
-      }
-
-      try {
-        return new Pair<>(ttk.getServer().toString(), createNewTransport(ttk));
-      } catch (TTransportException tte) {
-        log.debug("Failed to connect to {}", servers.get(index), tte);
-        servers.remove(index);
-        retryCount++;
+    if (serversSet.isEmpty()) {
+      return null;
+    }
+    Collections.shuffle(serversSet, RANDOM.get());
+    for (ThriftTransportKey ttk : serversSet) {
+      CachedConnection connection = connectionPool.reserveAny(ttk);
+      if (connection != null) {
+        final String serverAddr = ttk.getServer().toString();
+        log.trace("Using existing connection to {}", serverAddr);
+        return new Pair<>(serverAddr, connection.transport);
       }
     }
-
-    throw new TTransportException("Failed to connect to a server");
+    return null;
   }
 
   private TTransport createNewTransport(ThriftTransportKey cacheKey) throws TTransportException {
@@ -389,27 +358,6 @@ CachedConnection reserveAny(final ThriftTransportKey key) {
       return executeWithinLock(key, connections::reserveAny);
     }
 
-    /**
-     * Reserve and return a new {@link CachedConnection} from the {@link CachedConnections} mapped
-     * to the specified transport key. If a {@link CachedConnections} is not found, null will be
-     * returned.
-     *
-     * <p>
-     *
-     * This operation locks access to the mapping for the key in {@link ConnectionPool#connections}
-     * until the operation completes.
-     *
-     * @param key the transport key
-     * @return the reserved {@link CachedConnection}, or null if none were available.
-     */
-    CachedConnection reserveAnyIfPresent(final ThriftTransportKey key) {
-      // It's possible that multiple locks from executeWithinLock will overlap with a single lock
-      // inside the ConcurrentHashMap which can unnecessarily block threads. Access the
-      // ConcurrentHashMap outside of executeWithinLock to prevent this.
-      var connections = getCachedConnections(key);
-      return connections == null ? null : executeWithinLock(key, connections::reserveAny);
-    }
-
     /**
      * Puts the specified connection into the reserved map of the {@link CachedConnections} for the
      * specified transport key. If a {@link CachedConnections} is not found, one will be created.
@@ -515,10 +463,6 @@ Lock getLock(final ThriftTransportKey key) {
       return lock;
     }
 
-    CachedConnections getCachedConnections(final ThriftTransportKey key) {
-      return connections.get(key);
-    }
-
     CachedConnections getOrCreateCachedConnections(final ThriftTransportKey key) {
       return connections.computeIfAbsent(key, k -> new CachedConnections());
     }

core/src/main/java/org/apache/accumulo/core/rpc/ThriftUtil.java

@@ -111,8 +111,8 @@ public static <T extends TServiceClient> T getClientNoTimeout(ThriftClientTypes<
    */
   public static <T extends TServiceClient> T getClient(ThriftClientTypes<T> type,
       HostAndPort address, ClientContext context) throws TTransportException {
-    TTransport transport = context.getTransportPool().getTransport(address,
-        context.getClientTimeoutInMillis(), context);
+    TTransport transport = context.getTransportPool().getTransport(type, address,
+        context.getClientTimeoutInMillis(), context, true);
     return createClient(type, transport);
   }
 
@@ -127,7 +127,8 @@ public static <T extends TServiceClient> T getClient(ThriftClientTypes<T> type,
    */
   public static <T extends TServiceClient> T getClient(ThriftClientTypes<T> type,
       HostAndPort address, ClientContext context, long timeout) throws TTransportException {
-    TTransport transport = context.getTransportPool().getTransport(address, timeout, context);
+    TTransport transport =
+        context.getTransportPool().getTransport(type, address, timeout, context, true);
     return createClient(type, transport);
   }