common-utils/src/main/java/org/warp/commonutils/concurrency/executor/ScheduledTaskLifecycle.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());
	}
}