Resume lease renewal after restart.

Secret renewals are resumed after restarting secrets by re-registering LeaseRenewalScheduler.

Also, guard concurrent renewals that might have been in progress while the renewal scheduler has been in progress.

Closes gh-867

Author: mp911de

spring-vault-core/src/main/java/org/springframework/vault/core/lease/SecretLeaseContainer.java

@@ -549,6 +549,7 @@ void restartSecrets() {
 					Lease previousLease = getPreviousLease(previousLeases, requestedSecret);
 
 					try {
+						this.renewals.put(requestedSecret, renewalScheduler);
 						doStart(requestedSecret, renewalScheduler, (secrets, lease) -> {
 							onSecretsRotated(requestedSecret, previousLease, lease, secrets.getRequiredData());
 						}, () -> {
@@ -807,7 +808,19 @@ private Lease doRenew(Lease lease) {
 	protected void onLeaseExpired(RequestedSecret requestedSecret, Lease lease) {
 
 		if (requestedSecret.getMode() == Mode.ROTATE) {
-			doStart(requestedSecret, this.renewals.get(requestedSecret), (secrets, currentLease) -> {
+			LeaseRenewalScheduler renewalScheduler = this.renewals.get(requestedSecret);
+			// prevent races for concurrent renewals of the same secret using different
+			// leases
+			if (renewalScheduler == null || !renewalScheduler.leaseEquals(lease)) {
+				logger.debug(String.format(
+						"Skipping rotation after renewal expiry for secret %s with lease %s as no LeaseRenewalScheduler is found. This can happen if leases have been restarted while concurrent expiry processing.",
+						requestedSecret.getPath(), lease.getLeaseId()));
+
+				super.onLeaseExpired(requestedSecret, lease);
+				return;
+			}
+
+			doStart(requestedSecret, renewalScheduler, (secrets, currentLease) -> {
 				onSecretsRotated(requestedSecret, lease, currentLease, secrets.getRequiredData());
 			}, () -> super.onLeaseExpired(requestedSecret, lease));
 		}
@@ -1023,6 +1036,10 @@ private boolean isLeaseRotateOnly(Lease lease, RequestedSecret requestedSecret)
 					&& requestedSecret.getMode() == Mode.ROTATE;
 		}
 
+		boolean leaseEquals(Lease lease) {
+			return getLease() == lease;
+		}
+
 	}
 
 	private class LeaseAuthenticationEventListener implements AuthenticationListener, AuthenticationErrorListener {

spring-vault-core/src/test/java/org/springframework/vault/core/lease/SecretLeaseContainerUnitTests.java

@@ -37,6 +37,7 @@
 import org.springframework.http.HttpStatus;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.Trigger;
+import org.springframework.test.util.ReflectionTestUtils;
 import org.springframework.vault.VaultException;
 import org.springframework.vault.core.RestOperationsCallback;
 import org.springframework.vault.core.VaultOperations;
@@ -479,6 +480,78 @@ void shouldNotRotateExpiringLease() {
 		assertThat(((SecretLeaseCreatedEvent) events.get(1)).getSecrets()).containsOnlyKeys("foo");
 	}
 
+	@SuppressWarnings("unchecked")
+	@Test
+	void secretShouldContinueOnRestartSecrets() {
+
+		when(this.taskScheduler.schedule(any(Runnable.class), any(Trigger.class))).thenReturn(this.scheduledFuture);
+
+		VaultResponse first = createSecrets();
+		VaultResponse second = createSecrets();
+		second.setData(Collections.singletonMap("foo", (Object) "bar"));
+
+		when(this.vaultOperations.read(this.requestedSecret.getPath())).thenReturn(first, second);
+		when(this.vaultOperations.doWithSession(any(RestOperationsCallback.class)))
+				.thenReturn(Lease.of("after_restart", Duration.ofSeconds(1), true));
+
+		RequestedSecret secret = RequestedSecret.rotating("my-secret");
+
+		this.secretLeaseContainer.setExpiryThreshold(Duration.ofSeconds(1));
+		this.secretLeaseContainer.addRequestedSecret(secret);
+		this.secretLeaseContainer.start();
+
+		ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+		this.secretLeaseContainer.restartSecrets();
+
+		verify(this.taskScheduler, times(2)).schedule(captor.capture(), any(Trigger.class));
+		captor.getAllValues().get(0).run(); // old lease run
+		captor.getAllValues().get(1).run(); // new lease run
+
+		// one from restartSecrets, the second from the scheduler detecting expiry
+		verify(this.leaseListenerAdapter, times(2)).onLeaseEvent(any(SecretLeaseRotatedEvent.class));
+		verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(SecretLeaseExpiredEvent.class));
+		verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(SecretLeaseErrorEvent.class));
+	}
+
+	@SuppressWarnings("unchecked")
+	@Test
+	void concurrentRenewalAfterRestartSecretShouldContinueOnRestartSecrets() {
+
+		when(this.taskScheduler.schedule(any(Runnable.class), any(Trigger.class))).thenReturn(this.scheduledFuture);
+
+		VaultResponse first = createSecrets();
+		VaultResponse second = createSecrets();
+		second.setData(Collections.singletonMap("foo", (Object) "bar"));
+
+		when(this.vaultOperations.read(this.requestedSecret.getPath())).thenReturn(first, second);
+		when(this.vaultOperations.doWithSession(any(RestOperationsCallback.class)))
+				.thenReturn(Lease.of("after_restart", Duration.ofSeconds(1), true));
+
+		RequestedSecret secret = RequestedSecret.rotating("my-secret");
+
+		this.secretLeaseContainer.setExpiryThreshold(Duration.ofSeconds(1));
+		this.secretLeaseContainer.addRequestedSecret(secret);
+		this.secretLeaseContainer.start();
+
+		Map<RequestedSecret, SecretLeaseContainer.LeaseRenewalScheduler> renewals = (Map<RequestedSecret, SecretLeaseContainer.LeaseRenewalScheduler>) ReflectionTestUtils.getField(this.secretLeaseContainer, "renewals");
+
+		SecretLeaseContainer.LeaseRenewalScheduler scheduler = renewals.get(secret);
+		Lease lease = scheduler.getLease();
+		ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+		this.secretLeaseContainer.restartSecrets();
+
+		scheduler.associateLease(lease);
+
+		verify(this.taskScheduler, times(2)).schedule(captor.capture(), any(Trigger.class));
+		captor.getAllValues().get(1).run(); // new lease run
+		captor.getAllValues().get(0).run(); // old lease run
+
+		// one from restartSecrets, the second from the scheduler detecting expiry
+		verify(this.leaseListenerAdapter, times(2)).onLeaseEvent(any(SecretLeaseRotatedEvent.class));
+		verify(this.leaseListenerAdapter).onLeaseEvent(any(SecretLeaseExpiredEvent.class));
+		verify(this.leaseListenerAdapter, never()).onLeaseEvent(any(SecretLeaseErrorEvent.class));
+	}
+
 	@Test
 	void scheduleRenewalShouldApplyExpiryThreshold() {