Make asynchronous queries synchronous

Question

I have an underlying asynchronous out-of-order query/response system that I want to make synchronous. Queries and responses can be mapped by marking the queries with unique IDs that in turn will accompany the corresponding responses.

My attempt at making it synchronous uses two ConcurrentHashMaps: one that maps from IDs to results, and one that maps from the same IDs to CountDownLatches. The code goes like this when executing a query:

public Result execute(Query query) throws InterruptedException {
    int id = atomicInteger.incrementAndGet();
    CountDownLatch latch = new CountDownLatch(1);
    latchMap.put(id, latch);
    query.executeAsyncWithId(id); // Probably returns before result is ready
    try {
        latch.await(); // Blocks until result is ready
        return resultMap.remove(id);
    } catch (InterruptedException e) {
        latchMap.remove(id); // We are not waiting anymore
        throw e;
    }
}

And the code for handling an incoming result:

public void handleResult(Result result) {
    int id = result.getId();
    CountDownLatch latch = latchMap.remove(id);
    if (latch == null) {
        return; // Nobody wants result
    }
    resultMap.put(id, result);
    latch.countDown();
}

This method is called from a thread that reads all the incoming results from the underlying system (there is only one such reader thread).

First of all I am uncertain about the thread-safety, but it also seems unnecessary to use two HashMaps for this (particularly because IDs are never reused). Any ideas for improvements?

Answer 1

A new attempt inspired by this answer to a similar question:

public class ResultFuture {

    private volatile Result result = null;
    private final CountDownLatch latch = new CountDownLatch(1);

    public Result get() throws InterruptedException {
        latch.await();
        return result;
    }

    public void set(Result result) {
        this.result = result;
        latch.countDown();
    }
}

Now I need only one HashMap of these ResultFutures:

public Result execute(Query query) throws InterruptedException {
    int id = atomicInteger.incrementAndGet();
    ResultFuture resultFuture = new ResultFuture();
    resultFutureMap.put(id, resultFuture);
    query.executeAsyncWithId(id); // Probably returns before result is ready
    try {
        return resultFuture.get(); // Blocks until result is ready
    } finally {
        resultFutureMap.remove(id);
    }
}

public void handleResult(Result result) {
    int id = result.getId();
    ResultFuture resultFuture = resultFutureMap.get(id);
    if (resultFuture == null) {
        return; // Nobody wants result
    }
    resultFuture.set(result);
}