Use parallel calls when re-preparing statement on other hosts (JAVA-725).

Alexandre Dutra · olim7t · commit 17fb2f54d6ac · 2015-08-28T11:36:21.000+02:00
diff --git a/changelog/README.md b/changelog/README.md
@@ -41,6 +41,8 @@
 - [improvement] Allow hosts to remain partially up (JAVA-544)
 - [improvement] Remove internal blocking calls and expose async session
   creation (JAVA-821, JAVA-822)
+- [improvement] Use parallel calls when re-preparing statement on other
+  hosts (JAVA-725)
 
 Merged from 2.0.10_fixes branch:
 
diff --git a/driver-core/src/main/java/com/datastax/driver/core/SessionManager.java b/driver-core/src/main/java/com/datastax/driver/core/SessionManager.java
@@ -22,13 +22,15 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import com.google.common.base.Function;
+import com.google.common.base.Functions;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.datastax.driver.core.Message.Response;
 import com.datastax.driver.core.exceptions.DriverInternalError;
 import com.datastax.driver.core.exceptions.InvalidQueryException;
 import com.datastax.driver.core.exceptions.UnsupportedFeatureException;
@@ -179,8 +181,8 @@ public Session.State getState() {
     }
 
     private ListenableFuture<PreparedStatement> toPreparedStatement(final String query, final Connection.Future future) {
-        return Futures.transform(future, new Function<Message.Response, PreparedStatement>() {
-            public PreparedStatement apply(Message.Response response) {
+        return Futures.transform(future, new AsyncFunction<Response, PreparedStatement>() {
+            public ListenableFuture<PreparedStatement> apply(Response response) {
                 switch (response.type) {
                     case RESULT:
                         Responses.Result rm = (Responses.Result)response;
@@ -189,29 +191,26 @@ public PreparedStatement apply(Message.Response response) {
                                 Responses.Result.Prepared pmsg = (Responses.Result.Prepared)rm;
                                 PreparedStatement stmt = DefaultPreparedStatement.fromMessage(pmsg, cluster.getMetadata(), query, poolsState.keyspace);
                                 stmt = cluster.manager.addPrepared(stmt);
-                                if (cluster.getConfiguration().getQueryOptions().isPrepareOnAllHosts()){
-                                    try {
-                                        // All Sessions are connected to the same nodes so it's enough to prepare only the nodes of this session.
-                                        // If that changes, we'll have to make sure this propagate to other sessions too.
-                                        prepare(stmt.getQueryString(), future.getAddress());
-                                    } catch (InterruptedException e) {
-                                        Thread.currentThread().interrupt();
-                                        // This method doesn't propagate interruption, at least not for now. However, if we've
-                                        // interrupted preparing queries on other node it's not a problem as we'll re-prepare
-                                        // later if need be. So just ignore.
-                                    }
+                                if (cluster.getConfiguration().getQueryOptions().isPrepareOnAllHosts()) {
+                                    // All Sessions are connected to the same nodes so it's enough to prepare only the nodes of this session.
+                                    // If that changes, we'll have to make sure this propagate to other sessions too.
+                                    return prepare(stmt, future.getAddress());
+                                } else {
+                                    return Futures.immediateFuture(stmt);
                                 }
-                                return stmt;
                             default:
-                                throw new DriverInternalError(String.format("%s response received when prepared statement was expected", rm.kind));
+                                return Futures.immediateFailedFuture(
+                                    new DriverInternalError(String.format("%s response received when prepared statement was expected", rm.kind)));
                         }
                     case ERROR:
-                        throw ((Responses.Error)response).asException(future.getAddress());
+                        return Futures.immediateFailedFuture(
+                            ((Responses.Error)response).asException(future.getAddress()));
                     default:
-                        throw new DriverInternalError(String.format("%s response received when prepared statement was expected", response.type));
+                        return Futures.immediateFailedFuture(
+                            new DriverInternalError(String.format("%s response received when prepared statement was expected", response.type)));
                 }
             }
-        }, executor()); // Since the transformation involves querying other nodes, we should not do that in an I/O thread
+        });
     }
 
     Connection.Factory connectionFactory() {
@@ -581,35 +580,42 @@ public void run() {
             }, executor());
     }
 
-    private void prepare(String query, InetSocketAddress toExclude) throws InterruptedException {
-        for (Map.Entry<Host, HostConnectionPool> entry : pools.entrySet()) {
+    private ListenableFuture<PreparedStatement> prepare(final PreparedStatement statement, InetSocketAddress toExclude) {
+        final String query = statement.getQueryString();
+        List<ListenableFuture<Response>> futures = Lists.newArrayListWithExpectedSize(pools.size());
+        for (final Map.Entry<Host, HostConnectionPool> entry : pools.entrySet()) {
             if (entry.getKey().getSocketAddress().equals(toExclude))
                 continue;
 
-            // Let's not wait too long if we can't get a connection. Things
-            // will fix themselves once the user tries a query anyway.
-            Connection c = null;
-            boolean timedOut = false;
             try {
-                c = entry.getValue().borrowConnection(200, TimeUnit.MILLISECONDS);
-                c.write(new Requests.Prepare(query)).get();
-            } catch (ConnectionException e) {
+                // Preparing is not critical: if it fails, it will fix itself later when the user tries to execute
+                // the prepared query. So don't block if no connection is available, simply abort.
+                final Connection c = entry.getValue().borrowConnection(0, TimeUnit.MILLISECONDS);
+                ListenableFuture<Response> future = c.write(new Requests.Prepare(query));
+                Futures.addCallback(future, new FutureCallback<Response>() {
+                    @Override
+                    public void onSuccess(Response result) {
+                        c.release();
+                    }
+
+                    @Override
+                    public void onFailure(Throwable t) {
+                        logger.debug(String.format("Unexpected error while preparing query (%s) on %s", query, entry.getKey()), t);
+
+                        // If the query timed out, that already released the connection, otherwise do it now
+                        if (!(t instanceof OperationTimedOutException))
+                            c.release();
+                    }
+                });
+                futures.add(future);
+            } catch (Exception e) {
                 // Again, not being able to prepare the query right now is no big deal, so just ignore
-            } catch (BusyConnectionException e) {
-                // Same as above
-            } catch (TimeoutException e) {
-                // Same as above
-            } catch (ExecutionException e) {
-                // We shouldn't really get exception while preparing a
-                // query, so log this (but ignore otherwise as it's not a big deal)
-                logger.error(String.format("Unexpected error while preparing query (%s) on %s", query, entry.getKey()), e);
-                // If the query timed out, that already released the connection
-                timedOut = e.getCause() instanceof OperationTimedOutException;
-            } finally {
-                if (c != null && !timedOut)
-                    c.release();
             }
         }
+        // Return the statement when all futures are done
+        return Futures.transform(
+            Futures.successfulAsList(futures),
+            Functions.constant(statement));
     }
 
     ResultSetFuture executeQuery(Message.Request msg, Statement statement) {
