How to prevent CompletableFuture#whenComplete execution in context thread

Question

I have the following code:

ConcurrentHashMap taskMap= new ConcurrentHashMap();
....
taskMap.compute(key, (k, queue) -> {
        CompletableFuture<Void> future = (queue == null)
                ? CompletableFuture.runAsync(myTask, poolExecutor)
                : queue.whenCompleteAsync((r, e) -> myTask.run(), poolExecutor);
        //to prevent OutOfMemoryError in case if we will have too much keys
        future.whenComplete((r, e) -> taskMap.remove(key, future));            
        return future;
    });

The issue of this code that in case of future already completed whenComplete function argument invokes in the same thread as compute invokes. In the body of this method we remove entry from map. But compute method documentation forbid this and application freezes.

How can I fix this issue?

Answer 1

The most obvious solution is to use whenCompleteAsync instead of whenComplete , as the former guarantees to execute the action using the supplied Executor rather than the calling thread. Which can be demonstrated with

Executor ex = r -> { System.out.println("job scheduled"); new Thread(r).start(); };
for(int run = 0; run<2; run++) {
    boolean completed = run==0;
    System.out.println("*** "+(completed? "with already completed": "with async"));
    CompletableFuture<String> source = completed?
        CompletableFuture.completedFuture("created   in "+Thread.currentThread()):
        CompletableFuture.supplyAsync(() -> {
            LockSupport.parkNanos(TimeUnit.SECONDS.toNanos(1));
            return "created   in "+Thread.currentThread();
        }, ex);

    source.thenApplyAsync(s -> s+"\nprocessed in "+Thread.currentThread(), ex)
          .whenCompleteAsync((s,t) -> {
                if(t!=null) t.printStackTrace(); else System.out.println(s);
                System.out.println("consumed  in "+Thread.currentThread());
            }, ex)
          .join();
}

which will print something like

*** with already completed
job scheduled
job scheduled
created   in Thread[main,5,main]
processed in Thread[Thread-0,5,main]
consumed  in Thread[Thread-1,5,main]
*** with async
job scheduled
job scheduled
job scheduled
created   in Thread[Thread-2,5,main]
processed in Thread[Thread-3,5,main]
consumed  in Thread[Thread-4,5,main]

So you could just use

taskMap.compute(key, (k, queue) -> {
        CompletableFuture<Void> future = (queue == null)
                ? CompletableFuture.runAsync(myTask, poolExecutor)
                : queue.whenCompleteAsync((r, e) -> myTask.run(), poolExecutor);
        //to prevent OutOfMemoryError in case if we will have too much keys
        future.whenCompleteAsync((r, e) -> taskMap.remove(key, future), poolExecutor);
        return future;
    });

If early completion has a significant likelihood, you could reduce the overhead using

taskMap.compute(key, (k, queue) -> {
        CompletableFuture<Void> future = (queue == null)
                ? CompletableFuture.runAsync(myTask, poolExecutor)
                : queue.whenCompleteAsync((r, e) -> myTask.run(), poolExecutor);
        //to prevent OutOfMemoryError in case if we will have too much keys
        if(future.isDone()) future = null;
        else future.whenCompleteAsync((r, e) -> taskMap.remove(key, future), poolExecutor);
        return future;
    });

Maybe, you didn't come to this obvious solution, because you don't like that the dependent action will always be scheduled as a new task to the pool, even if the completion happens in a different task already. You could solve this with a specialized executor which will only reschedule the task when necessary:

Executor inPlace = Runnable::run;
Thread forbidden = Thread.currentThread();
Executor forceBackground
       = r -> (Thread.currentThread()==forbidden? poolExecutor: inPlace).execute(r);

…

future.whenCompleteAsync((r, e) -> taskMap.remove(key, future), forceBackground);

---

But you may rethink whether this complicated per-mapping cleanup logic is really desired. It's not only complicated but also may create a notable overhead, potentially scheduling lots of cleanup actions that are not really required when they are already outdated when executed.

It might be much simpler and even more efficient to execute

taskMap.values().removeIf(CompletableFuture::isDone);

from time to time to cleanup the entire map.