Fixing concurrent modification exceptions in the pooling logic.

fixes elastic#894

Author: jbaiera

mr/src/itest/java/org/elasticsearch/hadoop/rest/pooling/AbstractTransportPoolTest.java

@@ -0,0 +1,79 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you 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 org.elasticsearch.hadoop.rest.pooling;
+
+import org.elasticsearch.hadoop.cfg.Settings;
+import org.elasticsearch.hadoop.rest.Transport;
+import org.elasticsearch.hadoop.util.SettingsUtils;
+import org.elasticsearch.hadoop.util.TestSettings;
+import org.junit.Test;
+
+import java.util.UUID;
+
+import static org.elasticsearch.hadoop.cfg.ConfigurationOptions.ES_NET_TRANSPORT_POOLING_EXPIRATION_TIMEOUT;
+
+public class AbstractTransportPoolTest {
+
+    @Test
+    public void removeOldConnections() throws Exception {
+        Settings settings = new TestSettings();
+        settings.setProperty(ES_NET_TRANSPORT_POOLING_EXPIRATION_TIMEOUT, "2s");
+
+        String host = SettingsUtils.discoveredOrDeclaredNodes(settings).get(0);
+
+        TransportPool pool = new TransportPool(UUID.randomUUID().toString(), host, settings);
+
+        Transport transport1 = null;
+        Transport transport2 = null;
+        Transport transport3 = null;
+
+        try {
+            // Checkout three transports all at once, this should create three pooled transports.
+            transport1 = pool.borrowTransport();
+            transport2 = pool.borrowTransport();
+            transport3 = pool.borrowTransport();
+
+            // Close two of them.
+            transport1.close();
+            transport2.close();
+
+            // Wait the amount of time to close.
+            Thread.sleep(settings.getTransportPoolingExpirationTimeout().millis() + 1000L);
+
+            // Will need to remove 2 connections at this point
+            pool.removeOldConnections();
+
+        } finally {
+            // Close everything
+            if (transport1 != null) {
+                transport1.close();
+            }
+
+            if (transport2 != null) {
+                transport2.close();
+            }
+
+            if (transport3 != null) {
+                transport3.close();
+            }
+        }
+    }
+
+}

mr/src/itest/java/org/elasticsearch/hadoop/rest/pooling/PoolingSuite.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you 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 org.elasticsearch.hadoop.rest.pooling;
+
+import org.elasticsearch.hadoop.LocalEs;
+import org.junit.ClassRule;
+import org.junit.rules.ExternalResource;
+import org.junit.runner.RunWith;
+import org.junit.runners.Suite;
+
+@RunWith(Suite.class)
+@Suite.SuiteClasses({ AbstractTransportPoolTest.class })
+public class PoolingSuite {
+
+    @ClassRule
+    public static ExternalResource resource = new LocalEs();
+
+}

mr/src/main/java/org/elasticsearch/hadoop/rest/pooling/TransportPool.java

@@ -31,7 +31,9 @@
 import org.elasticsearch.hadoop.util.unit.TimeValue;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import static org.elasticsearch.hadoop.rest.Request.Method.GET;
@@ -115,20 +117,37 @@ private void release(PooledTransport transport) {
      */
     synchronized Transport borrowTransport() {
         long now = System.currentTimeMillis();
+
+        List<PooledTransport> garbageTransports = new ArrayList<PooledTransport>();
+        PooledTransport candidate = null;
+
+        // Grab a transport
         for (Map.Entry<PooledTransport, Long> entry : idle.entrySet()) {
             PooledTransport transport = entry.getKey();
             if (validate(transport)) {
-                idle.remove(transport);
-                leased.put(transport, now);
-                return new LeasedTransport(transport, this);
+                candidate = transport;
+                break;
             } else {
-                idle.remove(transport);
-                release(transport);
+                garbageTransports.add(transport);
             }
         }
-        PooledTransport transport = create();
-        leased.put(transport, now);
-        return new LeasedTransport(transport, this);
+
+        // Remove any dead connections found
+        for (PooledTransport transport : garbageTransports) {
+            idle.remove(transport);
+            release(transport);
+        }
+
+        // Create the connection if we didn't find any, remove it from the pool if we did.
+        if (candidate == null) {
+            candidate = create();
+        } else {
+            idle.remove(candidate);
+        }
+
+        // Lease.
+        leased.put(candidate, now);
+        return new LeasedTransport(candidate, this);
     }
 
     /**
@@ -165,6 +184,8 @@ private synchronized void returnTransport(Transport returning) {
     synchronized int removeOldConnections() {
         long now = System.currentTimeMillis();
         long expirationTime = now - idleTransportTimeout.millis();
+
+        List<PooledTransport> removeFromIdle = new ArrayList<PooledTransport>();
         for (Map.Entry<PooledTransport, Long> idleEntry : idle.entrySet()) {
             long lastUsed = idleEntry.getValue();
             if (lastUsed < expirationTime) {
@@ -175,9 +196,14 @@ synchronized int removeOldConnections() {
                             + idleTransportTimeout + "] ago.");
                 }
                 release(removed);
-                idle.remove(removed);
+                removeFromIdle.add(removed);
             }
         }
+
+        for (PooledTransport toRemove : removeFromIdle) {
+            idle.remove(toRemove);
+        }
+
         return idle.size() + leased.size();
     }