TimeoutException instead of CancellationException on timeout

src/main/java/io/nats/client/Options.java

@@ -444,6 +444,11 @@ public class Options {
      */
     @Deprecated
     public static final String PROP_UTF8_SUBJECTS = "allow.utf8.subjects";
+    /**
+     * Property used to throw {@link java.util.concurrent.TimeoutException} on timeout instead of {@link java.util.concurrent.CancellationException}.
+     * {@link Builder#useTimeoutException()}.
+     */
+    public static final String PROP_USE_TIMEOUT_EXCEPTION = PFX + "use.timeout.exception";
 
     // ----------------------------------------------------------------------------------------------------
     // PROTOCOL CONNECT OPTION CONSTANTS
@@ -571,6 +576,7 @@ public class Options {
     private final boolean discardMessagesWhenOutgoingQueueFull;
     private final boolean ignoreDiscoveredServers;
     private final boolean tlsFirst;
+    private final boolean useTimeoutException;
 
     private final AuthHandler authHandler;
     private final ReconnectDelayHandler reconnectDelayHandler;
@@ -679,6 +685,7 @@ public static class Builder {
         private boolean discardMessagesWhenOutgoingQueueFull = DEFAULT_DISCARD_MESSAGES_WHEN_OUTGOING_QUEUE_FULL;
         private boolean ignoreDiscoveredServers = false;
         private boolean tlsFirst = false;
+        private boolean useTimeoutException = false;
         private ServerPool serverPool = null;
         private DispatcherFactory dispatcherFactory = null;
 
@@ -804,6 +811,7 @@ public Builder properties(Properties props) {
 
             booleanProperty(props, PROP_IGNORE_DISCOVERED_SERVERS, b -> this.ignoreDiscoveredServers = b);
             booleanProperty(props, PROP_TLS_FIRST, b -> this.tlsFirst = b);
+            booleanProperty(props, PROP_USE_TIMEOUT_EXCEPTION, b -> this.useTimeoutException = b);
 
             classnameProperty(props, PROP_SERVERS_POOL_IMPLEMENTATION_CLASS, o -> this.serverPool = (ServerPool) o);
             classnameProperty(props, PROP_DISPATCHER_FACTORY_CLASS, o -> this.dispatcherFactory = (DispatcherFactory) o);
@@ -1489,6 +1497,15 @@ public Builder tlsFirst() {
             return this;
         }
 
+        /**
+         * Throw {@link java.util.concurrent.TimeoutException} on timeout instead of {@link java.util.concurrent.CancellationException}?
+         * @return the Builder for chaining
+         */
+        public Builder useTimeoutException() {
+            this.useTimeoutException = true;
+            return this;
+        }
+
         /**
          * Set the ServerPool implementation for connections to use instead of the default implementation
          * @param serverPool the implementation
@@ -1667,6 +1684,7 @@ public Builder(Options o) {
 
             this.ignoreDiscoveredServers = o.ignoreDiscoveredServers;
             this.tlsFirst = o.tlsFirst;
+            this.useTimeoutException = o.useTimeoutException;
 
             this.serverPool = o.serverPool;
             this.dispatcherFactory = o.dispatcherFactory;
@@ -1729,6 +1747,7 @@ private Options(Builder b) {
 
         this.ignoreDiscoveredServers = b.ignoreDiscoveredServers;
         this.tlsFirst = b.tlsFirst;
+        this.useTimeoutException = b.useTimeoutException;
 
         this.serverPool = b.serverPool;
         this.dispatcherFactory = b.dispatcherFactory;
@@ -2117,6 +2136,14 @@ public boolean isTlsFirst() {
         return tlsFirst;
     }
 
+    /**
+     * Get whether to throw {@link java.util.concurrent.TimeoutException} on timeout instead of {@link java.util.concurrent.CancellationException}.
+     * @return the flag
+     */
+    public boolean isUseTimeoutException() {
+        return useTimeoutException;
+    }
+
     /**
      * Get the ServerPool implementation. If null, a default implementation is used.
      * @return the ServerPool implementation

src/main/java/io/nats/client/impl/NatsConnection.java

@@ -1169,7 +1169,7 @@ CompletableFuture<Message> requestFutureInternal(String subject, Headers headers
         String responseToken = getResponseToken(responseInbox);
         NatsRequestCompletableFuture future =
             new NatsRequestCompletableFuture(cancelAction,
-                futureTimeout == null ? options.getRequestCleanupInterval() : futureTimeout);
+                futureTimeout == null ? options.getRequestCleanupInterval() : futureTimeout, options.isUseTimeoutException());
 
         if (!oldStyle) {
             responsesAwaiting.put(responseToken, future);

src/main/java/io/nats/client/support/NatsRequestCompletableFuture.java

@@ -6,6 +6,7 @@
 import java.time.Duration;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeoutException;
 
 /**
  * This is an internal class and is only public for access.
@@ -18,11 +19,13 @@ public enum CancelAction { CANCEL, REPORT, COMPLETE }
     private final long timeOutAfter;
     private boolean wasCancelledClosing;
     private boolean wasCancelledTimedOut;
+    private final boolean useTimeoutException;
 
-    public NatsRequestCompletableFuture(CancelAction cancelAction, Duration timeout) {
+    public NatsRequestCompletableFuture(CancelAction cancelAction, Duration timeout, boolean useTimeoutException) {
         this.cancelAction = cancelAction;
         timeOutAfter = System.currentTimeMillis() + 10 + (timeout == null ? DEFAULT_TIMEOUT : timeout.toMillis());
         // 10 extra millis allows for communication time, probably more than needed but...
+        this.useTimeoutException = useTimeoutException;
     }
 
     public void cancelClosing() {
@@ -32,7 +35,8 @@ public void cancelClosing() {
 
     public void cancelTimedOut() {
         wasCancelledTimedOut = true;
-        completeExceptionally(new CancellationException("Future cancelled, response not registered in time, likely due to server disconnect."));
+        final String message = "Future cancelled, response not registered in time, likely due to server disconnect.";
+        completeExceptionally(useTimeoutException ? new TimeoutException(message) : new CancellationException(message));
     }
 
     public CancelAction getCancelAction() {
@@ -50,4 +54,8 @@ public boolean wasCancelledClosing() {
     public boolean wasCancelledTimedOut() {
         return wasCancelledTimedOut;
     }
+
+    public boolean isUseTimeoutException() {
+        return useTimeoutException;
+    }
 }