Backport of InMemoryWebSession changes

- hooks to check expired sessions in both create and retrieve.
- maxSessions limit on the total number of sessions.
- getSessions method for management purposes
- removeExpiredSessions public API

Issue: SPR-17020, SPR-16713

Author: rstoyanchev

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

@@ -20,10 +20,11 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.time.ZoneId;
+import java.time.temporal.ChronoUnit;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -43,20 +44,37 @@
  */
 public class InMemoryWebSessionStore implements WebSessionStore {
 
-	/** Minimum period between expiration checks */
-	private static final Duration EXPIRATION_CHECK_PERIOD = Duration.ofSeconds(60);
-
 	private static final IdGenerator idGenerator = new JdkIdGenerator();
 
 
+	private int maxSessions = 10000;
+
 	private Clock clock = Clock.system(ZoneId.of("GMT"));
 
-	private final ConcurrentMap<String, InMemoryWebSession> sessions = new ConcurrentHashMap<>();
+	private final Map<String, InMemoryWebSession> sessions = new ConcurrentHashMap<>();
+
+	private final ExpiredSessionChecker expiredSessionChecker = new ExpiredSessionChecker();
 
-	private volatile Instant nextExpirationCheckTime = Instant.now(this.clock).plus(EXPIRATION_CHECK_PERIOD);
 
-	private final ReentrantLock expirationCheckLock = new ReentrantLock();
+	/**
+	 * Set the maximum number of sessions that can be stored. Once the limit is
+	 * reached, any attempt to store an additional session will result in an
+	 * {@link IllegalStateException}.
+	 * <p>By default set to 10000.
+	 * @param maxSessions the maximum number of sessions
+	 * @since 5.1
+	 */
+	public void setMaxSessions(int maxSessions) {
+		this.maxSessions = maxSessions;
+	}
 
+	/**
+	 * Return the maximum number of sessions that can be stored.
+	 * @since 5.1
+	 */
+	public int getMaxSessions() {
+		return this.maxSessions;
+	}
 
 	/**
 	 * Configure the {@link Clock} to use to set lastAccessTime on every created
@@ -70,8 +88,7 @@ public class InMemoryWebSessionStore implements WebSessionStore {
 	public void setClock(Clock clock) {
 		Assert.notNull(clock, "Clock is required");
 		this.clock = clock;
-		// Force a check when clock changes..
-		this.nextExpirationCheckTime = Instant.now(this.clock);
+		removeExpiredSessions();
 	}
 
 	/**
@@ -81,67 +98,67 @@ public Clock getClock() {
 		return this.clock;
 	}
 
+	/**
+	 * Return the map of sessions with an {@link Collections#unmodifiableMap
+	 * unmodifiable} wrapper. This could be used for management purposes, to
+	 * list active sessions, invalidate expired ones, etc.
+	 * @since 5.1
+	 */
+	public Map<String, InMemoryWebSession> getSessions() {
+		return Collections.unmodifiableMap(this.sessions);
+	}
+
 
 	@Override
 	public Mono<WebSession> createWebSession() {
-		return Mono.fromSupplier(InMemoryWebSession::new);
+		Instant now = this.clock.instant();
+		this.expiredSessionChecker.checkIfNecessary(now);
+		return Mono.fromSupplier(() -> new InMemoryWebSession(now));
 	}
 
 	@Override
 	public Mono<WebSession> retrieveSession(String id) {
-		Instant currentTime = Instant.now(this.clock);
-		if (!this.sessions.isEmpty() && !currentTime.isBefore(this.nextExpirationCheckTime)) {
-			checkExpiredSessions(currentTime);
-		}
-
+		Instant now = this.clock.instant();
+		this.expiredSessionChecker.checkIfNecessary(now);
 		InMemoryWebSession session = this.sessions.get(id);
 		if (session == null) {
 			return Mono.empty();
 		}
-		else if (session.isExpired(currentTime)) {
+		else if (session.isExpired(now)) {
 			this.sessions.remove(id);
 			return Mono.empty();
 		}
 		else {
-			session.updateLastAccessTime(currentTime);
+			session.updateLastAccessTime(now);
 			return Mono.just(session);
 		}
 	}
 
-	private void checkExpiredSessions(Instant currentTime) {
-		if (this.expirationCheckLock.tryLock()) {
-			try {
-				Iterator<InMemoryWebSession> iterator = this.sessions.values().iterator();
-				while (iterator.hasNext()) {
-					InMemoryWebSession session = iterator.next();
-					if (session.isExpired(currentTime)) {
-						iterator.remove();
-						session.invalidate();
-					}
-				}
-			}
-			finally {
-				this.nextExpirationCheckTime = currentTime.plus(EXPIRATION_CHECK_PERIOD);
-				this.expirationCheckLock.unlock();
-			}
-		}
-	}
-
 	@Override
 	public Mono<Void> removeSession(String id) {
 		this.sessions.remove(id);
 		return Mono.empty();
 	}
 
-	public Mono<WebSession> updateLastAccessTime(WebSession webSession) {
+	public Mono<WebSession> updateLastAccessTime(WebSession session) {
 		return Mono.fromSupplier(() -> {
-			Assert.isInstanceOf(InMemoryWebSession.class, webSession);
-			InMemoryWebSession session = (InMemoryWebSession) webSession;
-			session.updateLastAccessTime(Instant.now(getClock()));
+			Assert.isInstanceOf(InMemoryWebSession.class, session);
+			((InMemoryWebSession) session).updateLastAccessTime(this.clock.instant());
 			return session;
 		});
 	}
 
+	/**
+	 * Check for expired sessions and remove them. Typically such checks are
+	 * kicked off lazily during calls to {@link #createWebSession() create} or
+	 * {@link #retrieveSession retrieve}, no less than 60 seconds apart.
+	 * This method can be called to force a check at a specific time.
+	 * @since 5.1
+	 */
+	public void removeExpiredSessions() {
+		this.expiredSessionChecker.removeExpiredSessions(this.clock.instant());
+	}
+
 
 	private class InMemoryWebSession implements WebSession {
 
@@ -157,8 +174,9 @@ private class InMemoryWebSession implements WebSession {
 
 		private final AtomicReference<State> state = new AtomicReference<>(State.NEW);
 
-		public InMemoryWebSession() {
-			this.creationTime = Instant.now(getClock());
+
+		public InMemoryWebSession(Instant creationTime) {
+			this.creationTime = creationTime;
 			this.lastAccessTime = this.creationTime;
 		}
 
@@ -222,6 +240,12 @@ public Mono<Void> invalidate() {
 
 		@Override
 		public Mono<Void> save() {
+			if (sessions.size() >= maxSessions) {
+				expiredSessionChecker.removeExpiredSessions(clock.instant());
+				if (sessions.size() >= maxSessions) {
+					return Mono.error(new IllegalStateException("Max sessions limit reached: " + sessions.size()));
+				}
+			}
 			if (!getAttributes().isEmpty()) {
 				this.state.compareAndSet(State.NEW, State.STARTED);
 			}
@@ -231,14 +255,14 @@ public Mono<Void> save() {
 
 		@Override
 		public boolean isExpired() {
-			return isExpired(Instant.now(getClock()));
+			return isExpired(clock.instant());
 		}
 
-		private boolean isExpired(Instant currentTime) {
+		private boolean isExpired(Instant now) {
 			if (this.state.get().equals(State.EXPIRED)) {
 				return true;
 			}
-			if (checkExpired(currentTime)) {
+			if (checkExpired(now)) {
 				this.state.set(State.EXPIRED);
 				return true;
 			}
@@ -256,6 +280,47 @@ private void updateLastAccessTime(Instant currentTime) {
 	}
 
 
+	private class ExpiredSessionChecker {
+
+		/** Max time between expiration checks. */
+		private static final int CHECK_PERIOD = 60 * 1000;
+
+
+		private final ReentrantLock lock = new ReentrantLock();
+
+		private Instant checkTime = clock.instant().plus(CHECK_PERIOD, ChronoUnit.MILLIS);
+
+
+		public void checkIfNecessary(Instant now) {
+			if (this.checkTime.isBefore(now)) {
+				removeExpiredSessions(now);
+			}
+		}
+
+		public void removeExpiredSessions(Instant now) {
+			if (sessions.isEmpty()) {
+				return;
+			}
+			if (this.lock.tryLock()) {
+				try {
+					Iterator<InMemoryWebSession> iterator = sessions.values().iterator();
+					while (iterator.hasNext()) {
+						InMemoryWebSession session = iterator.next();
+						if (session.isExpired(now)) {
+							iterator.remove();
+							session.invalidate();
+						}
+					}
+				}
+				finally {
+					this.checkTime = now.plus(CHECK_PERIOD, ChronoUnit.MILLIS);
+					this.lock.unlock();
+				}
+			}
+		}
+	}
+
+
 	private enum State { NEW, STARTED, EXPIRED }
 
 }

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

@@ -18,15 +18,20 @@
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Map;
+import java.util.stream.IntStream;
 
 import org.junit.Test;
 
+import org.springframework.beans.DirectFieldAccessor;
 import org.springframework.web.server.WebSession;
 
 import static junit.framework.TestCase.assertSame;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 /**
  * Unit tests for {@link InMemoryWebSessionStore}.
@@ -55,7 +60,7 @@ public void startsSessionImplicitly() {
 	}
 
 	@Test
-	public void retrieveExpiredSession() throws Exception {
+	public void retrieveExpiredSession() {
 		WebSession session = this.store.createWebSession().block();
 		assertNotNull(session);
 		session.getAttributes().put("foo", "bar");
@@ -73,7 +78,7 @@ public void retrieveExpiredSession() throws Exception {
 	}
 
 	@Test
-	public void lastAccessTimeIsUpdatedOnRetrieve() throws Exception {
+	public void lastAccessTimeIsUpdatedOnRetrieve() {
 		WebSession session1 = this.store.createWebSession().block();
 		assertNotNull(session1);
 		String id = session1.getId();
@@ -91,46 +96,45 @@ public void lastAccessTimeIsUpdatedOnRetrieve() throws Exception {
 	}
 
 	@Test
-	public void expirationChecks() throws Exception {
-		// Create 3 sessions
-		WebSession session1 = this.store.createWebSession().block();
-		assertNotNull(session1);
-		session1.start();
-		session1.save().block();
+	public void expirationCheckPeriod() {
 
-		WebSession session2 = this.store.createWebSession().block();
-		assertNotNull(session2);
-		session2.start();
-		session2.save().block();
+		DirectFieldAccessor accessor = new DirectFieldAccessor(this.store);
+		Map<?,?> sessions = (Map<?, ?>) accessor.getPropertyValue("sessions");
+		assertNotNull(sessions);
 
-		WebSession session3 = this.store.createWebSession().block();
-		assertNotNull(session3);
-		session3.start();
-		session3.save().block();
+		// Create 100 sessions
+		IntStream.range(0, 100).forEach(i -> insertSession());
+		assertEquals(100, sessions.size());
 
-		// Fast-forward 31 minutes
-		this.store.setClock(Clock.offset(this.store.getClock(), Duration.ofMinutes(31)));
+		// Force a new clock (31 min later), don't use setter which would clean expired sessions
+		accessor.setPropertyValue("clock", Clock.offset(this.store.getClock(), Duration.ofMinutes(31)));
+		assertEquals(100, sessions.size());
 
-		// Create 2 more sessions
-		WebSession session4 = this.store.createWebSession().block();
-		assertNotNull(session4);
-		session4.start();
-		session4.save().block();
+		// Create 1 more which forces a time-based check (clock moved forward)
+		insertSession();
+		assertEquals(1, sessions.size());
+	}
 
-		WebSession session5 = this.store.createWebSession().block();
-		assertNotNull(session5);
-		session5.start();
-		session5.save().block();
+	@Test
+	public void maxSessions() {
 
-		// Retrieve, forcing cleanup of all expired..
-		assertNull(this.store.retrieveSession(session1.getId()).block());
-		assertNull(this.store.retrieveSession(session2.getId()).block());
-		assertNull(this.store.retrieveSession(session3.getId()).block());
+		IntStream.range(0, 10000).forEach(i -> insertSession());
 
-		assertNotNull(this.store.retrieveSession(session4.getId()).block());
-		assertNotNull(this.store.retrieveSession(session5.getId()).block());
+		try {
+			insertSession();
+			fail();
+		}
+		catch (IllegalStateException ex) {
+			assertEquals("Max sessions limit reached: 10000", ex.getMessage());
+		}
 	}
 
+	private WebSession insertSession() {
+		WebSession session = this.store.createWebSession().block();
+		assertNotNull(session);
+		session.start();
+		session.save().block();
+		return session;
+	}
 
-
-}
+}