net.dongliu.prettypb.rpc.common.TaskSet Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of prettypb-rpc Show documentation
proto rpc libs, compatible with proto-rpc-pro
There is a newer version: 0.3.5
package net.dongliu.prettypb.rpc.common;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * hold tasks with timeout checker
 *
 * @author Dong Liu
 */
public class TaskSet> {
    private final ConcurrentHashMap taskMap = new ConcurrentHashMap<>();

    //private PriorityBlockingQueue queue = new PriorityBlockingQueue<>();

    private final AtomicBoolean closed = new AtomicBoolean(false);

    private final Thread timeoutChecker;

    private static Logger logger = LoggerFactory.getLogger(TaskSet.class);

    private static final AtomicInteger threadNum = new AtomicInteger(1);

    public TaskSet() {
        timeoutChecker = startCheckTimeout();
    }

    /**
     * add one task
     *
     * @param task
     * @return true if add successfully
     */
    public boolean add(T task) {
        if (closed.get()) {
            task.onClosed();
            return false;
        }
        boolean success = taskMap.putIfAbsent(task.id(), task) == null;

        // check again
        if (success && closed.get()) {
            consume(task.id(), closeCallBack);
            return false;
        }
        return success;
    }

    /**
     * remove task with id
     *
     * @param id
     */
    public void consume(int id) {
        taskMap.remove(id);
    }

    /**
     * remove task with id, if not null, process with callback.
     * exceptions will be catched and logged
     *
     * @param id
     */
    public void consume(int id, TaskCallBack callBack) {
        T task = taskMap.remove(id);
        if (task != null) {
            try {
                callBack.onTask(task);
            } catch (Exception e) {
                logger.error("", e);
            }
        }
    }

    private Thread startCheckTimeout() {
        Thread timeoutChecker = new Thread(new Runnable() {
            @Override
            public void run() {
                checkTimeout();
            }
        }, "Rpc-timeout-checker-" + threadNum.getAndIncrement());
        timeoutChecker.start();
        return timeoutChecker;
    }

    /**
     * find timeout tasks and cancel
     */
    private void checkTimeout() {

        while (!closed.get()) {
            Collection tasks = taskMap.values();
            for (Task task : tasks) {
                if (task.timeout()) {
                    consume(task.id(), timeoutCallBack);
                }
            }
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                logger.debug("timeout checker interrupted");
                break;
            }
        }
    }

    /**
     * clear all task
     */
    public void close() {
        if (closed.compareAndSet(false, true)) {
            timeoutChecker.interrupt();
            while (!taskMap.isEmpty()) {
                Collection tasks = taskMap.values();
                for (T task : tasks) {
                    consume(task.id(), closeCallBack);
                }
                taskMap.clear();
            }
        }
    }

    private final TaskCallBack closeCallBack = new TaskCallBack() {
        @Override
        public void onTask(Task task) {
            task.onClosed();
        }
    };

    private final TaskCallBack timeoutCallBack = new TaskCallBack() {
        @Override
        public void onTask(Task task) {
            task.onTimeout();
        }
    };
}