traf-prcs-server/src/main/java/com/its/traf/scheduler/JobScheduler.java

package com.its.traf.scheduler;

import com.its.traf.config.SchedulerConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
import org.springframework.scheduling.support.CronTrigger;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ScheduledFuture;

@Slf4j
@Configuration
@Service
public class JobScheduler implements SchedulingConfigurer {

    private final SchedulerConfig schedulerConfig;

    private ScheduledTaskRegistrar scheduledTaskRegistrar;
    private Map<String, ScheduledFuture> futureMap = new HashMap<>();
    private Map<String, JobScheduleInfo> batchScheduleMap = new HashMap<>();

    public JobScheduler(SchedulerConfig schedulerConfig) {
        this.schedulerConfig = schedulerConfig;
    }

    @PostConstruct
    public void init() {
        /*List<JobScheduleInfo> jobs = this.jobSchedulerConfig.getJobs();
        for(JobScheduleInfo info : jobs) {
            addJob(info);
            this.scheduleMap.put(info.getName(), info);
        }*/
    }

    @Bean
    public TaskScheduler poolScheduler() {
        ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
        scheduler.setThreadNamePrefix("jobScheduler-");
        scheduler.setPoolSize(this.schedulerConfig.getPoolCore());
        scheduler.initialize();
        return scheduler;
    }

    @Override
    public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
        /*if (this.scheduledTaskRegistrar == null) {
            this.scheduledTaskRegistrar = taskRegistrar;
        }
        if (taskRegistrar.getScheduler() == null) {
            taskRegistrar.setScheduler(poolScheduler());
        }*/

        taskRegistrar.setScheduler(poolScheduler());

        List<JobScheduleInfo> jobs = this.schedulerConfig.getBatchJobs();
        for(JobScheduleInfo info : jobs) {
            this.batchScheduleMap.put(info.getName(), info);
            taskRegistrar.addTriggerTask(() -> batchJobScheduleCron(info.getName()), t -> {
                CronTrigger crontrigger = new CronTrigger(info.getCron());
                return crontrigger.nextExecutionTime(t);
            });
        }

        /*
        taskRegistrar.setScheduler(poolScheduler());

        // Random next execution time.
        taskRegistrar.addTriggerTask(() -> scheduleDynamically(), t -> {
            Calendar nextExecutionTime = new GregorianCalendar();
            Date lastActualExecutionTime = t.lastActualExecutionTime();
            nextExecutionTime.setTime(lastActualExecutionTime != null ? lastActualExecutionTime : new Date());
            nextExecutionTime.add(Calendar.SECOND, getNextExecutionTime()); // This is where we set the next execution time.
            return nextExecutionTime.getTime();
        });

        // Fixed next execution time.
        taskRegistrar.addTriggerTask(() -> scheduleFixed(), t -> {
            Calendar nextExecutionTime = new GregorianCalendar();
            Date lastActualExecutionTime = t.lastActualExecutionTime();
            nextExecutionTime.setTime(lastActualExecutionTime != null ? lastActualExecutionTime : new Date());
            nextExecutionTime.add(Calendar.SECOND, 7); // This is where we set the next execution time.
            return nextExecutionTime.getTime();
        });

        // Next execution time is taken from DB, so if the value in DB changes, next execution time will change too.
        taskRegistrar.addTriggerTask(() -> scheduledDatabase(repo.findById("next_exec_time").get().getConfigValue()), t -> {
            Calendar nextExecutionTime = new GregorianCalendar();
            Date lastActualExecutionTime = t.lastActualExecutionTime();
            nextExecutionTime.setTime(lastActualExecutionTime != null ? lastActualExecutionTime : new Date());
            nextExecutionTime.add(Calendar.SECOND, Integer.parseInt(repo.findById("next_exec_time").get().getConfigValue()));
            return nextExecutionTime.getTime();
        });

        // or cron way, you can also get the expression from DB or somewhere else just like we did above.
        taskRegistrar.addTriggerTask(() -> scheduleCron(repo.findById("next_exec_time").get().getConfigValue()), t -> {
            CronTrigger crontrigger = new CronTrigger(repo.findById("next_exec_time").get().getConfigValue());
            return crontrigger.nextExecutionTime(t);
        });*/
    }

    public boolean addJob(JobScheduleInfo jobInfo) {
        if (this.futureMap.containsKey(jobInfo.getName())) {
            return false;
        }

/*        this.scheduledTaskRegistrar.getScheduler().addTriggerTask(() -> scheduleCron(repo.findById("next_exec_time").get().getConfigValue()), t -> {
            CronTrigger crontrigger = new CronTrigger(repo.findById("next_exec_time").get().getConfigValue());
            return crontrigger.nextExecutionTime(t);
        });

        ScheduledFuture future = this.scheduledTaskRegistrar.getScheduler().schedule(() -> methodToBeExecuted(), t -> {
            Calendar nextExecutionTime = new GregorianCalendar();
            Date lastActualExecutionTime = t.lastActualExecutionTime();
            nextExecutionTime.setTime(lastActualExecutionTime != null ? lastActualExecutionTime : new Date());
            nextExecutionTime.add(Calendar.SECOND, 5);
            return nextExecutionTime.getTime();
        });*/

        configureTasks(this.scheduledTaskRegistrar);
        //this.futureMap.put(jobInfo.getName(), future);
        return true;
    }

    public boolean removeJob(String name) {
        if (!this.futureMap.containsKey(name)) {
            return false;
        }
        ScheduledFuture future = this.futureMap.get(name);
        future.cancel(true);
        this.futureMap.remove(name);
        return true;
    }

    public void scheduleDynamically() {
        log.info("scheduleDynamically: Next execution time of this changes every time between 1 and 5 seconds");
    }

    public void scheduleFixed() {
        log.info("scheduleFixed: Next execution time of this will always be 7 seconds");
    }

    public void scheduledDatabase(String time) {
        log.info("scheduledDatabase: Next execution time of this will be taken from DB -> {}", time);
    }

    public void batchJobScheduleCron(String jobName) {
        JobScheduleInfo job = this.batchScheduleMap.get(jobName);
        if (job != null) {
            log.info("scheduleCron: Next execution time of this taken from cron expression -> {}", job.toString());
        }
    }

    public int getNextExecutionTime() {
        return new Random().nextInt(5) + 1;
    }
}