From fb396b0448468ce11c0e46821f5531233b53eec6 Mon Sep 17 00:00:00 2001 From: Emil Boman Date: Tue, 26 Mar 2024 15:59:25 +0100 Subject: [PATCH] Fix fullstack render server context (#2139) * Fix fullstack render server context * only set the server context while polling futures --------- Co-authored-by: Evan Almloff --- packages/fullstack/src/axum_adapter.rs | 10 +++---- packages/fullstack/src/render.rs | 33 +++++++++++------------- packages/fullstack/src/server_context.rs | 18 ++++++------- 3 files changed, 26 insertions(+), 35 deletions(-) diff --git a/packages/fullstack/src/axum_adapter.rs b/packages/fullstack/src/axum_adapter.rs index dcafdb01d..cf4a2bdf5 100644 --- a/packages/fullstack/src/axum_adapter.rs +++ b/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 { @@ -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 diff --git a/packages/fullstack/src/render.rs b/packages/fullstack/src/render.rs index 42b6395ee..7a46a0fbd 100644 --- a/packages/fullstack/src/render.rs +++ b/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, diff --git a/packages/fullstack/src/server_context.rs b/packages/fullstack/src/server_context.rs index 7ddc67f3b..51f8b1998 100644 --- a/packages/fullstack/src/server_context.rs +++ b/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; @@ -159,22 +159,20 @@ pub async fn extract, I>() -> Result { E::from_request(&server_context()).await } -pub(crate) fn with_server_context( - context: Box, - f: impl FnOnce() -> O, -) -> (O, Box) { +pub(crate) fn with_server_context(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 { - context: Option>, + context: Option, #[pin] f: F, } @@ -183,7 +181,7 @@ impl ProvideServerContext { /// 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 std::future::Future for ProvideServerContext { ) -> std::task::Poll { 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 }