Merge pull request #4765 from sabify/perf/server_fn-hot-path

Performance optimization for `server_fn`
This commit is contained in:
Greg Johnston
2026-06-12 09:53:55 -04:00
committed by GitHub
10 changed files with 174 additions and 85 deletions

View File

@@ -1,6 +1,6 @@
use super::{Patch, Post, Put};
use crate::{ContentType, Decodes, Encodes, Format, FormatType};
use bytes::Bytes;
use bytes::{BufMut, Bytes, BytesMut};
use serde::{Serialize, de::DeserializeOwned};
/// Serializes and deserializes CBOR with [`ciborium`].
@@ -25,6 +25,10 @@ where
ciborium::ser::into_writer(value, &mut buffer)?;
Ok(Bytes::from(buffer))
}
fn encode_into(value: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
ciborium::ser::into_writer(value, buf.writer())
}
}
impl<T> Decodes<T> for CborEncoding

View File

@@ -1,6 +1,6 @@
use super::{Patch, Post, Put};
use crate::{ContentType, Decodes, Encodes, Format, FormatType};
use bytes::Bytes;
use bytes::{BufMut, Bytes, BytesMut};
use serde::{Serialize, de::DeserializeOwned};
/// Serializes and deserializes JSON with [`serde_json`].
@@ -23,6 +23,10 @@ where
fn encode(output: &T) -> Result<Bytes, Self::Error> {
serde_json::to_vec(output).map(Bytes::from)
}
fn encode_into(output: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
serde_json::to_writer(buf.writer(), output)
}
}
impl<T> Decodes<T> for JsonEncoding

View File

@@ -2,7 +2,7 @@ use crate::{
ContentType, Decodes, Encodes, Format, FormatType,
codec::{Patch, Post, Put},
};
use bytes::Bytes;
use bytes::{BufMut, Bytes, BytesMut};
use serde::{Serialize, de::DeserializeOwned};
/// Serializes and deserializes MessagePack with [`rmp_serde`].
@@ -25,6 +25,10 @@ where
fn encode(value: &T) -> Result<Bytes, Self::Error> {
rmp_serde::to_vec(value).map(Bytes::from)
}
fn encode_into(value: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
rmp_serde::encode::write(&mut buf.writer(), value)
}
}
impl<T> Decodes<T> for MsgPackEncoding

View File

@@ -2,7 +2,7 @@ use crate::{
ContentType, Decodes, Encodes, Format, FormatType,
codec::{Patch, Post, Put},
};
use bytes::Bytes;
use bytes::{Bytes, BytesMut};
use serde::{Serialize, de::DeserializeOwned};
/// A codec for Postcard.
@@ -25,6 +25,14 @@ where
fn encode(value: &T) -> Result<Bytes, Self::Error> {
postcard::to_allocvec(value).map(Bytes::from)
}
fn encode_into(value: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
// `to_extend` appends into the buffer we already own (which holds any
// framing written before it); `mem::take` hands it over by move, so the
// existing contents are not copied.
*buf = postcard::to_extend(value, std::mem::take(buf))?;
Ok(())
}
}
impl<T> Decodes<T> for PostcardEncoding

View File

@@ -2,7 +2,7 @@ use crate::{
ContentType, Decodes, Encodes, Format, FormatType,
codec::{Patch, Post, Put},
};
use bytes::Bytes;
use bytes::{Bytes, BytesMut};
use rkyv::{
Archive, Deserialize, Serialize,
api::high::{HighDeserializer, HighSerializer, HighValidator},
@@ -40,6 +40,14 @@ where
let encoded = rkyv::to_bytes::<rancor::Error>(value)?;
Ok(Bytes::copy_from_slice(encoded.as_ref()))
}
fn encode_into(value: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
// Copy the `AlignedVec` straight into the caller's buffer, skipping the
// intermediate `Bytes` that `encode` allocates and copies into.
let encoded = rkyv::to_bytes::<rancor::Error>(value)?;
buf.extend_from_slice(encoded.as_ref());
Ok(())
}
}
impl<T> Decodes<T> for RkyvEncoding

View File

@@ -3,7 +3,7 @@ use crate::{
codec::{Patch, Post, Put},
error::ServerFnErrorErr,
};
use bytes::Bytes;
use bytes::{BufMut, Bytes, BytesMut};
use serde_lite::{Deserialize, Serialize};
/// Pass arguments and receive responses as JSON in the body of a `POST` request.
@@ -32,6 +32,14 @@ where
.map_err(|e| ServerFnErrorErr::Serialization(e.to_string()))
.map(Bytes::from)
}
fn encode_into(value: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
let intermediate = value
.serialize()
.map_err(|e| ServerFnErrorErr::Serialization(e.to_string()))?;
serde_json::to_writer(buf.writer(), &intermediate)
.map_err(|e| ServerFnErrorErr::Serialization(e.to_string()))
}
}
impl<T> Decodes<T> for SerdeLiteEncoding

View File

@@ -222,33 +222,35 @@ where
CustErr: Display,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"{}",
match self {
ServerFnError::Registration(s) => format!(
"error while trying to register the server function: {s}"
),
ServerFnError::Request(s) => format!(
"error reaching server to call server function: {s}"
),
ServerFnError::ServerError(s) =>
format!("error running server function: {s}"),
ServerFnError::MiddlewareError(s) =>
format!("error running middleware: {s}"),
ServerFnError::Deserialization(s) =>
format!("error deserializing server function results: {s}"),
ServerFnError::Serialization(s) =>
format!("error serializing server function arguments: {s}"),
ServerFnError::Args(s) => format!(
"error deserializing server function arguments: {s}"
),
ServerFnError::MissingArg(s) => format!("missing argument {s}"),
ServerFnError::Response(s) =>
format!("error generating HTTP response: {s}"),
ServerFnError::WrappedServerError(e) => format!("{e}"),
match self {
ServerFnError::Registration(s) => write!(
f,
"error while trying to register the server function: {s}"
),
ServerFnError::Request(s) => {
write!(f, "error reaching server to call server function: {s}")
}
)
ServerFnError::ServerError(s) => {
write!(f, "error running server function: {s}")
}
ServerFnError::MiddlewareError(s) => {
write!(f, "error running middleware: {s}")
}
ServerFnError::Deserialization(s) => {
write!(f, "error deserializing server function results: {s}")
}
ServerFnError::Serialization(s) => {
write!(f, "error serializing server function arguments: {s}")
}
ServerFnError::Args(s) => {
write!(f, "error deserializing server function arguments: {s}")
}
ServerFnError::MissingArg(s) => write!(f, "missing argument {s}"),
ServerFnError::Response(s) => {
write!(f, "error generating HTTP response: {s}")
}
ServerFnError::WrappedServerError(e) => write!(f, "{e}"),
}
}
}
@@ -312,7 +314,7 @@ where
type Error = String;
fn decode(bytes: Bytes) -> Result<ServerFnError<CustErr>, Self::Error> {
let data = String::from_utf8(bytes.to_vec())
let data = String::from_utf8(bytes.into())
.map_err(|err| format!("UTF-8 conversion error: {err}"))?;
data.split_once('|')

View File

@@ -652,17 +652,11 @@ where
let output = server_fn(input.into()).await?;
let output = output.stream.map(|output| {
let result = match output {
Ok(output) => OutputEncoding::encode(&output).map_err(|e| {
OutputStreamError::from_server_fn_error(
ServerFnErrorErr::Serialization(e.to_string()),
)
.ser()
.body
}),
Err(err) => Err(err.ser().body),
};
serialize_result(result)
encode_websocket_frame::<
OutputEncoding,
OutputItem,
OutputStreamError,
>(output)
});
Server::spawn(async move {
@@ -694,21 +688,11 @@ where
pin_mut!(input);
pin_mut!(sink);
while let Some(input) = input.stream.next().await {
let result = match input {
Ok(input) => {
InputEncoding::encode(&input).map_err(|e| {
InputStreamError::from_server_fn_error(
ServerFnErrorErr::Serialization(
e.to_string(),
),
)
.ser()
.body
})
}
Err(err) => Err(err.ser().body),
};
let result = serialize_result(result);
let result = encode_websocket_frame::<
InputEncoding,
InputItem,
InputStreamError,
>(input);
if sink.send(result).await.is_err() {
break;
}
@@ -766,6 +750,36 @@ fn serialize_result(result: Result<Bytes, Bytes>) -> Bytes {
}
}
// Encodes one streamed websocket item directly into a tagged frame.
// Format: [tag: u8][content: Bytes] (see `serialize_result`).
//
// The Ok payload is serialized straight after the tag byte via
// `Encodes::encode_into`, so codecs with an incremental serializer avoid the
// extra allocation and full copy that prepending the tag to an already-encoded
// `Bytes` would otherwise require. The Err path (an already-serialized error
// body, the rare case) is framed through `serialize_result`.
fn encode_websocket_frame<Enc, T, E>(item: Result<T, E>) -> Bytes
where
Enc: Encodes<T>,
E: FromServerFnError,
{
match item {
Ok(value) => {
let mut buf = BytesMut::new();
buf.put_u8(0); // Tag for Ok variant
match Enc::encode_into(&value, &mut buf) {
Ok(()) => buf.freeze(),
Err(e) => serialize_result(Err(E::from_server_fn_error(
ServerFnErrorErr::Serialization(e.to_string()),
)
.ser()
.body)),
}
}
Err(err) => serialize_result(Err(err.ser().body)),
}
}
// Deserializes a Bytes instance back into a Result<Bytes, Bytes>.
fn deserialize_result<E: FromServerFnError>(
bytes: Bytes,
@@ -848,6 +862,21 @@ pub trait Encodes<T>: ContentType + FormatType {
/// Encodes the given value into a bytes.
fn encode(output: &T) -> Result<Bytes, Self::Error>;
/// Encodes the given value by appending it to an existing buffer.
///
/// This lets callers that need to prepend their own framing (for example the
/// websocket `[tag][payload]` format) serialize straight into a buffer they
/// already own, avoiding the extra allocation and full copy that prepending
/// to an already-encoded [`Bytes`] would require.
///
/// The default implementation falls back to [`encode`](Encodes::encode) and
/// copies the result; codecs whose serializer can write incrementally should
/// override it to write directly into `buf`.
fn encode_into(output: &T, buf: &mut BytesMut) -> Result<(), Self::Error> {
buf.extend_from_slice(&Self::encode(output)?);
Ok(())
}
}
/// A trait for types that can be decoded from a bytes for a response body.
@@ -868,14 +897,16 @@ pub use inventory;
macro_rules! initialize_server_fn_map {
($req:ty, $res:ty) => {
std::sync::LazyLock::new(|| {
std::sync::RwLock::new(
$crate::inventory::iter::<ServerFnTraitObj<$req, $res>>
.into_iter()
.map(|obj| {
((obj.path().to_string(), obj.method()), obj.clone())
})
.collect(),
)
let mut map: std::collections::HashMap<
Method,
std::collections::HashMap<String, ServerFnTraitObj<$req, $res>>,
> = std::collections::HashMap::new();
for obj in $crate::inventory::iter::<ServerFnTraitObj<$req, $res>> {
map.entry(obj.method())
.or_default()
.insert(obj.path().to_string(), obj.clone());
}
std::sync::RwLock::new(map)
})
};
}
@@ -987,8 +1018,9 @@ impl<Req, Res> Clone for ServerFnTraitObj<Req, Res> {
}
#[allow(unused)] // used by server integrations
type LazyServerFnMap<Req, Res> =
LazyLock<RwLock<HashMap<(String, Method), ServerFnTraitObj<Req, Res>>>>;
type LazyServerFnMap<Req, Res> = LazyLock<
RwLock<HashMap<Method, HashMap<String, ServerFnTraitObj<Req, Res>>>>,
>;
#[cfg(feature = "ssr")]
impl<Req: 'static, Res: 'static> inventory::Collect
@@ -1071,10 +1103,17 @@ pub mod axum {
>,
> + 'static,
{
REGISTERED_SERVER_FUNCTIONS.write().or_poisoned().insert(
(T::PATH.into(), T::Protocol::METHOD),
ServerFnTraitObj::new::<T>(|req| Box::pin(T::run_on_server(req))),
);
REGISTERED_SERVER_FUNCTIONS
.write()
.or_poisoned()
.entry(T::Protocol::METHOD)
.or_default()
.insert(
T::PATH.into(),
ServerFnTraitObj::new::<T>(|req| {
Box::pin(T::run_on_server(req))
}),
);
}
/// The set of all registered server function paths.
@@ -1083,6 +1122,7 @@ pub mod axum {
.read()
.unwrap()
.values()
.flat_map(|by_path| by_path.values())
.map(|item| (item.path(), item.method()))
.collect();
@@ -1119,11 +1159,11 @@ pub mod axum {
path: &str,
method: Method,
) -> Option<BoxedService<Request<Body>, Response<Body>>> {
let key = (path.into(), method);
REGISTERED_SERVER_FUNCTIONS
.read()
.or_poisoned()
.get(&key)
.get(&method)
.and_then(|by_path| by_path.get(path))
.map(|server_fn| {
let middleware = (server_fn.middleware)();
let mut service = server_fn.clone().boxed();
@@ -1193,10 +1233,17 @@ pub mod actix {
>,
> + 'static,
{
REGISTERED_SERVER_FUNCTIONS.write().or_poisoned().insert(
(T::PATH.into(), T::Protocol::METHOD),
ServerFnTraitObj::new::<T>(|req| Box::pin(T::run_on_server(req))),
);
REGISTERED_SERVER_FUNCTIONS
.write()
.or_poisoned()
.entry(T::Protocol::METHOD)
.or_default()
.insert(
T::PATH.into(),
ServerFnTraitObj::new::<T>(|req| {
Box::pin(T::run_on_server(req))
}),
);
}
/// The set of all registered server function paths.
@@ -1205,6 +1252,7 @@ pub mod actix {
.read()
.unwrap()
.values()
.flat_map(|by_path| by_path.values())
.map(|item| (item.path(), item.method()))
.collect();
@@ -1260,7 +1308,8 @@ pub mod actix {
REGISTERED_SERVER_FUNCTIONS
.read()
.or_poisoned()
.get(&(path.into(), method))
.get(&method)
.and_then(|by_path| by_path.get(path))
.map(|server_fn| {
let middleware = (server_fn.middleware)();
let mut service = server_fn.clone().boxed();

View File

@@ -57,7 +57,7 @@ where
async fn try_into_string(self) -> Result<String, Error> {
let bytes = Req::<Error>::try_into_bytes(self).await?;
String::from_utf8(bytes.to_vec()).map_err(|e| {
String::from_utf8(bytes.into()).map_err(|e| {
ServerFnErrorErr::Deserialization(e.to_string()).into_app_error()
})
}

View File

@@ -18,7 +18,7 @@ use crate::{
};
use bytes::Bytes;
use futures::{
Sink, StreamExt,
Sink,
stream::{self, Stream},
};
use http::{Request, Response};
@@ -47,9 +47,11 @@ where
self,
) -> Result<impl Stream<Item = Result<Bytes, Bytes>> + Send + 'static, Error>
{
Ok(stream::iter(self.into_body())
.ready_chunks(16)
.map(|chunk| Ok(Bytes::from(chunk))))
// The body is a single, contiguous, already-in-memory `Bytes`. Hand it
// back as one stream item instead of iterating it `u8`-by-`u8` and
// re-chunking, which would cost `N` polls plus `N/16` re-allocations and
// a full copy of a buffer we already own.
Ok(stream::once(async move { Ok(self.into_body()) }))
}
fn to_content_type(&self) -> Option<Cow<'_, str>> {