Implement ability to reuse session for initial render (#3679)

panel/command/serve.py

@@ -209,6 +209,10 @@ class Serve(_BkServe):
             help    = "The endpoint for the liveness API.",
             default = "liveness"
         )),
+        ('--reuse-sessions', dict(
+            action  = 'store_true',
+            help    = "Whether to reuse sessions when serving the initial request.",
+        )),
     )
 
     # Supported file extensions
@@ -277,6 +281,7 @@ def customize_kwargs(self, args, server_kwargs):
             raise ValueError("rest-provider %r not recognized." % args.rest_provider)
 
         config.autoreload = args.autoreload
+        config.reuse_sessions = args.reuse_sessions
 
         if config.autoreload:
             for f in files:

panel/config.py

@@ -151,6 +151,19 @@ class _config(_base_config):
         'pyinstrument', 'snakeviz', 'memray'], doc="""
         The profiler engine to enable.""")
 
+    reuse_sessions = param.Boolean(default=False, doc="""
+        Whether to reuse a session for the initial request to speed up
+        the initial page render. Note that if the initial page differs
+        between sessions, e.g. because it uses query parameters to modify
+        the rendered content, then this option will result in the wrong
+        content being rendered. Define a session_key_func to ensure that
+        reused sessions are only reused when appropriate.""")
+
+    session_key_func = param.Callable(default=None, doc="""
+        Used in conjunction with the reuse_sessions option, the
+        session_key_func is given a tornado.httputil.HTTPServerRequest
+        and should return a key that uniquely captures a session.""")
+
     safe_embed = param.Boolean(default=False, doc="""
         Ensure all bokeh property changes trigger events which are
         embedded. Useful when only partial updates are made in an

panel/io/server.py

@@ -54,6 +54,9 @@
 )
 from bokeh.server.views.doc_handler import DocHandler as BkDocHandler
 from bokeh.server.views.static_handler import StaticHandler
+from bokeh.util.token import (
+    generate_jwt_token, generate_session_id, get_token_payload,
+)
 # Tornado imports
 from tornado.ioloop import IOLoop
 from tornado.web import (
@@ -194,9 +197,12 @@ def _initialize_session_info(session_context: 'BokehSessionContext'):
 #---------------------------------------------------------------------
 
 def server_html_page_for_session(
-    session: 'ServerSession', resources: 'Resources', title: str,
+    session: 'ServerSession',
+    resources: 'Resources',
+    title: str,
+    token: str | None = None,
     template: str | Template = BASE_TEMPLATE,
-    template_variables: Optional[Dict[str, Any]] = None
+    template_variables: Optional[Dict[str, Any]] = None,
 ) -> str:
 
     # ALERT: Replace with better approach before Bokeh 3.x compatible release
@@ -214,7 +220,7 @@ def server_html_page_for_session(
         patch_model_css(root, dist_url=dist_url)
 
     render_item = RenderItem(
-        token = session.token,
+        token = token or session.token,
         roots = session.document.roots,
         use_for_title = False,
     )
@@ -317,7 +323,6 @@ def initialize_document(self, doc):
 
 bokeh.command.util.Application = Application # type: ignore
 
-
 class SessionPrefixHandler:
 
     @contextmanager
@@ -347,10 +352,55 @@ def _session_prefix(self):
 # Patch Bokeh DocHandler URL
 class DocHandler(BkDocHandler, SessionPrefixHandler):
 
+    _session_key_funcs = {}
+
+    @authenticated
+    async def get_session(self):
+        from ..config import config
+        path = self.request.path
+        session = None
+        if config.reuse_sessions and path in self._session_key_funcs:
+            key = self._session_key_funcs[path](self.request)
+            session = state._sessions.get(key)
+        if session is None:
+            session = await super().get_session()
+            with set_curdoc(session.document):
+                if config.reuse_sessions:
+                    key_func = config.session_key_func or (lambda r: path)
+                    if key_func:
+                        self._session_key_funcs[path] = key_func
+                        key = key_func(self.request)
+                    else:
+                        key = path
+                    state._sessions[key] = session
+                    session.block_expiration()
+        return session
+
     @authenticated
     async def get(self, *args, **kwargs):
+        app = self.application
         with self._session_prefix():
+            key_func = self._session_key_funcs.get(self.request.path)
+            if key_func:
+                old_request = key_func(self.request) in state._sessions
+            else:
+                old_request = False
             session = await self.get_session()
+            if old_request and state._sessions.get(key_func(self.request)) is session:
+                session_id = generate_session_id(
+                    secret_key=self.application.secret_key,
+                    signed=self.application.sign_sessions
+                )
+                payload = get_token_payload(session.token)
+                token = generate_jwt_token(
+                    session_id,
+                    secret_key=app.secret_key,
+                    signed=app.sign_sessions,
+                    expiration=app.session_token_expiration,
+                    extra_payload=payload
+                )
+            else:
+                token = session.token
             logger.info(LOG_SESSION_CREATED, id(session.document))
             with set_curdoc(session.document):
                 if config.authorize_callback and not config.authorize_callback(state.user_info):
@@ -370,8 +420,8 @@ async def get(self, *args, **kwargs):
                     resources = Resources.from_bokeh(self.application.resources())
                     page = server_html_page_for_session(
                         session, resources=resources, title=session.document.title,
-                        template=session.document.template,
-                        template_variables=session.document.template_variables
+                        token=token, template=session.document.template,
+                        template_variables=session.document.template_variables,
                     )
         self.set_header("Content-Type", 'text/html')
         self.write(page)

panel/io/state.py

@@ -207,6 +207,9 @@ class _state(param.Parameterized):
     # Locks
     _cache_locks: ClassVar[Dict[str, threading.Lock]] = {'main': threading.Lock()}
 
+    # Sessions
+    _sessions = {}
+
     def __repr__(self) -> str:
         server_info = []
         for server, panel, docs in self._servers.values():
@@ -731,6 +734,7 @@ def reset(self):
         if self._thread_pool is not None:
             self._thread_pool.shutdown(wait=False)
             self._thread_pool = None
+        self._sessions.clear()
 
     def schedule_task(
         self, name: str, callback: Callable[[], None], at: Tat =None,

panel/tests/conftest.py

@@ -359,6 +359,16 @@ def threads():
     finally:
         config.nthreads = None
 
+@pytest.fixture
+def reuse_sessions():
+    config.reuse_sessions = True
+    try:
+        yield
+    finally:
+        config.reuse_sessions = False
+        config.session_key_func = None
+        state._sessions.clear()
+
 @pytest.fixture
 def nothreads():
     yield

panel/tests/test_server.py

@@ -456,6 +456,61 @@ def app():
 
     server.stop()
 
+@pytest.mark.xdist_group(name="server")
+def test_server_reuse_sessions(port, reuse_sessions):
+    def app(counts=[0]):
+        content = f'# Count {counts[0]}'
+        counts[0] += 1
+        return content
+
+    serve(app, port=port, threaded=True, show=False)
+
+    # Wait for server to start
+    time.sleep(1)
+
+    r1 = requests.get(f"http://localhost:{port}/")
+    r2 = requests.get(f"http://localhost:{port}/")
+
+    assert len(state._sessions) == 1
+    assert '/' in state._sessions
+
+    session = state._sessions['/']
+
+    assert session.token in r1.content.decode('utf-8')
+    assert session.token not in r2.content.decode('utf-8')
+
+
+@pytest.mark.xdist_group(name="server")
+def test_server_reuse_sessions_with_session_key_func(port, reuse_sessions):
+    config.session_key_func = lambda r: (r.path, r.arguments.get('arg', [''])[0])
+    def app(counts=[0]):
+        if 'arg' in state.session_args:
+            title = state.session_args['arg'][0].decode('utf-8')
+        else:
+            title = 'Empty'
+        content = f"# Count {counts[0]}"
+        tmpl = BootstrapTemplate(title=title)
+        tmpl.main.append(content)
+        counts[0] += 1
+        return tmpl
+
+    serve(app, port=port, threaded=True, show=False)
+
+    # Wait for server to start
+    time.sleep(1)
+
+    r1 = requests.get(f"http://localhost:{port}/?arg=foo")
+    r2 = requests.get(f"http://localhost:{port}/?arg=bar")
+
+    assert len(state._sessions) == 2
+    assert ('/', b'foo') in state._sessions
+    assert ('/', b'bar') in state._sessions
+
+    session1, session2 = state._sessions.values()
+    assert session1.token in r1.content.decode('utf-8')
+    assert session2.token in r2.content.decode('utf-8')
+
+
 @pytest.mark.xdist_group(name="server")
 def test_show_server_info(html_server_session, markdown_server_session):
     *_, html_port = html_server_session

panel/tests/ui/io/test_server.py

@@ -0,0 +1,53 @@
+import time
+
+import pytest
+
+pytestmark = pytest.mark.ui
+
+from panel import config, state
+from panel.io.server import serve
+from panel.template import BootstrapTemplate
+
+
+def test_server_reuse_sessions(page, port, reuse_sessions):
+    def app(counts=[0]):
+        content = f'# Count {counts[0]}'
+        counts[0] += 1
+        return content
+
+    serve(app, port=port, threaded=True, show=False)
+
+    time.sleep(0.2)
+
+    page.goto(f"http://localhost:{port}")
+
+    assert page.text_content(".markdown h1") == 'Count 0'
+
+    page.goto(f"http://localhost:{port}")
+
+    assert page.text_content(".markdown h1") == 'Count 1'
+
+
+def test_server_reuse_sessions_with_session_key_func(page, port, reuse_sessions):
+    config.session_key_func = lambda r: (r.path, r.arguments.get('arg', [''])[0])
+    def app(counts=[0]):
+        title = state.session_args.get('arg', [b''])[0].decode('utf-8')
+        content = f"# Count {counts[0]}"
+        tmpl = BootstrapTemplate(title=title)
+        tmpl.main.append(content)
+        counts[0] += 1
+        return tmpl
+
+    serve(app, port=port, threaded=True, show=False)
+
+    time.sleep(0.2)
+
+    page.goto(f"http://localhost:{port}/?arg=foo")
+
+    assert page.text_content("title") == 'foo'
+    assert page.text_content(".markdown h1") == 'Count 0'
+
+    page.goto(f"http://localhost:{port}/?arg=bar")
+
+    assert page.text_content("title") == 'bar'
+    assert page.text_content(".markdown h1") == 'Count 1'