Add API to get and cancel timed out Futures (Azure#36773)

Add API to get and cancel timed out Futures

Author: alzimmermsft

eng/versioning/version_client.txt

@@ -441,6 +441,7 @@ com.azure.tools:azure-sdk-build-tool;1.0.0;1.1.0-beta.1
 # note: The unreleased dependencies will not be manipulated with the automatic PR creation code.
 # In the pom, the version update tag after the version should name the unreleased package and the dependency version:
 # <!-- {x-version-update;unreleased_com.azure:azure-core;dependency} -->
+unreleased_com.azure:azure-core;1.44.0-beta.1
 
 # Released Beta dependencies: Copy the entry from above, prepend "beta_", remove the current
 # version and set the version to the released beta. Released beta dependencies are only valid

sdk/core/azure-core/src/main/java/com/azure/core/util/CoreUtils.java

@@ -23,7 +23,12 @@
 import java.util.Objects;
 import java.util.Properties;
 import java.util.UUID;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -531,4 +536,29 @@ static UUID randomUuid(long msb, long lsb) {
         // For environments using Reactor's BlockHound this will raise an exception if called in non-blocking threads.
         return new UUID(msb, lsb);
     }
+
+    /**
+     * Calls {@link Future#get(long, TimeUnit)} and returns the value if the {@code future} completes before the timeout
+     * is triggered. If the timeout is triggered, the {@code future} is {@link Future#cancel(boolean) cancelled}
+     * interrupting the execution of the task that the {@link Future} represented.
+     *
+     * @param <T> The type of value returned by the {@code future}.
+     * @param future The {@link Future} to get the value from.
+     * @param timeout The timeout value.
+     * @param unit The {@link TimeUnit} of the timeout value.
+     * @return The value from the {@code future}.
+     * @throws CancellationException If the computation was cancelled.
+     * @throws ExecutionException If the computation threw an exception.
+     * @throws InterruptedException If the current thread was interrupted while waiting.
+     * @throws TimeoutException If the wait timed out.
+     */
+    public static <T> T getFutureWithCancellation(Future<T> future, long timeout, TimeUnit unit)
+        throws InterruptedException, ExecutionException, TimeoutException {
+        try {
+            return future.get(timeout, unit);
+        } catch (TimeoutException e) {
+            future.cancel(true);
+            throw e;
+        }
+    }
 }

sdk/core/azure-core/src/test/java/com/azure/core/util/CoreUtilsTests.java

@@ -6,7 +6,9 @@
 import com.azure.core.http.HttpHeaders;
 import com.azure.core.http.policy.HttpLogOptions;
 import com.azure.core.util.logging.ClientLogger;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
@@ -25,7 +27,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
@@ -36,6 +45,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 public class CoreUtilsTests {
     private static final byte[] BYTES = "Hello world!".getBytes(StandardCharsets.UTF_8);
@@ -49,6 +59,18 @@ public class CoreUtilsTests {
     private static final String TIMEOUT_PROPERTY_NAME = "TIMEOUT_PROPERTY_NAME";
     private static final ConfigurationSource EMPTY_SOURCE = new TestConfigurationSource();
 
+    private static ExecutorService executorService;
+
+    @BeforeAll
+    public static void setupClass() {
+        executorService = Executors.newCachedThreadPool();
+    }
+
+    @AfterAll
+    public static void teardownClass() {
+        executorService.shutdownNow();
+    }
+
     @Test
     public void findFirstOfTypeEmptyArgs() {
         assertNull(CoreUtils.findFirstOfType(null, Integer.class));
@@ -484,7 +506,7 @@ public void randomUuidIsCorrectlyType4() {
         bytes[6] &= 0x0f;  /* clear version        */
         bytes[6] |= 0x40;  /* set to version 4     */
         bytes[8] &= 0x3f;  /* clear variant        */
-        bytes[8] |= 0x80;  /* set to IETF variant  */
+        bytes[8] |= (byte) 0x80;  /* set to IETF variant  */
         long msbForJava = 0;
         long lsbForJava = 0;
         for (int i = 0; i < 8; i++) {
@@ -496,4 +518,48 @@ public void randomUuidIsCorrectlyType4() {
 
         assertEquals(new UUID(msbForJava, lsbForJava), CoreUtils.randomUuid(msb, lsb));
     }
+
+    @Test
+    public void futureCompletesBeforeTimeout() {
+        try {
+            AtomicBoolean completed = new AtomicBoolean(false);
+            Future<?> future = executorService.submit(() -> {
+                Thread.sleep(10);
+                completed.set(true);
+                return null;
+            });
+
+            future.get(1000, TimeUnit.MILLISECONDS);
+
+            assertTrue(completed.get());
+        } catch (InterruptedException | ExecutionException | TimeoutException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    public void futureTimesOutAndIsCancelled() {
+        try {
+            AtomicBoolean completed = new AtomicBoolean(false);
+            Future<?> future = executorService.submit(() -> {
+                Thread.sleep(1000);
+                completed.set(true);
+                return null;
+            });
+
+            try {
+                CoreUtils.getFutureWithCancellation(future, 100, TimeUnit.MILLISECONDS);
+                fail("Expected future to timout and be cancelled.");
+            } catch (TimeoutException e) {
+                // Expected.
+            }
+
+            // Give time for the future to complete if cancellation didn't work.
+            Thread.sleep(1000);
+
+            assertFalse(completed.get());
+        } catch (InterruptedException | ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

sdk/tables/azure-data-tables/pom.xml

@@ -46,7 +46,7 @@ Licensed under the MIT License.
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core</artifactId>
-      <version>1.43.0</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
+      <version>1.44.0-beta.1</version> <!-- {x-version-update;unreleased_com.azure:azure-core;dependency} -->
     </dependency>
     <dependency>
       <groupId>com.azure</groupId>