WebSessionStore updates lastAccessTime on retrieve

Now that WebSessionStore is in charge of expiration checks on retrieve
it makes sense to also update the lastAccessTime on retrieve at the
same time, saving the need to call it after a retrieve.

Issue: SPR-15963

Author: rstoyanchev

spring-web/src/main/java/org/springframework/web/server/session/DefaultWebSessionManager.java

@@ -78,11 +78,9 @@ public WebSessionStore getSessionStore() {
 
 	@Override
 	public Mono<WebSession> getSession(ServerWebExchange exchange) {
-		return Mono.defer(() ->
-				retrieveSession(exchange)
-						.flatMap(this.getSessionStore()::updateLastAccessTime)
-						.switchIfEmpty(this.sessionStore.createWebSession())
-						.doOnNext(session -> exchange.getResponse().beforeCommit(() -> save(exchange, session))));
+		return Mono.defer(() -> retrieveSession(exchange)
+				.switchIfEmpty(this.sessionStore.createWebSession())
+				.doOnNext(session -> exchange.getResponse().beforeCommit(() -> save(exchange, session))));
 	}
 
 	private Mono<WebSession> retrieveSession(ServerWebExchange exchange) {
@@ -93,11 +91,7 @@ private Mono<WebSession> retrieveSession(ServerWebExchange exchange) {
 
 	private Mono<Void> save(ServerWebExchange exchange, WebSession session) {
 		if (session.isExpired()) {
-			return Mono.error(new IllegalStateException(
-					"Sessions are checked for expiration and have their " +
-							"lastAccessTime updated when first accessed during request processing. " +
-							"However this session is expired meaning that maxIdleTime elapsed " +
-							"before the call to session.save()."));
+			return Mono.error(new IllegalStateException("Session='" + session.getId() + "' expired."));
 		}
 
 		if (!session.isStarted()) {

spring-web/src/main/java/org/springframework/web/server/session/InMemoryWebSessionStore.java

@@ -45,7 +45,7 @@ public class InMemoryWebSessionStore implements WebSessionStore {
 
 	private Clock clock = Clock.system(ZoneId.of("GMT"));
 
-	private final Map<String, WebSession> sessions = new ConcurrentHashMap<>();
+	private final Map<String, InMemoryWebSession> sessions = new ConcurrentHashMap<>();
 
 
 	/**
@@ -77,7 +77,7 @@ public Mono<WebSession> createWebSession() {
 
 	@Override
 	public Mono<WebSession> retrieveSession(String id) {
-		WebSession session = this.sessions.get(id);
+		InMemoryWebSession session = this.sessions.get(id);
 		if (session == null) {
 			return Mono.empty();
 		}
@@ -86,6 +86,7 @@ else if (session.isExpired()) {
 			return Mono.empty();
 		}
 		else {
+			session.updateLastAccessTime();
 			return Mono.just(session);
 		}
 	}
@@ -98,27 +99,14 @@ public Mono<Void> removeSession(String id) {
 
 	public Mono<WebSession> updateLastAccessTime(WebSession webSession) {
 		return Mono.fromSupplier(() -> {
+			Assert.isInstanceOf(InMemoryWebSession.class, webSession);
 			InMemoryWebSession session = (InMemoryWebSession) webSession;
-			Instant lastAccessTime = Instant.now(getClock());
-			return new InMemoryWebSession(session, lastAccessTime);
+			session.updateLastAccessTime();
+			return session;
 		});
 	}
 
 
-	// Private methods for InMemoryWebSession
-
-	private Mono<Void> changeSessionId(String oldId, WebSession session) {
-		this.sessions.remove(oldId);
-		this.sessions.put(session.getId(), session);
-		return Mono.empty();
-	}
-
-	private Mono<Void> storeSession(WebSession session) {
-		this.sessions.put(session.getId(), session);
-		return Mono.empty();
-	}
-
-
 	private class InMemoryWebSession implements WebSession {
 
 		private final AtomicReference<String> id;
@@ -127,12 +115,13 @@ private class InMemoryWebSession implements WebSession {
 
 		private final Instant creationTime;
 
-		private final Instant lastAccessTime;
+		private volatile Instant lastAccessTime;
 
 		private volatile Duration maxIdleTime;
 
 		private volatile boolean started;
 
+
 		InMemoryWebSession() {
 			this.id = new AtomicReference<>(String.valueOf(idGenerator.generateId()));
 			this.attributes = new ConcurrentHashMap<>();
@@ -141,14 +130,6 @@ private class InMemoryWebSession implements WebSession {
 			this.maxIdleTime = Duration.ofMinutes(30);
 		}
 
-		InMemoryWebSession(InMemoryWebSession existingSession, Instant lastAccessTime) {
-			this.id = existingSession.id;
-			this.attributes = existingSession.attributes;
-			this.creationTime = existingSession.creationTime;
-			this.lastAccessTime = lastAccessTime;
-			this.maxIdleTime = existingSession.maxIdleTime;
-			this.started = existingSession.isStarted(); // Use method (explicit or implicit start)
-		}
 
 		@Override
 		public String getId() {
@@ -192,22 +173,33 @@ public boolean isStarted() {
 
 		@Override
 		public Mono<Void> changeSessionId() {
-			String oldId = this.id.get();
+			String currentId = this.id.get();
+			if (InMemoryWebSessionStore.this.sessions.remove(currentId) == null) {
+				return Mono.error(new IllegalStateException(
+						"Failed to change session id: " + currentId +
+								" because the Session is no longer present in the store."));
+			}
 			String newId = String.valueOf(idGenerator.generateId());
 			this.id.set(newId);
-			return InMemoryWebSessionStore.this.changeSessionId(oldId, this).doOnError(ex -> this.id.set(oldId));
+			InMemoryWebSessionStore.this.sessions.put(this.getId(), this);
+			return Mono.empty();
 		}
 
 		@Override
 		public Mono<Void> save() {
-			return InMemoryWebSessionStore.this.storeSession(this);
+			InMemoryWebSessionStore.this.sessions.put(this.getId(), this);
+			return Mono.empty();
 		}
 
 		@Override
 		public boolean isExpired() {
 			return (isStarted() && !this.maxIdleTime.isNegative() &&
 					Instant.now(getClock()).minus(this.maxIdleTime).isAfter(this.lastAccessTime));
 		}
+
+		private void updateLastAccessTime() {
+			this.lastAccessTime = Instant.now(getClock());
+		}
 	}
 
 }

spring-web/src/main/java/org/springframework/web/server/session/WebSessionStore.java

@@ -41,7 +41,8 @@ public interface WebSessionStore {
 	/**
 	 * Return the WebSession for the given id.
 	 * <p><strong>Note:</strong> This method should perform an expiration check,
-	 * remove the session if it has expired and return empty.
+	 * and if it has expired remove the session and return empty. This method
+	 * should also update the lastAccessTime of retrieved sessions.
 	 * @param sessionId the session to load
 	 * @return the session, or an empty {@code Mono} .
 	 */

spring-web/src/test/java/org/springframework/web/server/session/InMemoryWebSessionStoreTests.java

@@ -17,6 +17,7 @@
 
 import java.time.Clock;
 import java.time.Duration;
+import java.time.Instant;
 
 import org.junit.Test;
 
@@ -28,36 +29,14 @@
 import static org.junit.Assert.assertTrue;
 
 /**
- * Unit tests.
+ * Unit tests for {@link InMemoryWebSessionStore}.
  * @author Rob Winch
  */
 public class InMemoryWebSessionStoreTests {
 
 	private InMemoryWebSessionStore store = new InMemoryWebSessionStore();
 
 
-	@Test
-	public void constructorWhenImplicitStartCopiedThenCopyIsStarted() {
-		WebSession original = this.store.createWebSession().block();
-		assertNotNull(original);
-		original.getAttributes().put("foo", "bar");
-
-		WebSession copy = this.store.updateLastAccessTime(original).block();
-		assertNotNull(copy);
-		assertTrue(copy.isStarted());
-	}
-
-	@Test
-	public void constructorWhenExplicitStartCopiedThenCopyIsStarted() {
-		WebSession original = this.store.createWebSession().block();
-		assertNotNull(original);
-		original.start();
-
-		WebSession copy = this.store.updateLastAccessTime(original).block();
-		assertNotNull(copy);
-		assertTrue(copy.isStarted());
-	}
-
 	@Test
 	public void startsSessionExplicitly() {
 		WebSession session = this.store.createWebSession().block();
@@ -87,8 +66,27 @@ public void retrieveExpiredSession() throws Exception {
 		assertNotNull(retrieved);
 		assertSame(session, retrieved);
 
+		// Fast-forward 31 minutes
 		this.store.setClock(Clock.offset(this.store.getClock(), Duration.ofMinutes(31)));
 		WebSession retrievedAgain = this.store.retrieveSession(id).block();
 		assertNull(retrievedAgain);
 	}
+
+	@Test
+	public void lastAccessTimeIsUpdatedOnRetrieve() throws Exception {
+		WebSession session1 = this.store.createWebSession().block();
+		assertNotNull(session1);
+		String id = session1.getId();
+		Instant time1 = session1.getLastAccessTime();
+		session1.save();
+
+		// Fast-forward a few seconds
+		this.store.setClock(Clock.offset(this.store.getClock(), Duration.ofSeconds(5)));
+
+		WebSession session2 = this.store.retrieveSession(id).block();
+		assertNotNull(session2);
+		assertSame(session1, session2);
+		Instant time2 = session2.getLastAccessTime();
+		assertTrue(time1.isBefore(time2));
+	}
 }