Jetty 12.0.x object identity ee8 and ee9

jetty-core/jetty-session/src/main/java/org/eclipse/jetty/session/AbstractSessionManager.java

@@ -1425,7 +1425,8 @@ public HttpCookie getSessionCookie(ManagedSession session, boolean requestIsSecu
 
     /**
      * StreamWrapper to intercept commit and complete events to ensure
-     * session handling happens in context, with request available.
+     * session handling happens in context, with the request available.
+     * This implementation assumes that a request only has a single session.
      */
     private class SessionStreamWrapper extends HttpStream.Wrapper
     {

jetty-ee9/jetty-ee9-nested/src/main/java/org/eclipse/jetty/ee9/nested/ContextHandler.java

@@ -2463,6 +2463,8 @@ public static class CoreContextRequest extends ContextRequest
         private final HttpChannel _httpChannel;
         private SessionManager _sessionManager;
         private ManagedSession _managedSession;
+        private List<ManagedSession> _managedSessions;
+
         AbstractSessionManager.RequestedSession _requestedSession;
 
         protected CoreContextRequest(org.eclipse.jetty.server.Request wrapped,
@@ -2501,20 +2503,123 @@ public ManagedSession getManagedSession()
             return _managedSession;
         }
 
+        /**
+         * Retrieve an existing session, if one exists, for a given SessionManager. A
+         * session belongs to a single SessionManager, and a context can only have a single
+         * SessionManager. Thus, calling this method is equivalent to asking
+         * "Does a ManagedSession already exist for the given context?".
+         *
+         * @param manager the SessionManager that should be associated with a ManagedSession
+         * @return the ManagedSession that already exists in the given context and is managed
+         * by the given SessionManager.
+         */
+        public ManagedSession getManagedSession(SessionManager manager)
+        {
+            if (_managedSessions == null)
+                return null;
+
+            for (ManagedSession s : _managedSessions)
+            {
+                if (manager == s.getSessionManager())
+                {
+                   if (s.isValid())
+                       return s;
+                }
+            }
+            return null;
+        }
+
         public void setManagedSession(ManagedSession managedSession)
         {
             _managedSession = managedSession;
+            addManagedSession(managedSession);
+        }
+
+        /**
+         * Add a session to the list of sessions maintained by this request.
+         * A session will be added whenever a request visits a new context
+         * that already has a session associated with it, or one is created
+         * during the dispatch.
+         *
+         * @param managedSession the session to add
+         */
+        private void addManagedSession(ManagedSession managedSession)
+        {
+            if (managedSession == null)
+                return;
+            if (_managedSessions == null)
+                _managedSessions = new ArrayList<>();
+            if (!_managedSessions.contains(managedSession))
+                _managedSessions.add(managedSession);
         }
 
         public SessionManager getSessionManager()
         {
             return _sessionManager;
         }
 
+        /**
+         * Remember the session that was extracted from the id in the request
+         *
+         * @param requestedSession info about the session matching the id in the request
+         */
         public void setRequestedSession(AbstractSessionManager.RequestedSession requestedSession)
         {
             _requestedSession = requestedSession;
-            _managedSession = requestedSession.session();
+        }
+
+        /**
+         * Release each of the sessions as the request is now complete
+         */
+        public void completeSessions()
+        {
+            if (_managedSessions != null)
+            {
+                for (ManagedSession s : _managedSessions)
+                {
+
+                    if (s.getSessionManager() == null)
+                        continue; //TODO log it
+                    s.getSessionManager().getContext().run(() -> completeSession(s), this);
+                }
+            }
+        }
+
+        /**
+         * Ensure that each session is committed - ie written out to storage if necessary -
+         * because the response is about to be returned to the client.
+         */
+        public void commitSessions()
+        {
+            if (_managedSessions != null)
+            {
+                for (ManagedSession s : _managedSessions)
+                {
+                    if (s.getSessionManager() == null)
+                        continue; //TODO log it
+                    s.getSessionManager().getContext().run(() -> commitSession(s), this);
+                }
+            }
+        }
+
+        private void commitSession(ManagedSession session)
+        {
+            if (session == null)
+                return;
+            SessionManager manager = session.getSessionManager();
+            if (manager == null)
+                return;
+            manager.commit(session);
+        }
+
+        private void completeSession(ManagedSession session)
+        {
+            if (session == null)
+                return;
+            SessionManager manager = session.getSessionManager();
+            if (manager == null)
+                return;
+            manager.complete(session);
         }
 
         public AbstractSessionManager.RequestedSession getRequestedSession()

jetty-ee9/jetty-ee9-nested/src/main/java/org/eclipse/jetty/ee9/nested/CrossContextDispatcher.java

@@ -248,6 +248,13 @@ public void forward(ServletRequest servletRequest, ServletResponse response) thr
         org.eclipse.jetty.server.Response coreResponse = coreContextRequest.getHttpChannel().getCoreResponse();
         baseResponse.resetForForward();
 
+        //if forwarding to the same environment we must mutate this request for Object wrapper identity
+        if (_targetContext.getTargetContext().getContextHandler() instanceof ContextHandler.CoreContextHandler coreContextHandler)
+        {
+            new Dispatcher(coreContextHandler.getContextHandler(), _uri, _decodedPathInContext).forward(httpServletRequest, httpServletResponse, DispatcherType.FORWARD);
+            return;
+        }
+
         ForwardRequest forwardRequest = new ForwardRequest(coreContextRequest, baseRequest, httpServletRequest);
         ServletCoreResponse servletCoreResponse = new ServletCoreResponse(forwardRequest, httpServletResponse, baseResponse, coreResponse, false);
 
@@ -289,6 +296,13 @@ public void include(ServletRequest servletRequest, ServletResponse servletRespon
         ContextHandler.CoreContextRequest coreContextRequest = baseRequest.getCoreRequest();
         org.eclipse.jetty.server.Response coreResponse = coreContextRequest.getHttpChannel().getCoreResponse();
 
+        //if including to the same environment we must mutate this request for Object wrapper identity
+        if (_targetContext.getTargetContext().getContextHandler() instanceof ContextHandler.CoreContextHandler coreContextHandler)
+        {
+            new Dispatcher(coreContextHandler.getContextHandler(), _uri, _decodedPathInContext).include(httpServletRequest, httpServletResponse);
+            return;
+        }
+
         IncludeRequest includeRequest = new IncludeRequest(coreContextRequest, baseRequest, httpServletRequest);
         IncludeResponse includeResponse = new IncludeResponse(includeRequest, httpServletResponse, baseResponse, coreResponse);
 

jetty-ee9/jetty-ee9-nested/src/main/java/org/eclipse/jetty/ee9/nested/SessionHandler.java

@@ -14,6 +14,7 @@
 package org.eclipse.jetty.ee9.nested;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.Enumeration;
@@ -41,6 +42,8 @@
 import jakarta.servlet.http.HttpSessionIdListener;
 import jakarta.servlet.http.HttpSessionListener;
 import org.eclipse.jetty.http.HttpCookie;
+import org.eclipse.jetty.http.MetaData;
+import org.eclipse.jetty.server.HttpStream;
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.Session;
 import org.eclipse.jetty.session.AbstractSessionManager;
@@ -49,6 +52,7 @@
 import org.eclipse.jetty.session.SessionConfig;
 import org.eclipse.jetty.session.SessionIdManager;
 import org.eclipse.jetty.session.SessionManager;
+import org.eclipse.jetty.util.Callback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -66,6 +70,56 @@ public class SessionHandler extends ScopedHandler implements SessionConfig.Mutab
 
     private ContextHandler _contextHandler;
 
+    /**
+     * Wrapper to ensure the correct lifecycle for all sessions in all
+     * contexts to which the request has been dispatched. When a response
+     * is about to be sent back to the client, all of the sessions are
+     * given the opportunity to persist themselves. When the response is
+     * finished being handled, then the sessions have their reference
+     * counts decremented, potentially leading to eviction from their cache,
+     * as appropriate to the configuration of that cache.
+     *
+     * This wrapper replaces the AbstractSessionManager.SessionStreamWrapper.
+     */
+    private static class SessionStreamWrapper extends HttpStream.Wrapper
+    {
+        private final ContextHandler.CoreContextRequest _request;
+        private final org.eclipse.jetty.server.Context _context;
+
+        public SessionStreamWrapper(HttpStream wrapped, ContextHandler.CoreContextRequest request)
+        {
+            super(wrapped);
+            _request = request;
+            _context = _request.getContext();
+        }
+
+        @Override
+        public void failed(Throwable x)
+        {
+             _request.completeSessions();
+            super.failed(x);
+        }
+
+        @Override
+        public void send(MetaData.Request metadataRequest, MetaData.Response metadataResponse, boolean last, ByteBuffer content, Callback callback)
+        {
+            if (metadataResponse != null)
+            {
+                // Write out all sessions
+                _request.commitSessions();
+            }
+            super.send(metadataRequest, metadataResponse, last, content, callback);
+        }
+
+        @Override
+        public void succeeded()
+        {
+            // Decrement usage count on all sessions
+             _request.completeSessions();
+            super.succeeded();
+        }
+    }
+
     public SessionHandler()
     {
         setSessionTrackingModes(DEFAULT_SESSION_TRACKING_MODES);
@@ -445,29 +499,100 @@ public void setUsingUriParameters(boolean value)
     @Override
     public void doScope(String target, Request baseRequest, HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException
     {
-        if (baseRequest.getDispatcherType() == DispatcherType.REQUEST)
-        {
-            ContextHandler.CoreContextRequest coreRequest = baseRequest.getHttpChannel().getCoreRequest();
+        ContextHandler.CoreContextRequest coreRequest = baseRequest.getHttpChannel().getCoreRequest();
+        SessionManager oldSessionManager = null;
+        ManagedSession oldSession = null;
+        AbstractSessionManager.RequestedSession oldRequestedSession = null;
 
-            // TODO should we use the Stream wrapper? Could use the HttpChannel#onCompleted mechanism instead.
-            _sessionManager.addSessionStreamWrapper(coreRequest);
+        ManagedSession currentSession = null;
+        AbstractSessionManager.RequestedSession currentRequestedSession = null;
 
-            // find and set the session if one exists
-            AbstractSessionManager.RequestedSession requestedSession = _sessionManager.resolveRequestedSessionId(coreRequest);
-
-            coreRequest.setSessionManager(_sessionManager);
-            coreRequest.setRequestedSession(requestedSession);
+        try
+        {
+            switch (baseRequest.getDispatcherType())
+            {
+                case REQUEST:
+                {
+                    _sessionManager.addSessionStreamWrapper(coreRequest);
+                    // find and set the session if one exists, along with an appropriate session manager
+                    currentRequestedSession = _sessionManager.resolveRequestedSessionId(coreRequest);
+                    coreRequest.setSessionManager(_sessionManager);
+                    coreRequest.setRequestedSession(currentRequestedSession);
+                    if (currentRequestedSession != null)
+                        coreRequest.setManagedSession(currentRequestedSession.session());
+                    break;
+                }
+                case ASYNC:
+                case ERROR:
+                case FORWARD:
+                case INCLUDE:
+                {
+                    oldSessionManager = coreRequest.getSessionManager();
+                    oldSession = coreRequest.getManagedSession();
+                    oldRequestedSession = coreRequest.getRequestedSession();
+
+                    //check if we have changed contexts during the dispatch
+                    if (oldSessionManager != _sessionManager)
+                    {
+                        _sessionManager.addSessionStreamWrapper(coreRequest);
+
+                        //find any existing session for this context that has already been accessed
+                        currentSession = coreRequest.getManagedSession(_sessionManager);
+
+                        if (currentSession == null)
+                        {
+                            //session for this context has not been already loaded, try getting it
+                            coreRequest.setManagedSession(null);
+                            currentRequestedSession = _sessionManager.resolveRequestedSessionId(coreRequest);
+                            currentSession = currentRequestedSession.session();
+                        }
+                        else
+                            currentRequestedSession = new AbstractSessionManager.RequestedSession(currentSession, currentSession.getId(), false /*TODO!!!*/);
+
+                        coreRequest.setManagedSession(currentSession);
+                        coreRequest.setRequestedSession(currentRequestedSession);
+                        coreRequest.setSessionManager(_sessionManager);
+                    }
+                    break;
+                }
+                default:
+                    break;
+            }
 
-            HttpCookie cookie = _sessionManager.access(requestedSession.session(), coreRequest.getConnectionMetaData().isSecure());
+            //first time the request has entered this context, or we have changed context
+            if ((currentSession != null) && (oldSessionManager != _sessionManager))
+            {
+                HttpCookie cookie = _sessionManager.access(currentRequestedSession.session(), coreRequest.getConnectionMetaData().isSecure());
 
-            // Handle changed ID or max-age refresh, but only if this is not a redispatched request
-            if (cookie != null)
+                // Handle changed ID or max-age refresh, but only if this is not a redispatched request
+                if ((cookie != null) &&
+                    (request.getDispatcherType() == DispatcherType.ASYNC ||
+                        request.getDispatcherType() == DispatcherType.REQUEST))
+                    baseRequest.getResponse().replaceCookie(cookie);
                 baseRequest.getResponse().replaceCookie(cookie);
-        }
+            }
+
+            if (LOG.isDebugEnabled())
+                LOG.debug("sessionHandler={} session={}", this, currentSession);
 
-        // TODO other dispatch type????
+            nextScope(target, baseRequest, request, response);
+        }
+        finally
+        {
+            if (LOG.isDebugEnabled())
+                LOG.debug("Leaving scope {} dispatch={}, async={}, session={}, oldsession={}, oldsessionhandler={}",
+                    this, baseRequest.getDispatcherType(), baseRequest.isAsyncStarted(), baseRequest.getSession(false),
+                    oldSession, oldSessionManager);
 
-        nextScope(target, baseRequest, request, response);
+            // revert the session handler to the previous, unless it was null, in which case remember it as
+            // the first session handler encountered.
+            if (oldSessionManager != null && oldSessionManager != _sessionManager)
+            {
+                coreRequest.setSessionManager(oldSessionManager);
+                coreRequest.setManagedSession(oldSession);
+                coreRequest.setRequestedSession(oldRequestedSession);
+            }
+        }
     }
 
     @Override
@@ -595,7 +720,8 @@ public Server getServer()
         @Override
         protected void addSessionStreamWrapper(org.eclipse.jetty.server.Request request)
         {
-            super.addSessionStreamWrapper(request);
+            final ContextHandler.CoreContextRequest coreRequest = (ContextHandler.CoreContextRequest)request;
+            coreRequest.addHttpStreamWrapper(s -> new SessionStreamWrapper(s, coreRequest));
         }
 
         @Override