Improve TimerSPI to support cancelling of previously scheduled tasks

services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java

@@ -146,8 +146,8 @@ public class KeycloakApplication extends Application {
         long interval = Config.scope("scheduled").getLong("interval", 60L) * 1000;
 
         TimerProvider timer = sessionFactory.create().getProvider(TimerProvider.class);
-        timer.schedule(new ScheduledTaskRunner(sessionFactory, new ClearExpiredAuditEvents()), interval);
+        timer.schedule(new ScheduledTaskRunner(sessionFactory, new ClearExpiredAuditEvents()), interval, "ClearExpiredAuditEvents");
-        timer.schedule(new ScheduledTaskRunner(sessionFactory, new ClearExpiredUserSessions()), interval);
+        timer.schedule(new ScheduledTaskRunner(sessionFactory, new ClearExpiredUserSessions()), interval, "ClearExpiredUserSessions");
     }
 
     public KeycloakSessionFactory getSessionFactory() {

timer/api/src/main/java/org/keycloak/timer/TimerProvider.java

@@ -7,6 +7,8 @@ import org.keycloak.provider.Provider;
  */
 public interface TimerProvider extends Provider {
 
-    public void schedule(Runnable runnable, long interval);
+    public void schedule(Runnable runnable, long interval, String taskName);
+
+    public void cancelTask(String taskName);
 
 }

timer/basic/pom.xml

@@ -31,6 +31,11 @@
             <version>${project.version}</version>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jboss.logging</groupId>
+            <artifactId>jboss-logging</artifactId>
+            <scope>provided</scope>
+        </dependency>
     </dependencies>
 
 </project>

timer/basic/src/main/java/org/keycloak/timer/basic/BasicTimerProvider.java

@@ -1,5 +1,6 @@
 package org.keycloak.timer.basic;
 
+import org.jboss.logging.Logger;
 import org.keycloak.timer.TimerProvider;
 
 import java.util.Timer;
@@ -10,15 +11,18 @@ import java.util.TimerTask;
  */
 public class BasicTimerProvider implements TimerProvider {
 
-    private Timer timer;
+    private static final Logger logger = Logger.getLogger(BasicTimerProvider.class);
 
-    public BasicTimerProvider(Timer timer) {
+    private final Timer timer;
+    private final BasicTimerProviderFactory factory;
 
+    public BasicTimerProvider(Timer timer, BasicTimerProviderFactory factory) {
         this.timer = timer;
+        this.factory = factory;
     }
 
     @Override
-    public void schedule(final Runnable runnable, final long interval) {
+    public void schedule(final Runnable runnable, final long interval, String taskName) {
         TimerTask task = new TimerTask() {
             @Override
             public void run() {
@@ -26,9 +30,25 @@ public class BasicTimerProvider implements TimerProvider {
             }
         };
 
+        TimerTask existingTask = factory.putTask(taskName, task);
+        if (existingTask != null) {
+            logger.infof("Existing timer task '%s' found. Cancelling it", taskName);
+            existingTask.cancel();
+        }
+
+        logger.infof("Starting task '%s' with interval '%d'", taskName, interval);
         timer.schedule(task, interval, interval);
     }
 
+    @Override
+    public void cancelTask(String taskName) {
+        TimerTask existingTask = factory.removeTask(taskName);
+        if (existingTask != null) {
+            logger.infof("Cancelling task '%s'", taskName);
+            existingTask.cancel();
+        }
+    }
+
     @Override
     public void close() {
         // do nothing

timer/basic/src/main/java/org/keycloak/timer/basic/BasicTimerProviderFactory.java

@@ -6,6 +6,9 @@ import org.keycloak.timer.TimerProvider;
 import org.keycloak.timer.TimerProviderFactory;
 
 import java.util.Timer;
+import java.util.TimerTask;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
@@ -14,9 +17,11 @@ public class BasicTimerProviderFactory implements TimerProviderFactory {
 
     private Timer timer;
 
+    private ConcurrentMap<String, TimerTask> scheduledTasks = new ConcurrentHashMap<String, TimerTask>();
+
     @Override
     public TimerProvider create(KeycloakSession session) {
-        return new BasicTimerProvider(timer);
+        return new BasicTimerProvider(timer, this);
     }
 
     @Override
@@ -35,4 +40,12 @@ public class BasicTimerProviderFactory implements TimerProviderFactory {
         return "basic";
     }
 
+    protected TimerTask putTask(String taskName, TimerTask task) {
+        return scheduledTasks.put(taskName, task);
+    }
+
+    protected TimerTask removeTask(String taskName) {
+        return scheduledTasks.remove(taskName);
+    }
+
 }