common-utils/src/main/java/org/warp/commonutils/concurrency/executor/ScheduledTaskLifecycle.java

52 lines
1.2 KiB
Java

package org.warp.commonutils.concurrency.executor;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.locks.StampedLock;
import org.warp.commonutils.concurrency.atomicity.Atomic;
@Atomic
public class ScheduledTaskLifecycle {
private final StampedLock lock;
private final ConcurrentHashMap<ScheduledFuture<?>, Object> tasks = new ConcurrentHashMap<>();
public ScheduledTaskLifecycle() {
this.lock = new StampedLock();
}
/**
* Register a scheduled task
*/
public void registerScheduledTask(ScheduledFuture<?> task) {
this.tasks.put(task, new Object());
}
/**
* Mark this task as running.
* After calling this method, please call {@method endScheduledTask} inside a finally block!
*/
public void startScheduledTask() {
this.lock.readLock();
}
/**
* Mark this task as ended. Must be called after {@method startScheduledTask}
*/
public void endScheduledTask() {
this.lock.tryUnlockRead();
}
/**
* Cancel all scheduled tasks and wait all running methods to finish
*/
public void cancelAndWait() {
tasks.forEach((task, obj) -> {
task.cancel(false);
});
// Acquire a write lock to wait all tasks to end
lock.unlockWrite(lock.writeLock());
}
}