rpc-binario streaming

oscartbeaumont · oscartbeaumont · commit 07ac2f033e3f · 2024-12-27T14:05:31.000+08:00
diff --git a/crates/binario/src/lib.rs b/crates/binario/src/lib.rs
@@ -1,4 +1,13 @@
 //! rspc-binario: Binario support for rspc
+//!
+//! TODO:
+//!  - Support for streaming the result. Right now we encode into a buffer.
+//!  - `BinarioDeserializeError` should end up as a `ProcedureError::Deserialize` not `ProcedureError::Resolver`
+//!  - Binario needs impl for `()` for procedures with no input.
+//!  - Client integration
+//!  - Cleanup HTTP endpoint on `example-binario`. Maybe don't use HTTP cause Axum's model doesn't mesh with Binario?
+//!  - Maybe actix-web example to show portability. Might be interesting with the fact that Binario depends on `tokio::AsyncRead`.
+//!
 #![forbid(unsafe_code)]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![doc(
@@ -9,7 +18,7 @@
 use std::{error, fmt, marker::PhantomData, pin::Pin};
 
 use binario::{encode, Decode, Encode};
-use futures_util::{stream, Stream};
+use futures_util::{stream, Stream, StreamExt};
 use rspc::{
     middleware::Middleware, DynInput, ProcedureError, ProcedureStream, ResolverInput,
     ResolverOutput,
@@ -46,23 +55,65 @@ impl<T: Decode + Type + Send + 'static> ResolverInput for TypedBinarioInput<T> {
     }
 }
 
-// TODO: Streaming instead of this
+// TODO: This should probs be a stream not a buffer.
+// Binario currently only supports `impl AsyncRead` not `impl Stream`
 pub struct BinarioOutput(pub Vec<u8>);
-pub struct TypedBinarioOutput<T>(pub T);
+pub struct TypedBinarioOutput<T, M>(pub T, pub PhantomData<fn() -> M>);
 
-impl<TError, T: Encode + Type + Send + Sync + 'static> ResolverOutput<TError>
-    for TypedBinarioOutput<T>
+pub(crate) mod sealed {
+    use super::*;
+
+    pub trait ValidBinarioOutput<M>: Send + 'static {
+        type T: Encode + Send + Sync + 'static;
+        fn data_type(types: &mut TypeCollection) -> DataType;
+        fn into_stream(
+            self,
+        ) -> impl Stream<Item = Result<Self::T, ProcedureError>> + Send + 'static;
+    }
+    pub enum ValueMarker {}
+    impl<T: Encode + Type + Send + Sync + 'static> ValidBinarioOutput<ValueMarker> for T {
+        type T = T;
+        fn data_type(types: &mut TypeCollection) -> DataType {
+            T::inline(types, Generics::Definition)
+        }
+        fn into_stream(
+            self,
+        ) -> impl Stream<Item = Result<Self::T, ProcedureError>> + Send + 'static {
+            stream::once(async move { Ok(self) })
+        }
+    }
+    pub enum StreamMarker {}
+    impl<S: Stream + Send + Sync + 'static> ValidBinarioOutput<StreamMarker> for rspc::Stream<S>
+    where
+        S::Item: Encode + Type + Send + Sync + 'static,
+    {
+        type T = S::Item;
+
+        fn data_type(types: &mut TypeCollection) -> DataType {
+            S::Item::inline(types, Generics::Definition)
+        }
+        fn into_stream(
+            self,
+        ) -> impl Stream<Item = Result<Self::T, ProcedureError>> + Send + 'static {
+            self.0.map(|v| Ok(v))
+        }
+    }
+}
+
+impl<TError, M: 'static, T: sealed::ValidBinarioOutput<M>> ResolverOutput<TError>
+    for TypedBinarioOutput<T, M>
 {
     type T = BinarioOutput;
 
     fn data_type(types: &mut TypeCollection) -> DataType {
-        T::inline(types, Generics::Definition)
+        T::data_type(types)
     }
 
     fn into_stream(self) -> impl Stream<Item = Result<Self::T, ProcedureError>> + Send + 'static {
-        stream::once(async move {
+        // TODO: Encoding into a buffer is not how Binario is intended to work but it's how rspc needs it.
+        self.0.into_stream().then(|v| async move {
             let mut buf = Vec::new();
-            encode(&self.0, &mut buf).await.unwrap(); // TODO: Error handling
+            encode(&v?, &mut buf).await.unwrap(); // TODO: Error handling
             Ok(BinarioOutput(buf))
         })
     }
@@ -74,11 +125,11 @@ impl<TError, T: Encode + Type + Send + Sync + 'static> ResolverOutput<TError>
     }
 }
 
-pub fn binario<TError, TCtx, TInput, TResult>() -> Middleware<
+pub fn binario<TError, TCtx, TInput, TResult, M>() -> Middleware<
     TError,
     TCtx,
     TypedBinarioInput<TInput>,
-    TypedBinarioOutput<TResult>,
+    TypedBinarioOutput<TResult, M>,
     TCtx,
     TInput,
     TResult,
@@ -87,7 +138,7 @@ where
     TError: From<DeserializeError> + Send + 'static,
     TCtx: Send + 'static,
     TInput: Decode + Send + 'static,
-    TResult: Encode + Send + Sync + 'static,
+    TResult: sealed::ValidBinarioOutput<M>,
 {
     Middleware::new(
         move |ctx: TCtx, input: TypedBinarioInput<TInput>, next| async move {
@@ -97,7 +148,9 @@ where
             }
             .map_err(DeserializeError)?;
 
-            next.exec(ctx, input).await.map(TypedBinarioOutput)
+            next.exec(ctx, input)
+                .await
+                .map(|v| TypedBinarioOutput(v, PhantomData))
         },
     )
 }
diff --git a/crates/procedure/src/dyn_input.rs b/crates/procedure/src/dyn_input.rs
@@ -21,6 +21,7 @@ enum Repr<'a, 'de> {
 }
 
 impl<'a, 'de> DynInput<'a, 'de> {
+    // TODO: Explain invariant on `Option` + enforce it
     pub fn new_value<T: Send + 'static>(value: &'a mut Option<T>) -> Self {
         Self {
             inner: Repr::Value(value),
diff --git a/examples/binario/Cargo.toml b/examples/binario/Cargo.toml
@@ -15,7 +15,6 @@ tower-http = { version = "0.6.2", default-features = false, features = [
 ] }
 futures = "0.3"
 form_urlencoded = "1.2.1"
-axum-extra = { version = "0.9.6", features = ["async-read-body"] }
 tokio-util = { version = "0.7.13", features = ["compat"] }
 binario = "0.0.2"
 pin-project = "1.1.7"
diff --git a/examples/binario/client.ts b/examples/binario/client.ts
@@ -21,9 +21,26 @@
     5, 0, 0, 0, 0, 0, 0, 0, 79, 115, 99, 97, 114,
   ]);
   if (!isEqualBytes(result, expected))
-    throw new Error("Result doesn't match expected value");
+    throw new Error(`Result doesn't match expected value. Got ${result}`);
 
   console.log("Success!", result);
+
+  const resp2 = await fetch(
+    "http://localhost:4000/rspc/binario?procedure=streaming",
+    {
+      method: "POST",
+      headers: {
+        "Content-Type": "text/x-binario",
+      },
+      // { name: "Oscar" }
+      body: new Uint8Array([5, 0, 0, 0, 0, 0, 0, 0, 79, 115, 99, 97, 114]),
+    },
+  );
+  if (!resp2.ok) throw new Error(`Failed to fetch ${resp2.status}`);
+  if (resp2.headers.get("content-type") !== "text/x-binario")
+    throw new Error("Invalid content type");
+
+  console.log(await resp2.arrayBuffer());
 })();
 
 function isEqualBytes(bytes1: Uint8Array, bytes2: Uint8Array): boolean {
diff --git a/examples/binario/src/main.rs b/examples/binario/src/main.rs
@@ -42,17 +42,29 @@ impl<TErr> BaseProcedure<TErr> {
     }
 }
 
+#[derive(Debug, Clone, binario::Encode, binario::Decode, Type)]
+pub struct Input {
+    name: String,
+}
+
 pub fn mount() -> rspc::Router<()> {
-    rspc::Router::new().procedure("binario", {
-        #[derive(Debug, binario::Encode, binario::Decode, Type)]
-        pub struct Input {
-            name: String,
-        }
-
-        <BaseProcedure>::builder()
-            .with(rspc_binario::binario())
-            .query(|_, input: Input| async move { Ok(input) })
-    })
+    rspc::Router::new()
+        .procedure("binario", {
+            <BaseProcedure>::builder()
+                .with(rspc_binario::binario())
+                .query(|_, input: Input| async move { Ok(input) })
+        })
+        .procedure("streaming", {
+            <BaseProcedure>::builder()
+                .with(rspc_binario::binario())
+                .query(|_, input: Input| async move {
+                    Ok(rspc::Stream(futures::stream::iter([
+                        input.clone(),
+                        input.clone(),
+                        input,
+                    ])))
+                })
+        })
 }
 
 #[tokio::main]
@@ -104,21 +116,37 @@ pub fn rspc_binario_handler(procedures: Procedures<()>) -> axum::Router {
                     ctx.clone(),
                     rspc_binario::BinarioInput::from_stream(
                         body.into_data_stream()
-                            .map_err(|err| todo!()) // TODO: Error handling
+                            .map_err(|_err| todo!()) // TODO: Error handling
                             .into_async_read()
                             .compat(),
                     ),
                 );
                 let mut headers = HeaderMap::new();
                 headers.insert(header::CONTENT_TYPE, "text/x-binario".parse().unwrap());
 
+                let mut first = true;
                 (
                     headers,
-                    Body::from_stream(stream.map(|v| match v {
-                        Ok(v) => Ok(Ok::<_, Infallible>(
-                            v.as_value::<BinarioOutput>().unwrap().0,
-                        )),
-                        Err(err) => todo!("{err:?}"),
+                    Body::from_stream(stream.map(move |v| {
+                        let buf = match v {
+                            Ok(v) => Ok(Ok::<_, Infallible>(
+                                v.as_value::<BinarioOutput>().unwrap().0,
+                            )),
+                            Err(err) => todo!("{err:?}"),
+                        };
+
+                        if first {
+                            first = false;
+                            buf
+                        } else {
+                            buf.map(|v| {
+                                v.map(|mut v| {
+                                    let mut buf = vec!['\n' as u8, '\n' as u8];
+                                    buf.append(&mut v);
+                                    buf
+                                })
+                            })
+                        }
                     })),
                 )
             }
diff --git a/examples/bindings.ts b/examples/bindings.ts
@@ -11,6 +11,7 @@ export type Procedures = {
 	newstuffser: { kind: "query", input: any, output: any, error: any },
 	sfmPost: { kind: "query", input: any, output: any, error: any },
 	sfmPostEdit: { kind: "query", input: any, output: any, error: any },
+	streamInStreamInStreamInStream: { kind: "query", input: any, output: any, error: any },
 	validator: { kind: "query", input: any, output: any, error: any },
 	withoutBaseProcedure: { kind: "query", input: any, output: any, error: any },
 }
diff --git a/examples/core/Cargo.toml b/examples/core/Cargo.toml
@@ -20,3 +20,4 @@ async-stream = "0.3.6"
 thiserror = "2.0.9"
 validator = { version = "0.19.0", features = ["derive"] }
 anyhow = "1.0.95"
+futures.workspace = true
diff --git a/examples/core/src/lib.rs b/examples/core/src/lib.rs
@@ -197,6 +197,14 @@ pub fn mount() -> Router<Ctx> {
         .procedure("me", {
             <BaseProcedure>::builder().query(|ctx, _: ()| async move { Ok(ctx.zer.session()?) })
         })
+        .procedure("streamInStreamInStreamInStream", {
+            // You would never actually do this but it's just checking how the system behaves
+            <BaseProcedure>::builder().query(|_, _: ()| async move {
+                Ok(rspc::Stream(rspc::Stream(rspc::Stream(
+                    futures::stream::once(async move { Ok(42) }),
+                ))))
+            })
+        })
 
     // .procedure("fileupload", {
     //     <BaseProcedure>::builder().query(|_, _: File| async { Ok(env!("CARGO_PKG_VERSION")) })
diff --git a/rspc/src/lib.rs b/rspc/src/lib.rs
@@ -6,7 +6,7 @@
 //!
 //! Checkout the official docs at <https://rspc.dev>. This documentation is generally written **for authors of middleware and adapter**.
 //!
-#![forbid(unsafe_code)]
+// #![forbid(unsafe_code)] // TODO
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![doc(
     html_logo_url = "https://github.com/specta-rs/rspc/raw/main/.github/logo.png",
diff --git a/rspc/src/procedure/resolver_output.rs b/rspc/src/procedure/resolver_output.rs
@@ -81,8 +81,6 @@ where
     TErr: Error,
     S: Stream<Item = Result<T, TErr>> + Send + 'static,
     T: ResolverOutput<TErr>,
-    // Should prevent nesting `Stream`s
-    T::T: Serialize + Send + Sync + 'static,
 {
     type T = T::T;
 
@@ -100,6 +98,6 @@ where
     fn into_procedure_stream(
         stream: impl Stream<Item = Result<Self::T, ProcedureError>> + Send + 'static,
     ) -> ProcedureStream {
-        ProcedureStream::from_stream(stream)
+        T::into_procedure_stream(stream)
     }
 }
diff --git a/rspc/src/stream.rs b/rspc/src/stream.rs
@@ -41,12 +41,12 @@ impl<S: futures_util::Stream + Clone> Clone for Stream<S> {
     }
 }
 
-// TODO: I hate this requiring `Unpin` but we couldn't use `pin-project-lite` with the tuple variant.
-impl<S: futures_util::Stream + Unpin> futures_util::Stream for Stream<S> {
+impl<S: futures_util::Stream> futures_util::Stream for Stream<S> {
     type Item = S::Item;
 
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        self.0.poll_next_unpin(cx)
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        // TODO: Using `pin-project-lite` would be nice but I don't think it supports tuple variants and I don't want the macros of `pin-project`.
+        unsafe { self.map_unchecked_mut(|v| &mut v.0) }.poll_next_unpin(cx)
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ enum Repr<'a, 'de> {`
`21`	`21`	`}`
`22`	`22`
`23`	`23`	`impl<'a, 'de> DynInput<'a, 'de> {`
	`24`	+ // TODO: Explain invariant on `Option` + enforce it
`24`	`25`	`pub fn new_value<T: Send + 'static>(value: &'a mut Option<T>) -> Self {`
`25`	`26`	`Self {`
`26`	`27`	`inner: Repr::Value(value),`
Original file line number	Diff line number	Diff line change
`@@ -11,6 +11,7 @@ export type Procedures = {`
`11`	`11`	`newstuffser: { kind: "query", input: any, output: any, error: any },`
`12`	`12`	`sfmPost: { kind: "query", input: any, output: any, error: any },`
`13`	`13`	`sfmPostEdit: { kind: "query", input: any, output: any, error: any },`
	`14`	`+ streamInStreamInStreamInStream: { kind: "query", input: any, output: any, error: any },`
`14`	`15`	`validator: { kind: "query", input: any, output: any, error: any },`
`15`	`16`	`withoutBaseProcedure: { kind: "query", input: any, output: any, error: any },`
`16`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -81,8 +81,6 @@ where`
`81`	`81`	`TErr: Error,`
`82`	`82`	`S: Stream<Item = Result<T, TErr>> + Send + 'static,`
`83`	`83`	`T: ResolverOutput<TErr>,`
`84`		- // Should prevent nesting `Stream`s
`85`		`- T::T: Serialize + Send + Sync + 'static,`
`86`	`84`	`{`
`87`	`85`	`type T = T::T;`
`88`	`86`
`@@ -100,6 +98,6 @@ where`
`100`	`98`	`fn into_procedure_stream(`
`101`	`99`	`stream: impl Stream<Item = Result<Self::T, ProcedureError>> + Send + 'static,`
`102`	`100`	`) -> ProcedureStream {`
`103`		`- ProcedureStream::from_stream(stream)`
	`101`	`+ T::into_procedure_stream(stream)`
`104`	`102`	`}`
`105`	`103`	`}`