core(utils): Added utility class for methods related to CompletableFuture

Author: shortthirdman

src/main/java/com/shortthirdman/primekit/essentials/common/util/CompletableFutureUtils.java

@@ -0,0 +1,103 @@
+package com.shortthirdman.primekit.essentials.common.util;
+
+import org.jetbrains.annotations.NotNull;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.BiConsumer;
+
+/**
+ * A utility class to add JDK 9+ timeout functionality to CompletableFuture in JDK.
+ * Inspired by JDK 9's CompletableFuture implementation.
+ *
+ * @author ShortThirdMan
+ */
+public final class CompletableFutureUtils {
+
+    private CompletableFutureUtils() {}
+
+    /**
+     * If not already completed, causes this CompletableFuture to be
+     * completed exceptionally with a {@link TimeoutException} after the
+     * given timeout.
+     *
+     * @param future  the CompletableFuture to apply the timeout to
+     * @param timeout how long to wait before completing exceptionally
+     * @param unit    the time unit of the timeout argument
+     * @return the original CompletableFuture
+     */
+    public static <T> CompletableFuture<T> orTimeout(CompletableFuture<T> future, long timeout, TimeUnit unit) {
+        if (unit == null) {
+            throw new NullPointerException("Time unit cannot be null");
+        }
+        if (future == null) {
+            throw new NullPointerException("CompletableFuture cannot be null");
+        }
+        // If the future is already done, just return it.
+        if (future.isDone()) {
+            return future;
+        }
+        // Schedule a task to complete the future exceptionally after the timeout.
+        // The Canceller will cancel this scheduled task if the future completes normally.
+        return future.whenComplete(new Canceller(Delayer.delay(new Timeout(future), timeout, unit)));
+    }
+
+    /**
+     * Inner class to handle the actual timeout logic
+     */
+    static final class Timeout implements Runnable {
+        final CompletableFuture<?> future;
+        Timeout(CompletableFuture<?> future) {
+            this.future = future;
+        }
+        public void run() {
+            if (future != null && !future.isDone()) {
+                future.completeExceptionally(new TimeoutException());
+            }
+        }
+    }
+
+    /**
+     * Inner class to handle cancellation of the timeout task
+     */
+    static final class Canceller implements BiConsumer<Object, Throwable> {
+        final Future<?> future;
+        Canceller(Future<?> future) {
+            this.future = future;
+        }
+        public void accept(Object ignore, Throwable ex) {
+            // If the original future completed (ex is null) and the timeout task
+            // hasn't run yet, cancel the timeout task.
+            if (ex == null && future != null && !future.isDone()) {
+                future.cancel(false);
+            }
+        }
+    }
+
+    /**
+     * Singleton delayed scheduler
+     */
+    static final class Delayer {
+        static ScheduledFuture<?> delay(Runnable command, long delay, TimeUnit unit) {
+            return delayer.schedule(command, delay, unit);
+        }
+        private static final ScheduledThreadPoolExecutor delayer;
+        static {
+            delayer = new ScheduledThreadPoolExecutor(1, new DaemonThreadFactory());
+            delayer.setRemoveOnCancelPolicy(true);
+        }
+        static final class DaemonThreadFactory implements ThreadFactory {
+            public Thread newThread(@NotNull Runnable r) {
+                Thread t = new Thread(r);
+                t.setDaemon(true);
+                t.setName("CompletableFutureUtilsDelayScheduler");
+                return t;
+            }
+        }
+    }
+}

src/test/java/com/shortthirdman/primekit/essentials/common/util/CompletableFutureUtilsTest.java

@@ -0,0 +1,63 @@
+package com.shortthirdman.primekit.essentials.common.util;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class CompletableFutureUtilsTest {
+
+    @Test
+    void completesNormallyBeforeTimeout_cancelsTimeoutTask() {
+        CompletableFuture<String> f = new CompletableFuture<>();
+        CompletableFutureUtils.orTimeout(f, 200, TimeUnit.MILLISECONDS);
+
+        // Complete before 200ms
+        f.complete("ok");
+
+        assertTrue(f.isDone());
+        assertFalse(f.isCompletedExceptionally());
+        assertEquals("ok", f.join());
+
+        // Sleep past the timeout to ensure the scheduled timeout would have fired if not cancelled
+        try { Thread.sleep(300); } catch (InterruptedException ignored) {}
+        // If the timeout wasn't cancelled, the future would be completed exceptionally already, which is not the case.
+        assertFalse(f.isCompletedExceptionally());
+    }
+
+    @Test
+    void timesOutWhenNotCompletedInTime_completesExceptionallyWithTimeoutException() {
+        CompletableFuture<String> f = new CompletableFuture<>();
+        CompletableFutureUtils.orTimeout(f, 50, TimeUnit.MILLISECONDS);
+
+        // Do not complete the future; wait to let timeout trigger
+        try { Thread.sleep(120); } catch (InterruptedException ignored) {}
+
+        assertTrue(f.isCompletedExceptionally());
+        CompletionException ex = assertThrows(CompletionException.class, f::join);
+        assertInstanceOf(TimeoutException.class, ex.getCause());
+    }
+
+    @Test
+    void handlesNullArgumentsAndAlreadyCompletedFuture() {
+        // Null unit
+        CompletableFuture<String> f1 = new CompletableFuture<>();
+        assertThrows(NullPointerException.class, () -> CompletableFutureUtils.orTimeout(f1, 10, null));
+
+        // Null future
+        assertThrows(NullPointerException.class, () -> CompletableFutureUtils.orTimeout(null, 10, TimeUnit.MILLISECONDS));
+
+        // Already completed future should be returned as is and not be modified
+        CompletableFuture<Integer> done = CompletableFuture.completedFuture(42);
+        CompletableFuture<Integer> returned = CompletableFutureUtils.orTimeout(done, 10, TimeUnit.MILLISECONDS);
+        assertSame(done, returned);
+        assertEquals(42, returned.join());
+        // Sleep beyond timeout to confirm nothing changes
+        try { Thread.sleep(50); } catch (InterruptedException ignored) {}
+        assertFalse(returned.isCompletedExceptionally());
+    }
+}