3.x: handle synchronous request retry failures

Author: dkropachev

driver-core/src/main/java/com/datastax/driver/core/Connection.java

@@ -791,6 +791,11 @@ boolean isDefunct() {
     return isDefunct.get();
   }
 
+  @VisibleForTesting
+  void markDefunctForTest() {
+    isDefunct.set(true);
+  }
+
   int maxAvailableStreams() {
     return dispatcher.streamIdHandler.maxAvailableStreams();
   }

driver-core/src/main/java/com/datastax/driver/core/RequestHandler.java

@@ -812,7 +812,21 @@ public void onSet(
                     toPrepare.getQueryKeyspace(),
                     connection.endPoint);
 
-                write(connection, prepareAndRetry(toPrepare.getQueryString()));
+                try {
+                  write(connection, prepareAndRetry(toPrepare.getQueryString()));
+                } catch (ConnectionException e) {
+                  if (metricsEnabled()) metrics().getErrorMetrics().getConnectionErrors().inc();
+                  connection.release();
+                  logError(connection.endPoint, e);
+                  retry(false, null);
+                } catch (BusyConnectionException e) {
+                  connection.release(true);
+                  logError(connection.endPoint, e);
+                  retry(false, null);
+                } catch (RuntimeException e) {
+                  connection.release();
+                  throw e;
+                }
                 // we're done for now, the prepareAndRetry callback will handle the rest
                 return;
               case READ_FAILURE:

driver-core/src/test/java/com/datastax/driver/core/RequestHandlerTest.java

@@ -22,20 +22,133 @@
 package com.datastax.driver.core;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyLong;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.scassandra.http.client.PrimingRequest.preparedStatementBuilder;
+import static org.scassandra.http.client.PrimingRequest.queryBuilder;
 import static org.scassandra.http.client.PrimingRequest.then;
+import static org.scassandra.http.client.Result.unprepared;
+import static org.testng.Assert.fail;
 
+import com.datastax.driver.core.exceptions.ConnectionException;
+import com.datastax.driver.core.exceptions.NoHostAvailableException;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.Futures;
+import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import org.scassandra.Scassandra;
-import org.scassandra.http.client.PrimingRequest;
+import org.scassandra.http.client.UnpreparedConfig;
 import org.testng.annotations.Test;
 
 public class RequestHandlerTest {
 
+  @Test(groups = "short")
+  public void should_try_next_host_when_borrow_connection_future_fails() throws Exception {
+    ScassandraCluster scassandra = ScassandraCluster.builder().withNodes(2).build();
+    Cluster cluster = null;
+
+    try {
+      scassandra.init();
+      scassandra
+          .node(2)
+          .primingClient()
+          .prime(
+              queryBuilder()
+                  .withQuery("mock query")
+                  .withThen(then().withRows(row("result", "result2")))
+                  .build());
+
+      cluster = newCluster(scassandra);
+      Session session = cluster.connect();
+      Host host1 = TestUtils.findHost(cluster, 1);
+      Host host2 = TestUtils.findHost(cluster, 2);
+
+      failBorrowingFrom(session, host1, new ConnectionException(host1.getEndPoint(), "mock"));
+
+      ResultSet rs = session.execute("mock query");
+
+      assertThat(rs.one().getString("result")).isEqualTo("result2");
+      assertThat(rs.getExecutionInfo().getQueriedHost()).isEqualTo(host2);
+    } finally {
+      if (cluster != null) cluster.close();
+      scassandra.stop();
+    }
+  }
+
+  @Test(groups = "short")
+  public void should_report_no_host_when_unprepared_prepare_write_fails_synchronously()
+      throws Exception {
+    final Scassandra scassandra = TestUtils.createScassandraServer();
+    Cluster cluster = null;
+
+    try {
+      scassandra.start();
+      ScassandraCluster.primeSystemLocalRow(scassandra);
+      String query = "SELECT v FROM mock_table";
+      scassandra
+          .primingClient()
+          .prime(
+              preparedStatementBuilder()
+                  .withQuery(query)
+                  .withThen(then().withRows(row("result", "result1")))
+                  .build());
+
+      cluster =
+          Cluster.builder()
+              .addContactPoint(TestUtils.ipOfNode(1))
+              .withPort(scassandra.getBinaryPort())
+              .withPoolingOptions(
+                  new PoolingOptions()
+                      .setCoreConnectionsPerHost(HostDistance.LOCAL, 1)
+                      .setMaxConnectionsPerHost(HostDistance.LOCAL, 1)
+                      .setHeartbeatIntervalSeconds(0))
+              .build();
+      Session session = cluster.connect();
+      Host host = TestUtils.findHost(cluster, 1);
+      PreparedStatement prepared = session.prepare(query);
+
+      scassandra.primingClient().clearPreparedPrimes();
+      scassandra
+          .primingClient()
+          .prime(
+              preparedStatementBuilder()
+                  .withQuery(query)
+                  .withThen(
+                      then()
+                          .withResult(unprepared)
+                          .withFixedDelay(200L)
+                          .withConfig(
+                              new UnpreparedConfig(
+                                  prepared.getPreparedId().boundValuesMetadata.id.toString())))
+                  .build());
+
+      Connection connection = getSingleConnection(session);
+      ResultSetFuture future = session.executeAsync(prepared.bind());
+      waitForInFlight(connection);
+      markDefunctWithoutClosing(connection);
+
+      try {
+        future.getUninterruptibly(5, TimeUnit.SECONDS);
+        fail("expected a NoHostAvailableException");
+      } catch (NoHostAvailableException e) {
+        assertThat(e.getErrors()).containsKey(host.getEndPoint());
+        assertThat(e.getErrors().get(host.getEndPoint()))
+            .isInstanceOf(ConnectionException.class)
+            .hasMessageContaining("Write attempt on defunct connection");
+      }
+    } finally {
+      if (cluster != null) cluster.close();
+      scassandra.stop();
+    }
+  }
+
   @Test(groups = "short")
   public void should_handle_race_between_response_and_cancellation() {
     final Scassandra scassandra = TestUtils.createScassandraServer();
@@ -50,7 +163,7 @@ public void should_handle_race_between_response_and_cancellation() {
       scassandra
           .primingClient()
           .prime(
-              PrimingRequest.queryBuilder()
+              queryBuilder()
                   .withQuery("mock query")
                   .withThen(then().withRows(rows).withFixedDelay(10L))
                   .build());
@@ -100,4 +213,56 @@ private Connection getSingleConnection(Session session) {
     HostConnectionPool pool = ((SessionManager) session).pools.values().iterator().next();
     return pool.connections[0].get(0);
   }
+
+  private Cluster newCluster(ScassandraCluster scassandra) {
+    return Cluster.builder()
+        .addContactPoints(scassandra.address(1).getAddress())
+        .withPort(scassandra.getBinaryPort())
+        .withLoadBalancingPolicy(new SortingLoadBalancingPolicy())
+        .withPoolingOptions(
+            new PoolingOptions()
+                .setCoreConnectionsPerHost(HostDistance.LOCAL, 1)
+                .setMaxConnectionsPerHost(HostDistance.LOCAL, 1)
+                .setHeartbeatIntervalSeconds(0))
+        .build();
+  }
+
+  private void failBorrowingFrom(Session session, Host host, ConnectionException failure)
+      throws Exception {
+    SessionManager manager = (SessionManager) session;
+    HostConnectionPool pool = spy(manager.pools.get(host));
+    manager.pools.put(host, pool);
+    doReturn(Futures.<Connection>immediateFailedFuture(failure))
+        .when(pool)
+        .borrowConnection(
+            anyLong(),
+            (TimeUnit) any(),
+            anyInt(),
+            (Token.Factory) any(),
+            (ByteBuffer) any(),
+            (String) any(),
+            (String) any());
+  }
+
+  private void waitForInFlight(Connection connection) throws InterruptedException {
+    long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+    boolean observedInFlight = false;
+    while (System.nanoTime() < deadline) {
+      if (connection.inFlight.get() > 0) {
+        observedInFlight = true;
+        break;
+      }
+      Thread.sleep(1);
+    }
+    assertThat(observedInFlight).isTrue();
+  }
+
+  private void markDefunctWithoutClosing(Connection connection) {
+    // Keep the pending EXECUTE alive; only the following PREPARE write should see defunct.
+    connection.markDefunctForTest();
+  }
+
+  private static List<Map<String, ?>> row(String key, String value) {
+    return Collections.<Map<String, ?>>singletonList(ImmutableMap.of(key, value));
+  }
 }