From 76e52bb7579a7bac842490f86829b7340f958089 Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Thu, 6 Oct 2016 20:46:33 +0200 Subject: [PATCH] Polishing. Document scheduler using blocked threads if AsyncTaskExecutor is not a scheduler. Improve scheduling with a one-shot-trigger. See gh-13, gh-15. --- .../LifecycleAwareSessionManager.java | 53 +++++++++++++++---- ...LifecycleAwareSessionManagerUnitTests.java | 20 +++++++ 2 files changed, 64 insertions(+), 9 deletions(-) diff --git a/spring-vault-core/src/main/java/org/springframework/vault/authentication/LifecycleAwareSessionManager.java b/spring-vault-core/src/main/java/org/springframework/vault/authentication/LifecycleAwareSessionManager.java index 6e2d4b24..a7c0cd31 100644 --- a/spring-vault-core/src/main/java/org/springframework/vault/authentication/LifecycleAwareSessionManager.java +++ b/spring-vault-core/src/main/java/org/springframework/vault/authentication/LifecycleAwareSessionManager.java @@ -15,14 +15,18 @@ */ package org.springframework.vault.authentication; +import java.util.Date; import java.util.Map; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.factory.DisposableBean; import org.springframework.core.task.AsyncTaskExecutor; import org.springframework.scheduling.TaskScheduler; +import org.springframework.scheduling.Trigger; +import org.springframework.scheduling.TriggerContext; import org.springframework.util.Assert; import org.springframework.util.NumberUtils; import org.springframework.util.StringUtils; @@ -49,9 +53,10 @@ import org.springframework.vault.support.VaultToken; */ public class LifecycleAwareSessionManager implements SessionManager, DisposableBean { - private final static Logger logger = LoggerFactory.getLogger(LifecycleAwareSessionManager.class); public static final int REFRESH_PERIOD_BEFORE_EXPIRY = 5; + private final static Logger logger = LoggerFactory.getLogger(LifecycleAwareSessionManager.class); + private final ClientAuthentication clientAuthentication; private final VaultClient vaultClient; private final AsyncTaskExecutor taskExecutor; @@ -104,6 +109,8 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB */ protected boolean renewToken() { + logger.info("Renewing token"); + if (token == null) { getSessionToken(); return false; @@ -137,7 +144,7 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB token = clientAuthentication.login(); if (isTokenRenewable()) { - scheduleRefresh(); + scheduleRenewal(); } } } @@ -157,9 +164,9 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB return false; } - private void scheduleRefresh() { + private void scheduleRenewal() { - logger.debug("Refreshing token"); + logger.info("Scheduling Token renewal"); LoginToken loginToken = (LoginToken) token; final int seconds = NumberUtils.convertNumberToTargetClass( @@ -171,19 +178,17 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB try { if (LifecycleAwareSessionManager.this.token != null && isTokenRenewable()) { if (renewToken()) { - scheduleRefresh(); + scheduleRenewal(); } } } catch (Exception e) { - logger.error("Cannot refresh VaultToken", e); + logger.error("Cannot renew VaultToken", e); } } }; if (taskExecutor instanceof TaskScheduler) { - - TaskScheduler taskScheduler = (TaskScheduler) taskExecutor; - taskScheduler.scheduleWithFixedDelay(task, TimeUnit.SECONDS.toMillis(seconds)); + scheduleTask((TaskScheduler) taskExecutor, seconds, task); return; } @@ -191,6 +196,9 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB @Override public void run() { try { + // TODO: Revisit this approach since it blocks a thread. Spinning up a managed + // TaskScheduler just for once-in-a-while token renewal seemed a bit over-sophisticated + // that's why we emulate a scheduler by blocking a Thread resource Thread.sleep(TimeUnit.SECONDS.toMillis(seconds)); task.run(); } catch (InterruptedException e) { @@ -200,6 +208,10 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB }); } + private void scheduleTask(TaskScheduler taskScheduler, int seconds, Runnable task) { + taskScheduler.schedule(task, new OneShotTrigger(seconds)); + } + private static String buildExceptionMessage(VaultResponseEntity response) { if (StringUtils.hasText(response.getMessage())) { @@ -208,4 +220,27 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB return String.format("Status %s URI %s", response.getStatusCode(), response.getUri()); } + + /** + * This one-shot trigger creates only one execution time to trigger an execution only once. + */ + private static class OneShotTrigger implements Trigger { + + private final AtomicBoolean fired = new AtomicBoolean(); + private final int seconds; + + OneShotTrigger(int seconds) { + this.seconds = seconds; + } + + @Override + public Date nextExecutionTime(TriggerContext triggerContext) { + + if (fired.compareAndSet(false, true)) { + return new Date(System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(seconds)); + } + + return null; + } + } } diff --git a/spring-vault-core/src/test/java/org/springframework/vault/authentication/LifecycleAwareSessionManagerUnitTests.java b/spring-vault-core/src/test/java/org/springframework/vault/authentication/LifecycleAwareSessionManagerUnitTests.java index e7f601c6..631d106d 100644 --- a/spring-vault-core/src/test/java/org/springframework/vault/authentication/LifecycleAwareSessionManagerUnitTests.java +++ b/spring-vault-core/src/test/java/org/springframework/vault/authentication/LifecycleAwareSessionManagerUnitTests.java @@ -29,6 +29,8 @@ import org.mockito.Mock; import org.mockito.runners.MockitoJUnitRunner; import org.springframework.core.task.AsyncTaskExecutor; import org.springframework.http.HttpStatus; +import org.springframework.scheduling.Trigger; +import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.vault.client.VaultClient; import org.springframework.vault.client.VaultResponseEntity; import org.springframework.vault.support.VaultToken; @@ -45,6 +47,8 @@ public class LifecycleAwareSessionManagerUnitTests { @Mock private AsyncTaskExecutor taskExecutor; + @Mock private ThreadPoolTaskScheduler taskScheduler; + @Mock private VaultClient vaultClient; private LifecycleAwareSessionManager sessionManager; @@ -150,6 +154,22 @@ public class LifecycleAwareSessionManagerUnitTests { verify(taskExecutor, times(2)).execute(any(Runnable.class)); } + @Test + public void shouldUseTaskScheduler() { + + sessionManager = new LifecycleAwareSessionManager(clientAuthentication, taskScheduler, vaultClient); + + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); + + ArgumentCaptor triggerCaptor = ArgumentCaptor.forClass(Trigger.class); + + sessionManager.getSessionToken(); + verify(taskScheduler).schedule(any(Runnable.class), triggerCaptor.capture()); + + assertThat(triggerCaptor.getValue().nextExecutionTime(null)).isNotNull(); + assertThat(triggerCaptor.getValue().nextExecutionTime(null)).isNull(); + } + @Test public void shouldNotReScheduleTokenRenewalAfterFailedRenewal() {