From 110b4c4a9440c0cafab8f4b49cda128ff1fce8c6 Mon Sep 17 00:00:00 2001
From: madness-inc <matthias.mueller@appng.org>
Date: Mon, 2 Sep 2019 09:44:11 +0200
Subject: [PATCH] ATS-14

---
 .../session/mongo/MongoPersistentManager.java |  5 +-
 .../tomcat/session/mongo/MongoStore.java      | 96 ++++++++++---------
 2 files changed, 54 insertions(+), 47 deletions(-)

diff --git a/src/main/java/org/appng/tomcat/session/mongo/MongoPersistentManager.java b/src/main/java/org/appng/tomcat/session/mongo/MongoPersistentManager.java
index 7e88917..286affb 100644
--- a/src/main/java/org/appng/tomcat/session/mongo/MongoPersistentManager.java
+++ b/src/main/java/org/appng/tomcat/session/mongo/MongoPersistentManager.java
@@ -65,18 +65,19 @@ public MongoStore getStore() {
 	@Override
 	public Session createSession(String sessionId) {
 		Session session = super.createSession(sessionId);
+		getStore().setSessionActive(session);
 		try {
 			save(session);
 		} catch (IOException e) {
 			log.warn(String.format("Error creating session: %s", session.getIdInternal()));
 			session = null;
+			getStore().setSessionInactive();
 		}
-		getStore().setThreadLocalSession(session);
 		return session;
 	}
 
 	public void afterRequest() {
-		getStore().removeThreadLocalSession();
+		getStore().setSessionInactive();
 	}
 
 	public void save(Session session) throws IOException {
diff --git a/src/main/java/org/appng/tomcat/session/mongo/MongoStore.java b/src/main/java/org/appng/tomcat/session/mongo/MongoStore.java
index 000ea17..d57c890 100644
--- a/src/main/java/org/appng/tomcat/session/mongo/MongoStore.java
+++ b/src/main/java/org/appng/tomcat/session/mongo/MongoStore.java
@@ -57,7 +57,11 @@
  */
 public class MongoStore extends StoreBase {
 
+	/** The currently active session for this thread */
 	protected ThreadLocal<Session> currentSession = new ThreadLocal<>();
+	
+	/** The ID of the currently active session for this thread */
+	protected ThreadLocal<String> currentSessionId = new ThreadLocal<>();
 
 	private final Log log = Utils.getLog(MongoStore.class);
 
@@ -269,51 +273,48 @@ public void processExpires() {
 	 * {@inheritDoc}
 	 */
 	public StandardSession load(String id) throws ClassNotFoundException, IOException {
-
-		StandardSession session = (StandardSession) currentSession.get();
-		if (null != session) {
-			if (session.getIdInternal().equals(id)) {
-				trace("Session from ThreadLocal: %s", id);
-				return session;
+		StandardSession session = null;
+
+		if (null == id) {
+			setSessionInactive();
+		} else if (id.equals(currentSessionId.get())) {			
+			session = (StandardSession) currentSession.get();
+			debug("Returning currently active session %s for thread %s", currentSessionId.get(),
+					Thread.currentThread().getName());
+		} else {
+			long start = System.currentTimeMillis();
+			BasicDBObject sessionQuery = sessionQuery(id);
+			DBObject mongoSession = this.collection.findOne(sessionQuery);
+			if (null == mongoSession) {
+				info("Session %s not found, returning null!", id);
+				setSessionInactive();
 			} else {
-				debug("Session from ThreadLocal differed! Requested: %s, found: %s", id, session.getIdInternal());
-				removeThreadLocalSession();
-				session = null;
-			}
-		}
-
-		long start = System.currentTimeMillis();
-
-		BasicDBObject sessionQuery = sessionQuery(id);
-		DBObject mongoSession = this.collection.findOne(sessionQuery);
-		if (null == mongoSession) {
-			info("Session %s not found, returning null!", id);
-			return null;
-		}
-
-		Container container = manager.getContext();
-		Context context = (Context) container;
-		ClassLoader appContextLoader = context.getLoader().getClassLoader();
-
-		byte[] data = (byte[]) mongoSession.get(sessionDataProperty);
-		if (data != null) {
-			try (ObjectInputStream ois = Utils.getObjectInputStream(appContextLoader,
-					manager.getContext().getServletContext(), data)) {
-
-				session = (StandardSession) this.manager.createEmptySession();
-				session.readObjectData(ois);
-				session.setManager(this.manager);
-
-				debug("Loaded session %s with query %s in %s ms (lastModified %s)", id, sessionQuery,
-						System.currentTimeMillis() - start, new Date(session.getLastAccessedTime()));
-
-				setThreadLocalSession(session);
-			} catch (ReflectiveOperationException roe) {
-				getLog().error(String.format("Error loading session: %s", id), roe);
-				throw roe;
+				Container container = manager.getContext();
+				Context context = (Context) container;
+				ClassLoader appContextLoader = context.getLoader().getClassLoader();
+
+				byte[] data = (byte[]) mongoSession.get(sessionDataProperty);
+				if (data != null) {
+					try (ObjectInputStream ois = Utils.getObjectInputStream(appContextLoader,
+							manager.getContext().getServletContext(), data)) {
+
+						session = (StandardSession) this.manager.createEmptySession();
+						session.readObjectData(ois);
+						session.setManager(this.manager);
+
+						debug("Loaded session %s with query %s in %s ms (lastModified %s)", id, sessionQuery,
+								System.currentTimeMillis() - start, new Date(session.getLastAccessedTime()));
+
+						setSessionActive(session);
+					} catch (ReflectiveOperationException roe) {
+						getLog().error(String.format("Error loading session: %s", id), roe);
+						throw roe;
+					}
+				} else {
+					warn("No data for session: %s, returning null!", id);
+					setSessionInactive();
+				}
 			}
-		} else {
-			warn("No data for session: %s", id);
 		}
 		return session;
 	}
@@ -533,12 +534,17 @@ private Log getLog() {
 		return log;
 	}
 
-	void removeThreadLocalSession() {
+	void setSessionInactive() {
+		String id = currentSessionId.get();
+		debug("Removing session from thread %s: %s", Thread.currentThread().getName(), id);
 		currentSession.remove();
+		currentSessionId.remove();
 	}
 
-	void setThreadLocalSession(Session session) {
+	void setSessionActive(Session session) {
 		currentSession.set(session);
+		currentSessionId.set(session.getId());
+		debug("Setting session for thread %s: %s", Thread.currentThread().getName(), session.getId());
 	}
 
 	// Setters