Request Coalescer (Waiting room pattern)

RequestCoalescer.java

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * Coalesces concurrent requests for the same key into a single execution.
 *
 * <p>
 * <b>Contract for callers:</b> The {@code task} supplied by follower threads is
 * silently dropped — only the current leader's task runs. Do NOT rely on side
 * effects
 * (metrics, audit logs, etc.) inside {@code task} for every caller.
 *
 * <p>
 * <b>Timeout semantics:</b> {@code timeoutMillis} controls how long a follower
 * waits
 * before attempting to become the new leader. The original leader runs with no
 * timeout.
 * Under sustained slowness a new leader is elected every {@code timeoutMillis},
 * so the
 * task may execute concurrently across multiple leaders. Ensure {@code task} is
 * safe
 * to run concurrently for the same key.
 *
 * <p>
 * <b>Max re-elections:</b> At most {@code maxReelections} leadership handoffs
 * are
 * attempted before giving up with a {@link RuntimeException}.
 */
class RequestCoalescer<K, V> {

    private static final int DEFAULT_MAX_REELECTIONS = 3;

    private final ConcurrentHashMap<K, CompletableFuture<V>> inProgress = new ConcurrentHashMap<>();
    private final long timeoutMillis;
    private final int maxReelections;

    public RequestCoalescer(long timeoutMillis) {
        this(timeoutMillis, DEFAULT_MAX_REELECTIONS);
    }

    public RequestCoalescer(long timeoutMillis, int maxReelections) {
        this.timeoutMillis = timeoutMillis;
        this.maxReelections = maxReelections;
    }

    public V getCoalesced(K key, Callable<V> task) throws Exception {

        CompletableFuture<V> newFuture = new CompletableFuture<>();
        CompletableFuture<V> existing = inProgress.putIfAbsent(key, newFuture);

        if (existing == null) {
            return runAsLeader(key, newFuture, task);
        } else {
            return waitFor(existing, key, task);
        }
    }

    private V runAsLeader(K key, CompletableFuture<V> myFuture, Callable<V> task) throws Exception {
        System.out.println("hit run(" + myFuture.hashCode() + ") As Leader by thread " + Thread.currentThread().getName());
        try {
            V result = task.call();
            myFuture.complete(result);
            return result;
        } catch (Exception e) {
            myFuture.completeExceptionally(e);
            throw e;
        } finally {
            inProgress.remove(key, myFuture);
        }
    }

    /**
     * Waits for {@code future} to complete, re-electing as leader on timeout.
     * Uses an iterative loop instead of recursion to avoid stack overflow under
     * sustained timeouts with high concurrency.
     */
    private V waitFor(CompletableFuture<V> initialFuture, K key, Callable<V> task) throws Exception {
        CompletableFuture<V> current = initialFuture;
        int reelections = 0;

        while (true) {
            System.out.println("hit waitFor(" + current.hashCode() + ") by thread " + Thread.currentThread().getName());
            try {
                return current.get(timeoutMillis, TimeUnit.MILLISECONDS);
            } catch (TimeoutException te) {
                System.err.println("Timeout in thread :" + Thread.currentThread().getName());
                if (reelections >= maxReelections) {
                    throw new RuntimeException("Timed out waiting for key=" + key + " after "
                            + reelections + " re-elections", te);
                }
                // Attempt to become the new leader via atomic CAS.
                CompletableFuture<V> newFuture = new CompletableFuture<>();
                if (inProgress.replace(key, current, newFuture)) {
                    /*
                     * Won the election — run as new leader.
                     * Other followers waiting on newFuture get the result automatically.
                     */
                    return runAsLeader(key, newFuture, task);
                }

                /*
                 * Lost the election — another follower became leader.
                 * Wait on whatever is current; if the key is gone the new leader already
                 * finished (result not yet in caller's cache), so retry getCoalesced() from
                 * scratch.
                 * NOTE: getCoalesced() re-runs task only if no leader is active — task must be
                 * idempotent.
                 */
                CompletableFuture<V> next = inProgress.get(key);
                if (next == null) {
                    return getCoalesced(key, task);
                }
                current = next;
                reelections++;

            } catch (ExecutionException ee) {
                Throwable cause = ee.getCause();
                if (cause instanceof Exception) {
                    throw (Exception) cause;
                }
                throw new RuntimeException("Computation failed for key=" + key, cause);

            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
                throw new RuntimeException("Thread interrupted while waiting for key=" + key, ie);
            }
        }
    }
}

class UserService {

    private final RequestCoalescer<String, String> coalescer = new RequestCoalescer<>(500); // 500 ms timeout
    private final Map<String, String> cache = new ConcurrentHashMap<>();

    public static void main(String[] args) throws Exception {
        int numThreads = 5;
        String key = "user123";
        CountDownLatch startGate = new CountDownLatch(1); // ensures all threads start together
        CountDownLatch finishGate = new CountDownLatch(numThreads);
        ConcurrentLinkedQueue<String> results = new ConcurrentLinkedQueue<>();

        System.out.printf("Launching %d concurrent requests for '%s'...\n", numThreads, key);
        UserService userService = new UserService();
        for (int i = 0; i < numThreads; i++) {
            new Thread(() -> {
                try {
                    startGate.await(); // all threads wait here until released
                    String result = userService.getUserProfile(key);
                    results.add(String.format("    Thread-%s got: %s", Thread.currentThread().getName(), result));
                } catch (Exception e) {
                    results.add(String.format("    Thread-%s FAILED: %s", Thread.currentThread().getName(), e.getMessage()));
                } finally {
                    finishGate.countDown();
                }
            }, "caller-" + i).start();
        }

        // Release all threads simultaneously
        startGate.countDown();
        finishGate.await(10, TimeUnit.SECONDS);
        System.out.println("\nResults:");
        results.stream().sorted().forEach(System.out::println);
    }

    public String getUserProfile(String userId) {
        // Step 1: Cache lookup
        String cached = cache.get(userId);
        if (cached != null) {
            return cached;
        }
        try {
            // Step 2: Coalesced fetch
            String value = coalescer.getCoalesced(userId, () -> fetchFromDB(userId));

            // Step 3: Populate cache — putIfAbsent avoids redundant writes from
            // multiple followers returning concurrently for the same key.
            cache.putIfAbsent(userId, value);
            return value;

        } catch (Exception e) {
            // Step 4: Fallback strategy
            return handleFailure(userId, e);
        }
    }

    volatile boolean flag = true;

    private String fetchFromDB(String userId) throws InterruptedException {
        // Simulate timeout only for 1st time
        if (flag) {
            flag = false;
            Thread.sleep(1000);
        } else {
            Thread.sleep(200);
        }
        return "UserProfile-" + userId;
    }

    private String handleFailure(String userId, Exception e) {
        System.err.println("Error fetching user=" + userId + ", reason=" + e.getMessage());
        return "DEFAULT_PROFILE";
    }
}

/*
Output
------
Launching 5 concurrent requests for 'user123'...

hit waitFor(2123916292) by thread caller-3
hit run(2123916292) As Leader by thread caller-4
hit waitFor(2123916292) by thread caller-2
hit waitFor(2123916292) by thread caller-0
hit waitFor(2123916292) by thread caller-1
Timeout in thread :caller-0
Timeout in thread :caller-3
Timeout in thread :caller-2
Timeout in thread :caller-1
hit run(722320618) As Leader by thread caller-0
hit waitFor(722320618) by thread caller-3
hit waitFor(722320618) by thread caller-2
hit waitFor(722320618) by thread caller-1

Results:
    Thread-caller-0 got: UserProfile-user123
    Thread-caller-1 got: UserProfile-user123
    Thread-caller-2 got: UserProfile-user123
    Thread-caller-3 got: UserProfile-user123
    Thread-caller-4 got: UserProfile-user123
*/