fix: enrich timeout diagnostics

Author: dkropachev

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

@@ -411,6 +411,9 @@ public Message.Request request() {
             @Override
             public void register(RequestHandler handler) {}
 
+            @Override
+            public void registerReadTimeoutMillis(long readTimeoutMillis) {}
+
             @Override
             public void onSet(
                 Connection connection,

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

@@ -962,6 +962,9 @@ ResponseHandler write(
       throws ConnectionException, BusyConnectionException {
 
     ResponseHandler handler = new ResponseHandler(this, statementReadTimeoutMillis, callback);
+    if (callback instanceof RequestHandler.Callback) {
+      ((RequestHandler.Callback) callback).registerReadTimeoutMillis(handler.readTimeoutMillis);
+    }
     dispatcher.add(handler);
 
     Message.Request request = callback.request().setStreamId(handler.streamId);
@@ -1062,6 +1065,97 @@ public ProtocolFeatureStore getProtocolFeatureStore() {
     return protocolFeatureStore;
   }
 
+  OperationTimedOutException newTimeoutException(long elapsedTimeoutNanos, int retryCount) {
+    return newTimeoutException(
+        null,
+        factory.getReadTimeoutMillis(),
+        elapsedTimeoutNanos,
+        retryCount,
+        OperationTimedOutException.UNAVAILABLE);
+  }
+
+  OperationTimedOutException newTimeoutException(
+      long configuredTimeoutMs,
+      long elapsedTimeoutNanos,
+      int retryCount,
+      int speculativeExecutionIndex) {
+    return newTimeoutException(
+        null, configuredTimeoutMs, elapsedTimeoutNanos, retryCount, speculativeExecutionIndex);
+  }
+
+  OperationTimedOutException newTimeoutException(
+      String message,
+      long configuredTimeoutMs,
+      long elapsedTimeoutNanos,
+      int retryCount,
+      int speculativeExecutionIndex) {
+    HostConnectionPool pool = ownerPool();
+    ShardingInfo.ConnectionShardingInfo connectionShardingInfo =
+        protocolFeatureStore == null ? null : protocolFeatureStore.getConnectionShardingInfo();
+    ShardingInfo hostShardingInfo =
+        connectionShardingInfo != null
+            ? connectionShardingInfo.shardingInfo
+            : pool != null ? pool.host.getShardingInfo() : null;
+    int connectionShardId =
+        connectionShardingInfo != null ? shardId() : OperationTimedOutException.UNAVAILABLE;
+    int hostShardsCount =
+        hostShardingInfo != null
+            ? hostShardingInfo.getShardsCount()
+            : OperationTimedOutException.UNAVAILABLE;
+    int poolPendingBorrows =
+        pool != null ? pool.pendingBorrowCount.get() : OperationTimedOutException.UNAVAILABLE;
+    int poolPendingBorrowsForShard =
+        pool != null && connectionShardId != OperationTimedOutException.UNAVAILABLE
+            ? pool.pendingBorrowCountForShard(connectionShardId)
+            : OperationTimedOutException.UNAVAILABLE;
+    int poolTotalInFlight =
+        pool != null ? pool.totalInFlight.get() : OperationTimedOutException.UNAVAILABLE;
+    int poolOpenConnectionsForShard =
+        pool != null && connectionShardId != OperationTimedOutException.UNAVAILABLE
+            ? pool.openConnectionCountForShard(connectionShardId)
+            : OperationTimedOutException.UNAVAILABLE;
+    int poolMaxConnectionsPerShard =
+        pool != null ? pool.maxConnectionsPerShard() : OperationTimedOutException.UNAVAILABLE;
+    long elapsedTimeoutMs = TimeUnit.NANOSECONDS.toMillis(elapsedTimeoutNanos);
+
+    if (message == null) {
+      return new OperationTimedOutException(
+          endPoint,
+          configuredTimeoutMs,
+          elapsedTimeoutMs,
+          retryCount,
+          speculativeExecutionIndex,
+          inFlight.get(),
+          connectionShardId,
+          hostShardsCount,
+          poolPendingBorrows,
+          poolPendingBorrowsForShard,
+          poolTotalInFlight,
+          poolOpenConnectionsForShard,
+          poolMaxConnectionsPerShard);
+    }
+    return new OperationTimedOutException(
+        endPoint,
+        message,
+        configuredTimeoutMs,
+        elapsedTimeoutMs,
+        retryCount,
+        speculativeExecutionIndex,
+        inFlight.get(),
+        connectionShardId,
+        hostShardsCount,
+        poolPendingBorrows,
+        poolPendingBorrowsForShard,
+        poolTotalInFlight,
+        poolOpenConnectionsForShard,
+        poolMaxConnectionsPerShard);
+  }
+
+  private HostConnectionPool ownerPool() {
+    Owner owner = ownerRef.get();
+    return owner instanceof HostConnectionPool ? (HostConnectionPool) owner : null;
+  }
+
   /**
    * If the connection is part of a pool, return it to the pool. The connection should generally not
    * be reused after that.
@@ -1814,6 +1908,11 @@ public void register(RequestHandler handler) {
       // noop, we don't care about the handler here so far
     }
 
+    @Override
+    public void registerReadTimeoutMillis(long readTimeoutMillis) {
+      // noop, direct connection writes keep the legacy timeout exception shape
+    }
+
     @Override
     public Message.Request request() {
       return request;

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

@@ -41,6 +41,7 @@ class DefaultResultSetFuture extends AbstractFuture<ResultSet>
   private final ProtocolVersion protocolVersion;
   private final Message.Request request;
   private volatile RequestHandler handler;
+  private volatile long readTimeoutMillis = OperationTimedOutException.UNAVAILABLE;
 
   DefaultResultSetFuture(
       SessionManager session, ProtocolVersion protocolVersion, Message.Request request) {
@@ -54,6 +55,11 @@ public void register(RequestHandler handler) {
     this.handler = handler;
   }
 
+  @Override
+  public void registerReadTimeoutMillis(long readTimeoutMillis) {
+    this.readTimeoutMillis = readTimeoutMillis;
+  }
+
   @Override
   public Message.Request request() {
     return request;
@@ -293,7 +299,9 @@ public boolean onTimeout(Connection connection, long latency, int retryCount) {
     // RequestHandler).
     // So just set an exception for the final result, which should be handled correctly by said
     // internal call.
-    setException(new OperationTimedOutException(connection.endPoint));
+    setException(
+        connection.newTimeoutException(
+            readTimeoutMillis, latency, retryCount, OperationTimedOutException.UNAVAILABLE));
     return true;
   }
 

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

@@ -29,6 +29,7 @@
 import com.datastax.driver.core.exceptions.AuthenticationException;
 import com.datastax.driver.core.exceptions.BusyPoolException;
 import com.datastax.driver.core.exceptions.ConnectionException;
+import com.datastax.driver.core.exceptions.OperationTimedOutException;
 import com.datastax.driver.core.exceptions.UnsupportedProtocolVersionException;
 import com.datastax.driver.core.utils.MoreFutures;
 import com.google.common.annotations.VisibleForTesting;
@@ -90,6 +91,7 @@ class HostConnectionPool implements Connection.Owner {
   @VisibleForTesting Set<Connection>[] trash;
 
   private Queue<PendingBorrow>[] pendingBorrows;
+  private AtomicInteger[] pendingBorrowCountByShard;
   final AtomicInteger pendingBorrowCount = new AtomicInteger();
 
   private AtomicInteger[] scheduledForCreation;
@@ -208,6 +210,29 @@ public void tempBlockAdvShardAwareness(long millis) {
         Math.max(System.currentTimeMillis() + millis, advShardAwarenessBlockedUntil);
   }
 
+  int pendingBorrowCountForShard(int shardId) {
+    if (pendingBorrowCountByShard == null
+        || shardId < 0
+        || shardId >= pendingBorrowCountByShard.length) {
+      return OperationTimedOutException.UNAVAILABLE;
+    }
+    return pendingBorrowCountByShard[shardId].get();
+  }
+
+  int openConnectionCountForShard(int shardId) {
+    if (connections == null || shardId < 0 || shardId >= connections.length) {
+      return OperationTimedOutException.UNAVAILABLE;
+    }
+    return connections[shardId].size();
+  }
+
+  int maxConnectionsPerShard() {
+    if (connections == null) {
+      return OperationTimedOutException.UNAVAILABLE;
+    }
+    return maxConnectionsPerShard;
+  }
+
   private final ConnectionTasksSharedState connectionTasksSharedState =
       new ConnectionTasksSharedState();
 
@@ -306,11 +331,13 @@ ListenableFuture<Void> initAsyncWithConnection(Connection reusedConnection) {
     open = new AtomicInteger[shardsCount];
     trash = new Set[shardsCount];
     pendingBorrows = new Queue[shardsCount];
+    pendingBorrowCountByShard = new AtomicInteger[shardsCount];
     for (int i = 0; i < shardsCount; ++i) {
       this.connections[i] = new CopyOnWriteArrayList<Connection>();
       scheduledForCreation[i] = new AtomicInteger();
       open[i] = new AtomicInteger();
       trash[i] = new CopyOnWriteArraySet<Connection>();
+      pendingBorrowCountByShard[i] = new AtomicInteger();
       pendingBorrows[i] = new ConcurrentLinkedQueue<PendingBorrow>();
     }
 
@@ -626,6 +653,7 @@ private ListenableFuture<Connection> enqueue(
     }
 
     PendingBorrow pendingBorrow = new PendingBorrow(timeout, unit, timeoutsExecutor);
+    pendingBorrowCountByShard[shardId].incrementAndGet();
     pendingBorrows[shardId].add(pendingBorrow);
 
     // If we raced with shutdown, make sure the future will be completed. This has no effect if it
@@ -690,6 +718,7 @@ private void dequeue(final Connection connection) {
         connection.inFlight.decrementAndGet();
       } else {
         pendingBorrowCount.decrementAndGet();
+        pendingBorrowCountByShard[connection.shardId()].decrementAndGet();
         // Ensure that the keyspace set on the connection is the one set on the pool state, in the
         // general case it will be.
         ListenableFuture<Connection> setKeyspaceFuture =

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

@@ -319,6 +319,8 @@ void onSet(
         long latency);
 
     void register(RequestHandler handler);
+
+    void registerReadTimeoutMillis(long readTimeoutMillis);
   }
 
   /**
@@ -943,11 +945,19 @@ public boolean onTimeout(Connection connection, long latency, int retryCount) {
                 queryStateRef.get());
             return false;
           }
+          long configuredTimeoutMs =
+              connectionHandler != null
+                  ? connectionHandler.readTimeoutMillis
+                  : OperationTimedOutException.UNAVAILABLE;
+          OperationTimedOutException timeoutException =
+              connection.newTimeoutException(
+                  "Timed out waiting for response to PREPARE message",
+                  configuredTimeoutMs,
+                  latency,
+                  retryCount,
+                  position);
           connection.release();
-          logError(
-              connection.endPoint,
-              new OperationTimedOutException(
-                  connection.endPoint, "Timed out waiting for response to PREPARE message"));
+          logError(connection.endPoint, timeoutException);
           retry(false, null);
           return true;
         }
@@ -1006,25 +1016,13 @@ public boolean onTimeout(Connection connection, long latency, int retryCount) {
       }
 
       Host queriedHost = current;
-
-      HostConnectionPool pool = queriedHost == null ? null : manager.pools.get(queriedHost);
       long configuredTimeoutMs =
           connectionHandler != null
               ? connectionHandler.readTimeoutMillis
               : OperationTimedOutException.UNAVAILABLE;
-      int connInFlight = connection.inFlight.get();
-      int poolPendingBorrows =
-          pool != null ? pool.pendingBorrowCount.get() : OperationTimedOutException.UNAVAILABLE;
-      int poolTotalInFlight =
-          pool != null ? pool.totalInFlight.get() : OperationTimedOutException.UNAVAILABLE;
 
       OperationTimedOutException timeoutException =
-          new OperationTimedOutException(
-              connection.endPoint,
-              configuredTimeoutMs,
-              connInFlight,
-              poolPendingBorrows,
-              poolTotalInFlight);
+          connection.newTimeoutException(configuredTimeoutMs, latency, retryCount, position);
 
       try {
         connection.release();