001/*
002 * Copyright (C) Photon Vision.
003 *
004 * This program is free software: you can redistribute it and/or modify
005 * it under the terms of the GNU General Public License as published by
006 * the Free Software Foundation, either version 3 of the License, or
007 * (at your option) any later version.
008 *
009 * This program is distributed in the hope that it will be useful,
010 * but WITHOUT ANY WARRANTY; without even the implied warranty of
011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
012 * GNU General Public License for more details.
013 *
014 * You should have received a copy of the GNU General Public License
015 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
016 */
017
018package org.photonvision.common.util;
019
020import java.util.concurrent.*;
021import org.jetbrains.annotations.NotNull;
022import org.photonvision.common.logging.LogGroup;
023import org.photonvision.common.logging.Logger;
024
025public class TimedTaskManager {
026    private static final Logger logger = new Logger(TimedTaskManager.class, LogGroup.General);
027
028    private static class Singleton {
029        public static final TimedTaskManager INSTANCE = new TimedTaskManager();
030    }
031
032    public static TimedTaskManager getInstance() {
033        return Singleton.INSTANCE;
034    }
035
036    private static class CaughtThreadFactory implements ThreadFactory {
037        private static final ThreadFactory defaultThreadFactory = Executors.defaultThreadFactory();
038
039        @Override
040        public Thread newThread(@NotNull Runnable r) {
041            Thread thread = defaultThreadFactory.newThread(r);
042            thread.setUncaughtExceptionHandler(
043                    (t, e) -> logger.error("TimedTask threw uncaught exception!", e));
044            return thread;
045        }
046    }
047
048    private static class CaughtScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor {
049        public CaughtScheduledThreadPoolExecutor(int corePoolSize, ThreadFactory threadFactory) {
050            super(corePoolSize, threadFactory);
051        }
052
053        public Runnable wrap(Runnable runnable) {
054            return () -> {
055                try {
056                    runnable.run();
057                } catch (Throwable t) {
058                    logger.error("Exception thrown by threadpool: " + t.getMessage(), t);
059                }
060            };
061        }
062
063        public <V> Callable<V> wrap(Callable<V> runnable) {
064            return () -> {
065                try {
066                    return runnable.call();
067                } catch (Throwable t) {
068                    logger.error("Exception thrown by threadpool: " + t.getMessage(), t);
069                    return null;
070                }
071            };
072        }
073
074        @Override
075        public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
076            return super.schedule(wrap(callable), delay, unit);
077        }
078
079        @Override
080        public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
081            return super.schedule(wrap(command), delay, unit);
082        }
083
084        @Override
085        public ScheduledFuture<?> scheduleAtFixedRate(
086                Runnable command, long initialDelay, long period, TimeUnit unit) {
087            return super.scheduleAtFixedRate(wrap(command), initialDelay, period, unit);
088        }
089
090        @Override
091        public ScheduledFuture<?> scheduleWithFixedDelay(
092                Runnable command, long initialDelay, long delay, TimeUnit unit) {
093            return super.scheduleWithFixedDelay(wrap(command), initialDelay, delay, unit);
094        }
095    }
096
097    private final CaughtScheduledThreadPoolExecutor timedTaskExecutorPool =
098            new CaughtScheduledThreadPoolExecutor(2, new CaughtThreadFactory());
099    private final ConcurrentHashMap<String, Future<?>> activeTasks = new ConcurrentHashMap<>();
100
101    public void addTask(String identifier, Runnable runnable, long millisInterval) {
102        if (!activeTasks.containsKey(identifier)) {
103            var future =
104                    timedTaskExecutorPool.scheduleAtFixedRate(
105                            runnable, 0, millisInterval, TimeUnit.MILLISECONDS);
106            activeTasks.put(identifier, future);
107        }
108    }
109
110    public void addTask(
111            String identifier, Runnable runnable, long millisStartDelay, long millisInterval) {
112        if (!activeTasks.containsKey(identifier)) {
113            var future =
114                    timedTaskExecutorPool.scheduleAtFixedRate(
115                            runnable, millisStartDelay, millisInterval, TimeUnit.MILLISECONDS);
116            activeTasks.put(identifier, future);
117        }
118    }
119
120    public void addOneShotTask(Runnable runnable, long millisStartDelay) {
121        timedTaskExecutorPool.schedule(runnable, millisStartDelay, TimeUnit.MILLISECONDS);
122    }
123
124    public void cancelTask(String identifier) {
125        var future = activeTasks.getOrDefault(identifier, null);
126        if (future != null) {
127            future.cancel(true);
128            activeTasks.remove(identifier);
129        }
130    }
131
132    public boolean taskActive(String identifier) {
133        return activeTasks.containsKey(identifier);
134    }
135}