How to cancel ShceduledFuture and wait for runnable to stop, if runnable is in progress at the moment of cancellation?

Question

When any command scheduled with fixed rate at any ScheduledExecutorService, it returns ScheduledFuture which can be cancelled as well. But "cancel" does not provide guarantee that command is not still executing after cancel returns, for example because command was already in the middle of execution when "cancell" was called.

For mostly use cases it is enough functionality. But I have deal with usecase when need to block current thread after cancel, if command already is in progress, and wait until command done. In other words thread which called cancel should not go forward if command still executing. Cancelling with mayInterruptIfRunning=true also is not suitable, because I do not want to broke current executions, I just need to wait for normal complete.

I did not found how to achieve this requirements via standard JDK classes. Question1: Was I wrong and this kind of functionality exists?

So I decided to implement it by itself: import java.util.concurrent.*;

public class GracefullyStoppingScheduledFutureDecorator implements ScheduledFuture {

/**
 * @return the scheduled future with method special implementation of "cancel" method, 
 * which in additional to standard implementation, 
 * provides strongly guarantee that command is not in the middle of progress when "cancel" returns  
 */
public static ScheduledFuture schedule(Runnable command, long initialDelay, long period, TimeUnit unit, ScheduledExecutorService scheduler) {
    CancellableCommand cancellableCommand = new CancellableCommand(command);
    ScheduledFuture future = scheduler.scheduleAtFixedRate(cancellableCommand, initialDelay, period, unit);
    return new GracefullyStoppingScheduledFutureDecorator(future, cancellableCommand);
}

private GracefullyStoppingScheduledFutureDecorator(ScheduledFuture targetFuture, CancellableCommand command) {
    this.targetFuture = targetFuture;
    this.runnable = command;
}

private final ScheduledFuture targetFuture;
private final CancellableCommand runnable;

@Override
public boolean cancel(boolean mayInterruptIfRunning) {
    runnable.cancel();
    return targetFuture.cancel(mayInterruptIfRunning);
}

@Override
public long getDelay(TimeUnit unit) {
    return targetFuture.getDelay(unit);
}

@Override
public int compareTo(Delayed o) {
    return targetFuture.compareTo(o);
}

@Override
public boolean isCancelled() {
    return targetFuture.isCancelled();
}

@Override
public boolean isDone() {
    return targetFuture.isDone();
}

@Override
public Object get() throws InterruptedException, ExecutionException {
    return targetFuture.get();
}

@Override
public Object get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
    return targetFuture.get(timeout, unit);
}

private static class CancellableCommand implements Runnable {

    private final Object monitor = new Object();
    private final Runnable target;
    private boolean cancelled = false;

    private CancellableCommand(Runnable target) {
        this.target = target;
    }

        public void cancel() {
            synchronized (monitor) {
                cancelled = true;
            }
        }

        @Override
        public void run() {
            synchronized (monitor) {
                if (!cancelled) {
                    target.run();
                }
            }
        }

    }

}

Question2: Could anybody find errors in the code above?

Answer 1

Question2: Could anybody find errors in the code above?

There is hypothetical deadlock which can be reproduced by following scenario:

1. Having thread T1 which holds monitor M1
2. Scheduled task is executing(holds its monitor M2) on thread T2 and wants to enter to M1, so T2 need to wait until T1 exits monitor M1.
3. T1 decided to cancel task, but because its monitor M2 is locked by task itself we have the deadlock.

Most likely scenario abovr is unreal, but to protect from all possible cases, I decided to rewrite code in lock-free manner:

public class GracefullyStoppingScheduledFuture {

/**
 * @return the scheduled future with method special implementation of "cancel" method,
 * which in additional to standard implementation,
 * provides strongly guarantee that command is not in the middle of progress when "cancel" returns
 */
public static GracefullyStoppingScheduledFuture cheduleAtFixedRate(Runnable command, long initialDelay, long period, TimeUnit unit, ScheduledExecutorService scheduler) {
    CancellableCommand cancellableCommand = new CancellableCommand(command);
    ScheduledFuture future = scheduler.scheduleAtFixedRate(cancellableCommand, initialDelay, period, unit);
    return new GracefullyStoppingScheduledFuture(future, cancellableCommand);
}

private GracefullyStoppingScheduledFuture(ScheduledFuture targetFuture, CancellableCommand command) {
    this.targetFuture = targetFuture;
    this.runnable = command;
}

private final ScheduledFuture targetFuture;
private final CancellableCommand runnable;

public void cancelAndBeSureOfTermination(boolean mayInterruptIfRunning) throws InterruptedException, ExecutionException {
    try {
        targetFuture.cancel(mayInterruptIfRunning);
    } finally {
        runnable.cancel();
    }
}

private static class CancellableCommand implements Runnable {

    private static final int NOT_EXECUTING = 0;
    private static final int IN_PROGRESS = 1;
    private static final int CANCELLED_WITHOUT_OBSTRUCTION = 2;
    private static final int CANCELLED_IN_MIDDLE_OF_PROGRESS = 3;

    private final AtomicInteger state = new AtomicInteger(NOT_EXECUTING);
    private final AtomicReference<Thread> executionThread = new AtomicReference<>();
    private final CompletableFuture<Void> cancellationFuture = new CompletableFuture<>();
    private final Runnable target;

    private CancellableCommand(Runnable target) {
        this.target = target;
    }

    public void cancel() throws ExecutionException, InterruptedException {
        if (executionThread.get() == Thread.currentThread()) {
            // cancel method was called from target by itself
            state.set(CANCELLED_IN_MIDDLE_OF_PROGRESS);
            return;
        }
        while (true) {
            if (state.get() == CANCELLED_WITHOUT_OBSTRUCTION) {
                return;
            }
            if (state.get() == CANCELLED_IN_MIDDLE_OF_PROGRESS) {
                cancellationFuture.get();
                return;
            }
            if (state.compareAndSet(NOT_EXECUTING, CANCELLED_WITHOUT_OBSTRUCTION)) {
                return;
            }
            if (state.compareAndSet(IN_PROGRESS, CANCELLED_IN_MIDDLE_OF_PROGRESS)) {
                cancellationFuture.get();
                return;
            }
        }
    }

    @Override
    public void run() {
        if (!state.compareAndSet(NOT_EXECUTING, IN_PROGRESS)) {
            notifyWaiters();
            return;
        }

        try {
            executionThread.set(Thread.currentThread());
            target.run();
        } finally {
            executionThread.set(null);
            if (!state.compareAndSet(IN_PROGRESS, NOT_EXECUTING)) {
                notifyWaiters();
            }
        }
    }

    private void notifyWaiters() {
        if (state.get() == CANCELLED_WITHOUT_OBSTRUCTION) {
            // no need to notify anything
            return;
        }
        // someone waits for cancelling
        cancellationFuture.complete(null);
        return;
    }

}