From 4100d9f88657cce68b042a9b155dd4881c7c0102 Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Thu, 6 Oct 2016 17:23:52 +0200 Subject: [PATCH] Consider delay for token renewal. Adopt TaskExecutor capabilities. Use scheduler API if executor is a `TaskScheduler`, delay task execution otherwise. Fixes gh-15. --- .../LifecycleAwareSessionManager.java | 28 +++++++++++++-- ...LifecycleAwareSessionManagerUnitTests.java | 34 +++++++++---------- 2 files changed, 42 insertions(+), 20 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 aeb10653..6e2d4b24 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 @@ -22,6 +22,7 @@ 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.util.Assert; import org.springframework.util.NumberUtils; import org.springframework.util.StringUtils; @@ -158,11 +159,13 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB private void scheduleRefresh() { + logger.debug("Refreshing token"); + LoginToken loginToken = (LoginToken) token; - int seconds = NumberUtils.convertNumberToTargetClass( + final int seconds = NumberUtils.convertNumberToTargetClass( Math.max(1, loginToken.getLeaseDuration() - REFRESH_PERIOD_BEFORE_EXPIRY), Integer.class); - taskExecutor.execute(new Runnable() { + final Runnable task = new Runnable() { @Override public void run() { try { @@ -175,7 +178,26 @@ public class LifecycleAwareSessionManager implements SessionManager, DisposableB logger.error("Cannot refresh VaultToken", e); } } - }, TimeUnit.SECONDS.toMillis(seconds)); + }; + + if (taskExecutor instanceof TaskScheduler) { + + TaskScheduler taskScheduler = (TaskScheduler) taskExecutor; + taskScheduler.scheduleWithFixedDelay(task, TimeUnit.SECONDS.toMillis(seconds)); + return; + } + + taskExecutor.execute(new Runnable() { + @Override + public void run() { + try { + Thread.sleep(TimeUnit.SECONDS.toMillis(seconds)); + task.run(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + }); } private static String buildExceptionMessage(VaultResponseEntity response) { 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 dced5e8c..e7f601c6 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 @@ -106,28 +106,28 @@ public class LifecycleAwareSessionManagerUnitTests { @Test public void shouldScheduleTokenRenewal() { - when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 10)); + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); sessionManager.getSessionToken(); - verify(taskExecutor).execute(any(Runnable.class), eq(5000L)); + verify(taskExecutor).execute(any(Runnable.class)); } @Test public void shouldRunTokenRenewal() { - when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 10)); - when(vaultClient.postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 10)), + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); + when(vaultClient.postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 5)), ArgumentMatchers.any(), any(Class.class))) .thenReturn(new ResponseEntity(null, HttpStatus.OK, null, null)); ArgumentCaptor runnableCaptor = ArgumentCaptor.forClass(Runnable.class); sessionManager.getSessionToken(); - verify(taskExecutor).execute(runnableCaptor.capture(), eq(5000L)); + verify(taskExecutor).execute(runnableCaptor.capture()); runnableCaptor.getValue().run(); - verify(vaultClient).postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 10)), + verify(vaultClient).postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 5)), ArgumentMatchers.any(), any(Class.class)); verify(clientAuthentication, times(1)).login(); } @@ -135,54 +135,54 @@ public class LifecycleAwareSessionManagerUnitTests { @Test public void shouldReScheduleTokenRenewalAfterSucessfulRenewal() { - when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 10)); - when(vaultClient.postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 10)), + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); + when(vaultClient.postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 5)), ArgumentMatchers.any(), any(Class.class))) .thenReturn(new ResponseEntity(null, HttpStatus.OK, null, null)); ArgumentCaptor runnableCaptor = ArgumentCaptor.forClass(Runnable.class); sessionManager.getSessionToken(); - verify(taskExecutor).execute(runnableCaptor.capture(), eq(5000L)); + verify(taskExecutor).execute(runnableCaptor.capture()); runnableCaptor.getValue().run(); - verify(taskExecutor, times(2)).execute(any(Runnable.class), anyLong()); + verify(taskExecutor, times(2)).execute(any(Runnable.class)); } @Test public void shouldNotReScheduleTokenRenewalAfterFailedRenewal() { - when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 10)); - when(vaultClient.postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 10)), + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); + when(vaultClient.postForEntity(eq("auth/token/renew-self"), eq(LoginToken.renewable("login", 5)), ArgumentMatchers.any(), any(Class.class))) .thenReturn(new ResponseEntity(null, HttpStatus.INTERNAL_SERVER_ERROR, null, null)); ArgumentCaptor runnableCaptor = ArgumentCaptor.forClass(Runnable.class); sessionManager.getSessionToken(); - verify(taskExecutor).execute(runnableCaptor.capture(), eq(5000L)); + verify(taskExecutor).execute(runnableCaptor.capture()); runnableCaptor.getValue().run(); - verify(taskExecutor, times(1)).execute(any(Runnable.class), anyLong()); + verify(taskExecutor, times(1)).execute(any(Runnable.class)); } @Test public void shouldObtainTokenIfNoTokenAvailable() { - when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 10)); + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); sessionManager.renewToken(); - assertThat(sessionManager.getSessionToken()).isEqualTo(LoginToken.renewable("login", 10)); + assertThat(sessionManager.getSessionToken()).isEqualTo(LoginToken.renewable("login", 5)); verify(clientAuthentication, times(1)).login(); } @Test public void renewShouldReportFalseIfTokenRenewalFails() { - when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 10)); + when(clientAuthentication.login()).thenReturn(LoginToken.renewable("login", 5)); when(vaultClient.postForEntity(anyString(), any(VaultToken.class), ArgumentMatchers.any(), any(Class.class))) .thenReturn(new ResponseEntity(null, HttpStatus.BAD_REQUEST, null, null));