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 useTimeoutException() {
+        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.useTimeoutException());
 
         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() {

src/test/java/io/nats/client/impl/RequestTests.java

@@ -680,7 +680,7 @@ public void testThrowsEmptySubject() {
 
     @Test
     public void testNatsRequestCompletableFuture() throws InterruptedException {
-        NatsRequestCompletableFuture f = new NatsRequestCompletableFuture(CancelAction.CANCEL, Duration.ofHours(-1));
+        NatsRequestCompletableFuture f = new NatsRequestCompletableFuture(CancelAction.CANCEL, Duration.ofHours(-1), true);
         assertEquals(CancelAction.CANCEL, f.getCancelAction());
         assertTrue(f.hasExceededTimeout());
         assertFalse(f.wasCancelledClosing());
@@ -690,14 +690,14 @@ public void testNatsRequestCompletableFuture() throws InterruptedException {
         assertTrue(f.wasCancelledClosing());
         assertTrue(f.wasCancelledTimedOut());
 
-        f = new NatsRequestCompletableFuture(CancelAction.COMPLETE, Duration.ofHours(-1));
+        f = new NatsRequestCompletableFuture(CancelAction.COMPLETE, Duration.ofHours(-1), true);
         assertEquals(CancelAction.COMPLETE, f.getCancelAction());
 
-        f = new NatsRequestCompletableFuture(CancelAction.REPORT, Duration.ofHours(-1));
+        f = new NatsRequestCompletableFuture(CancelAction.REPORT, Duration.ofHours(-1), true);
         assertEquals(CancelAction.REPORT, f.getCancelAction());
 
         // coverage for null timeout
-        f = new NatsRequestCompletableFuture(CancelAction.CANCEL, null);
+        f = new NatsRequestCompletableFuture(CancelAction.CANCEL, null, true);
         Thread.sleep(Options.DEFAULT_REQUEST_CLEANUP_INTERVAL.toMillis() + 100);
         assertTrue(f.hasExceededTimeout());
     }