Fix fullstack render server context (#2139)

* Fix fullstack render server context

* only set the server context while polling futures

---------

Co-authored-by: Evan Almloff <evanalmloff@gmail.com>

packages/fullstack/src/axum_adapter.rs

@@ -60,8 +60,6 @@ use axum::{
     extract::State,
     http::{Request, Response, StatusCode},
     response::IntoResponse,
-    routing::{get, post},
-    Router,
 };
 use dioxus_lib::prelude::VirtualDom;
 use futures_util::Future;
@@ -69,9 +67,7 @@ use http::header::*;
 
 use std::sync::Arc;
 
-use crate::{
-    prelude::*, render::SSRState, serve_config::ServeConfig, server_context::DioxusServerContext,
-};
+use crate::prelude::*;
 
 /// A extension trait with utilities for integrating Dioxus with your Axum router.
 pub trait DioxusRouterExt<S> {
@@ -508,8 +504,8 @@ async fn handle_server_fns_inner(
                 .unwrap_or(false);
             let referrer = req.headers().get(REFERER).cloned();
 
-            // actually run the server fn
-            let mut res = service.run(req).await;
+            // actually run the server fn (which may use the server context)
+            let mut res = ProvideServerContext::new(service.run(req), server_context.clone()).await;
 
             // it it accepts text/html (i.e., is a plain form post) and doesn't already have a
             // Location set, then redirect to Referer

packages/fullstack/src/render.rs

@@ -1,9 +1,7 @@
 //! A shared pool of renderers for efficient server side rendering.
-use crate::render::dioxus_core::NoOpMutations;
-use crate::server_context::SERVER_CONTEXT;
-use dioxus_lib::prelude::VirtualDom;
+use crate::{render::dioxus_core::NoOpMutations, server_context::with_server_context};
 use dioxus_ssr::{
-    incremental::{IncrementalRendererConfig, RenderFreshness, WrapBody},
+    incremental::{RenderFreshness, WrapBody},
     Renderer,
 };
 use std::future::Future;
@@ -53,7 +51,7 @@ impl SsrRendererPool {
         };
         match self {
             Self::Renderer(pool) => {
-                let server_context = Box::new(server_context.clone());
+                let server_context = server_context.clone();
                 let mut renderer = pool.write().unwrap().pop().unwrap_or_else(pre_renderer);
 
                 let (tx, rx) = tokio::sync::oneshot::channel();
@@ -61,15 +59,13 @@ impl SsrRendererPool {
                 spawn_platform(move || async move {
                     let mut vdom = virtual_dom_factory();
                     let mut to = WriteBuffer { buffer: Vec::new() };
-                    // before polling the future, we need to set the context
-                    let prev_context = SERVER_CONTEXT.with(|ctx| ctx.replace(server_context));
                     // poll the future, which may call server_context()
                     tracing::info!("Rebuilding vdom");
-                    block_in_place(|| vdom.rebuild(&mut NoOpMutations));
-                    vdom.wait_for_suspense().await;
+                    with_server_context(server_context.clone(), || {
+                        block_in_place(|| vdom.rebuild(&mut NoOpMutations));
+                    });
+                    ProvideServerContext::new(vdom.wait_for_suspense(), server_context).await;
                     tracing::info!("Suspense resolved");
-                    // after polling the future, we need to restore the context
-                    SERVER_CONTEXT.with(|ctx| ctx.replace(prev_context));
 
                     if let Err(err) = wrapper.render_before_body(&mut *to) {
                         let _ = tx.send(Err(err));
@@ -120,16 +116,17 @@ impl SsrRendererPool {
                             &mut *to,
                             |vdom| {
                                 Box::pin(async move {
-                                    // before polling the future, we need to set the context
-                                    let prev_context = SERVER_CONTEXT
-                                        .with(|ctx| ctx.replace(Box::new(server_context)));
                                     // poll the future, which may call server_context()
                                     tracing::info!("Rebuilding vdom");
-                                    block_in_place(|| vdom.rebuild(&mut NoOpMutations));
-                                    vdom.wait_for_suspense().await;
+                                    with_server_context(server_context.clone(), || {
+                                        block_in_place(|| vdom.rebuild(&mut NoOpMutations));
+                                    });
+                                    ProvideServerContext::new(
+                                        vdom.wait_for_suspense(),
+                                        server_context,
+                                    )
+                                    .await;
                                     tracing::info!("Suspense resolved");
-                                    // after polling the future, we need to restore the context
-                                    SERVER_CONTEXT.with(|ctx| ctx.replace(prev_context));
                                 })
                             },
                             &wrapper,

packages/fullstack/src/server_context.rs

@@ -35,7 +35,7 @@ impl Default for DioxusServerContext {
 mod server_fn_impl {
     use super::*;
     use std::sync::LockResult;
-    use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
+    use std::sync::{PoisonError, RwLockReadGuard, RwLockWriteGuard};
 
     use anymap::{any::Any, Map};
     type SendSyncAnyMap = Map<dyn Any + Send + Sync + 'static>;
@@ -159,22 +159,20 @@ pub async fn extract<E: FromServerContext<I>, I>() -> Result<E, E::Rejection> {
     E::from_request(&server_context()).await
 }
 
-pub(crate) fn with_server_context<O>(
-    context: Box<DioxusServerContext>,
-    f: impl FnOnce() -> O,
-) -> (O, Box<DioxusServerContext>) {
+pub(crate) fn with_server_context<O>(context: DioxusServerContext, f: impl FnOnce() -> O) -> O {
     // before polling the future, we need to set the context
-    let prev_context = SERVER_CONTEXT.with(|ctx| ctx.replace(context));
+    let prev_context = SERVER_CONTEXT.with(|ctx| ctx.replace(Box::new(context)));
     // poll the future, which may call server_context()
     let result = f();
     // after polling the future, we need to restore the context
-    (result, SERVER_CONTEXT.with(|ctx| ctx.replace(prev_context)))
+    SERVER_CONTEXT.with(|ctx| ctx.replace(prev_context));
+    result
 }
 
 /// A future that provides the server context to the inner future
 #[pin_project::pin_project]
 pub struct ProvideServerContext<F: std::future::Future> {
-    context: Option<Box<DioxusServerContext>>,
+    context: Option<DioxusServerContext>,
     #[pin]
     f: F,
 }
@@ -183,7 +181,7 @@ impl<F: std::future::Future> ProvideServerContext<F> {
     /// Create a new future that provides the server context to the inner future
     pub fn new(f: F, context: DioxusServerContext) -> Self {
         Self {
-            context: Some(Box::new(context)),
+            context: Some(context),
             f,
         }
     }
@@ -198,7 +196,7 @@ impl<F: std::future::Future> std::future::Future for ProvideServerContext<F> {
     ) -> std::task::Poll<Self::Output> {
         let this = self.project();
         let context = this.context.take().unwrap();
-        let (result, context) = with_server_context(context, || this.f.poll(cx));
+        let result = with_server_context(context.clone(), || this.f.poll(cx));
         *this.context = Some(context);
         result
     }