HADOOP-12475. Replace guava Cache with ConcurrentHashMap for caching Connection in ipc Client (Walter Su via sjlee)

hadoop-common-project/hadoop-common/CHANGES.txt

@@ -510,6 +510,9 @@ Trunk (Unreleased)
     HADOOP-12364. Deleting pid file after stop is causing the daemons to
     keep restarting (Siqi Li via aw)
 
+    HADOOP-12475. Replace guava Cache with ConcurrentHashMap for caching
+    Connection in ipc Client (Walter Su via sjlee)
+
   OPTIMIZATIONS
 
     HADOOP-7761. Improve the performance of raw comparisons. (todd)

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Client.java

@@ -43,7 +43,8 @@ import java.util.Iterator;
 import java.util.Map.Entry;
 import java.util.Random;
 import java.util.Set;
-import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -57,8 +58,6 @@ import java.util.concurrent.atomic.AtomicLong;
 import javax.net.SocketFactory;
 import javax.security.sasl.Sasl;
 
-import com.google.common.cache.Cache;
-import com.google.common.cache.CacheBuilder;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.classification.InterfaceAudience;
@@ -128,8 +127,8 @@ public class Client {
     retryCount.set(rc);
   }
 
-  private final Cache<ConnectionId, Connection> connections =
-      CacheBuilder.newBuilder().build();
+  private ConcurrentMap<ConnectionId, Connection> connections =
+      new ConcurrentHashMap<>();
 
   private Class<? extends Writable> valueClass;   // class of call values
   private AtomicBoolean running = new AtomicBoolean(true); // if client runs
@@ -1178,7 +1177,10 @@ public class Client {
         return;
       }
 
-      connections.invalidate(remoteId);
+      // We have marked this connection as closed. Other thread could have
+      // already known it and replace this closedConnection with a new one.
+      // We should only remove this closedConnection.
+      connections.remove(remoteId, this);
 
       // close the streams and therefore the socket
       IOUtils.closeStream(out);
@@ -1265,12 +1267,12 @@ public class Client {
     }
     
     // wake up all connections
-    for (Connection conn : connections.asMap().values()) {
+    for (Connection conn : connections.values()) {
       conn.interrupt();
     }
     
     // wait until all connections are closed
-    while (connections.size() > 0) {
+    while (!connections.isEmpty()) {
       try {
         Thread.sleep(100);
       } catch (InterruptedException e) {
@@ -1289,7 +1291,6 @@ public class Client {
     ConnectionId remoteId = ConnectionId.getConnectionId(address, null, null, 0,
         conf);
     return call(RpcKind.RPC_BUILTIN, param, remoteId);
-
   }
 
   /**
@@ -1465,14 +1466,13 @@ public class Client {
   @InterfaceAudience.Private
   @InterfaceStability.Unstable
   Set<ConnectionId> getConnectionIds() {
-    return connections.asMap().keySet();
+    return connections.keySet();
   }
   
   /** Get a connection from the pool, or create a new one and add it to the
    * pool.  Connections to a given ConnectionId are reused. */
-  private Connection getConnection(
-      final ConnectionId remoteId,
-      Call call, final int serviceClass, AtomicBoolean fallbackToSimpleAuth)
+  private Connection getConnection(ConnectionId remoteId,
+      Call call, int serviceClass, AtomicBoolean fallbackToSimpleAuth)
       throws IOException {
     if (!running.get()) {
       // the client is stopped
@@ -1484,33 +1484,29 @@ public class Client {
      * refs for keys in HashMap properly. For now its ok.
      */
     while (true) {
-      try {
-        connection = connections.get(remoteId, new Callable<Connection>() {
-          @Override
-          public Connection call() throws Exception {
-            return new Connection(remoteId, serviceClass);
-          }
-        });
-      } catch (ExecutionException e) {
-        Throwable cause = e.getCause();
-        // the underlying exception should normally be IOException
-        if (cause instanceof IOException) {
-          throw (IOException) cause;
-        } else {
-          throw new IOException(cause);
-        }
-      }
-      if (connection.addCall(call)) {
-        break;
-      } else {
-        connections.invalidate(remoteId);
+      // These lines below can be shorten with computeIfAbsent in Java8
+      connection = connections.get(remoteId);
+      if (connection == null) {
+        connection = new Connection(remoteId, serviceClass);
+        Connection existing = connections.putIfAbsent(remoteId, connection);
+        if (existing != null) {
+          connection = existing;
         }
       }
 
-    //we don't invoke the method below inside "synchronized (connections)"
-    //block above. The reason for that is if the server happens to be slow,
-    //it will take longer to establish a connection and that will slow the
-    //entire system down.
+      if (connection.addCall(call)) {
+        break;
+      } else {
+        // This connection is closed, should be removed. But other thread could
+        // have already known this closedConnection, and replace it with a new
+        // connection. So we should call conditional remove to make sure we only
+        // remove this closedConnection.
+        connections.remove(remoteId, connection);
+      }
+    }
+
+    // If the server happens to be slow, the method below will take longer to
+    // establish a connection.
     connection.setupIOstreams(fallbackToSimpleAuth);
     return connection;
   }