// Path: src/main/java/com/baidu/unbiz/multitask/exception/TaskTimeoutException.java
// public class TaskTimeoutException extends RuntimeException {
// 
//     private static final long serialVersionUID = -7375423850222016116L;
// 
//     public TaskTimeoutException(String msg) {
//         super(msg);
//     }
// 
//     public TaskTimeoutException(Throwable cause) {
//         super(cause);
//     }
// 
//     public TaskTimeoutException(String message, Throwable cause) {
//         super(message, cause);
//     }
// }


import java.util.Map;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import com.baidu.unbiz.multitask.exception.TaskTimeoutException;
package com.baidu.unbiz.multitask.task.thread;




/**
 * 一组工作单元,可视为一组原子操作
 *
 * @author wangchongjie
 * @since 2015-12-3 下午3:56:36
 */
public class WorkUnit {

    private static final Log LOG = LogFactory.getLog(TaskManager.class);

    private CompletionService<Void> completion;
    private Map<Future<?>, Runnable> futureMap = new ConcurrentHashMap<Future<?>, Runnable>();
    private int taskCount;
    private Thread parentThread;

    /**
     * 构造函数，线程池委托给调用方
     *
     * @param pool
     */
    public WorkUnit(Executor pool) {
        parentThread = Thread.currentThread();
        completion = new ExecutorCompletionService<Void>(pool);
    }

    /**
     * 提交可并行的任务
     *
     * @param runnable
     */
    public void submit(Runnable runnable) {
        futureMap.put(completion.submit(runnable, null), runnable);
        taskCount++;
    }

    /**
     * 等待结果返回，带默认超时时间
     */
    public void waitForCompletion() {
        this.waitForCompletion(TaskManager.config().taskTimeoutMillSeconds());
    }

    /**
     * 等待结果返回，显示指定超时时间，若超时则取消任务并释放资源
     */
    public void waitForCompletion(long timeoutMillSeconds) {
        for (int i = 0; i < taskCount; i++) {
            Future<Void> future;
            try {
                long timeout =
                        timeoutMillSeconds > 0 ? timeoutMillSeconds : TaskManager.config().taskTimeoutMillSeconds();
                future = completion.poll(timeout, TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                LOG.error("wait for execute completion failed,e=" + e, e);
// Path: src/main/java/com/baidu/unbiz/multitask/exception/TaskTimeoutException.java
// public class TaskTimeoutException extends RuntimeException {
// 
//     private static final long serialVersionUID = -7375423850222016116L;
// 
//     public TaskTimeoutException(String msg) {
//         super(msg);
//     }
// 
//     public TaskTimeoutException(Throwable cause) {
//         super(cause);
//     }
// 
//     public TaskTimeoutException(String message, Throwable cause) {
//         super(message, cause);
//     }
// }


// Path: src/main/java/com/baidu/unbiz/multitask/task/thread/WorkUnit.java
import java.util.Map;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import com.baidu.unbiz.multitask.exception.TaskTimeoutException;

package com.baidu.unbiz.multitask.task.thread;




/**
 * 一组工作单元,可视为一组原子操作
 *
 * @author wangchongjie
 * @since 2015-12-3 下午3:56:36
 */
public class WorkUnit {

    private static final Log LOG = LogFactory.getLog(TaskManager.class);

    private CompletionService<Void> completion;
    private Map<Future<?>, Runnable> futureMap = new ConcurrentHashMap<Future<?>, Runnable>();
    private int taskCount;
    private Thread parentThread;

    /**
     * 构造函数，线程池委托给调用方
     *
     * @param pool
     */
    public WorkUnit(Executor pool) {
        parentThread = Thread.currentThread();
        completion = new ExecutorCompletionService<Void>(pool);
    }

    /**
     * 提交可并行的任务
     *
     * @param runnable
     */
    public void submit(Runnable runnable) {
        futureMap.put(completion.submit(runnable, null), runnable);
        taskCount++;
    }

    /**
     * 等待结果返回，带默认超时时间
     */
    public void waitForCompletion() {
        this.waitForCompletion(TaskManager.config().taskTimeoutMillSeconds());
    }

    /**
     * 等待结果返回，显示指定超时时间，若超时则取消任务并释放资源
     */
    public void waitForCompletion(long timeoutMillSeconds) {
        for (int i = 0; i < taskCount; i++) {
            Future<Void> future;
            try {
                long timeout =
                        timeoutMillSeconds > 0 ? timeoutMillSeconds : TaskManager.config().taskTimeoutMillSeconds();
                future = completion.poll(timeout, TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                LOG.error("wait for execute completion failed,e=" + e, e);
                throw new TaskTimeoutException(e);