package com.lifeonwalden.springscheduling.concurrent;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.scheduling.concurrent.ExecutorConfigurationSupport;

import com.lifeonwalden.springscheduling.task.Task;
import com.lifeonwalden.springscheduling.task.Work;

public class ThreadPoolTaskSchedulerWithRetry extends ExecutorConfigurationSupport {
    private final static Logger logger = LogManager.getLogger(ThreadPoolTaskSchedulerWithRetry.class);

    private static final long serialVersionUID = -4703424603462186353L;

    private volatile ScheduledThreadPoolExecutor scheduledExecutor;

    private volatile int poolSize = 1;

    private ConcurrentHashMap<String, Task> taskMap = new ConcurrentHashMap<>();

    private ConcurrentHashMap<String, Work> workMap = new ConcurrentHashMap<>();

    public void setPoolSize(int poolSize) {
        this.poolSize = poolSize;
        this.scheduledExecutor.setCorePoolSize(poolSize);
    }

    public int getPoolSize() {
        if (this.scheduledExecutor == null) {
            // Not initialized yet: assume initial pool size.
            return this.poolSize;
        }
        return getScheduledThreadPoolExecutor().getPoolSize();
    }

    @Override
    protected ExecutorService initializeExecutor(ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
        this.scheduledExecutor = createExecutor(this.poolSize, threadFactory, rejectedExecutionHandler);

        return this.scheduledExecutor;
    }

    protected ScheduledThreadPoolExecutor createExecutor(int poolSize, ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
        return new ScheduledThreadPoolExecutor(poolSize, threadFactory, rejectedExecutionHandler);
    }

    public ScheduledThreadPoolExecutor getScheduledThreadPoolExecutor() {
        return this.scheduledExecutor;
    }

    public ScheduledFuture<?> schedule(Task task) {
        logger.info("Load task [{}]", task.getName());

        ScheduledThreadPoolExecutor executor = getScheduledThreadPoolExecutor();
        task.setExecutor(executor);
        taskMap.put(task.getId(), task);
        return task.schedule();
    }

    public Task getTask(String taskId) {
        return taskMap.get(taskId);
    }

    public void addWork(Work work) {
        logger.info("Load work [{}]", work.getName());

        this.workMap.put(work.getId(), work);
    }

    public void removeTask(Task task) {
        logger.info("remove task [{}]", task.getName());

        Task preTask = taskMap.remove(task.getId());
        if (null != preTask) {
            preTask.setStopped(true);
            preTask.cancel(false);
        }
    }

    public void updateTask(Task task) {
        logger.info("update task [{}]", task.getName());

        removeTask(task);
        schedule(task);
    }

    /**
     * only accept work
     * 
     * @param taskId
     * @param param
     * @param async 是否异步执行
     */
    public void execute(String taskId, Map<String, Object> param, boolean async) {
        Work work = null;
        if (null == taskId || taskId.length() == 0 || (work = this.workMap.get(taskId)) == null) {
            logger.error("Can't find task for [{}]", taskId);

            return;
        }

        if (async) {
            final Work _work = work;
            new Thread(new Runnable() {

                @Override
                public void run() {
                    _work.doJob(param);
                }
            }).start();
        } else {
            work.doJob(param);
        }
    }

    public Work getWork(String workId) {
        return workMap.get(workId);
    }
}

@@ -61,321 +61,105 @@
 ort 
-java.util.Map;%0D%0Aimport java.util.concurrent.ConcurrentHashMap;%0D%0Aimport java.util.concurrent.ExecutorService;%0D%0Aimport java.util.concurrent.RejectedExecutionHandler;%0D%0Aimport java.util.concurrent.ScheduledFuture;%0D%0Aimport java.util.concurrent.ScheduledThreadPoolExecutor;%0D%0Aimport java.util.concurrent.ThreadFactory;%0D%0A
+com.lifeonwalden.springscheduling.task.Task;%0D%0Aimport com.lifeonwalden.springscheduling.task.Work;
 %0D%0Aim
@@ -335,104 +335,53 @@
 ort 
-com.lifeonwalden.springscheduling.task.Task;%0D%0Aimport com.lifeonwalden.springscheduling.task.Work
+java.util.Map;%0D%0Aimport java.util.concurrent.*
 ;%0D%0A%0D
@@ -944,14 +944,13 @@
 lic 
-void s
+int g
 etPo
@@ -956,28 +956,16 @@
 oolSize(
-int poolSize
 ) %7B%0D%0A   
@@ -973,73 +973,144 @@
     
+if (
 this.
-poolSize = poolSize;%0D%0A        this.scheduledExecutor.setCore
+scheduledExecutor == null) %7B%0D%0A            return this.poolSize;%0D%0A        %7D%0D%0A        return getScheduledThreadPoolExecutor().get
 Pool
@@ -1114,24 +1114,16 @@
 oolSize(
-poolSize
 );%0D%0A    
@@ -1130,37 +1130,38 @@
 %7D%0D%0A%0D%0A    public 
-int g
+void s
 etPoolSize() %7B%0D%0A
@@ -1159,220 +1159,98 @@
 ize(
-) %7B%0D%0A        if (this.scheduledExecutor == null) %7B%0D%0A            // Not initialized yet: assume initial 
+int poolSize) %7B%0D%0A        this.
 pool
- size.%0D%0A            return this.poolSize;%0D%0A        %7D%0D%0A        return getScheduledThreadPoolExecutor().get
+Size = poolSize;%0D%0A        this.scheduledExecutor.setCore
 Pool
@@ -1246,32 +1246,40 @@
 setCorePoolSize(
+poolSize
 );%0D%0A    %7D%0D%0A%0D%0A   
@@ -3004,17 +3004,16 @@
 %0D%0A     *
- 
 %0D%0A     *
@@ -3069,16 +3069,17 @@
 m async 
+ 
 %E6%98%AF%E5%90%A6%E5%BC%82%E6%AD%A5%E6%89%A7%E8%A1%8C%0D%0A

add getTask