denoland · littledivy · Oct 30, 2023 · Sep 19, 2023 · Oct 28, 2023 · Oct 28, 2023
diff --git a/Cargo.toml b/Cargo.toml
@@ -37,6 +37,7 @@ thiserror = "1.0.40"
 default = ["simd"]
 simd = ["simdutf8/aarch64_neon"]
 upgrade = ["hyper", "pin-project", "base64", "sha1"]
+unstable-split = []
 
 [dev-dependencies]
 tokio = { version = "1.25.0", features = ["full", "macros"] }

diff --git a/examples/echo_server_split.rs b/examples/echo_server_split.rs
@@ -0,0 +1,84 @@
+// Copyright 2023 Divy Srivastava <[email protected]>
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use fastwebsockets::upgrade;
+use fastwebsockets::FragmentCollectorRead;
+use fastwebsockets::OpCode;
+use fastwebsockets::WebSocketError;
+use hyper::server::conn::Http;
+use hyper::service::service_fn;
+use hyper::Body;
+use hyper::Request;
+use hyper::Response;
+use tokio::net::TcpListener;
+
+async fn handle_client(fut: upgrade::UpgradeFut) -> Result<(), WebSocketError> {
+  let ws = fut.await?;
+  let (rx, mut tx) = ws.split(|ws| tokio::io::split(ws));
+  let mut rx = FragmentCollectorRead::new(rx);
+  loop {
+    // Empty send_fn is fine because the benchmark does not create obligated writes.
+    let frame = rx
+      .read_frame::<_, WebSocketError>(&mut move |_| async {
+        unreachable!();
+      })
+      .await?;
+    match frame.opcode {
+      OpCode::Close => break,
+      OpCode::Text | OpCode::Binary => {
+        tx.write_frame(frame).await?;
+      }
+      _ => {}
+    }
+  }
+
+  Ok(())
+}
+async fn server_upgrade(
+  mut req: Request<Body>,
+) -> Result<Response<Body>, WebSocketError> {
+  let (response, fut) = upgrade::upgrade(&mut req)?;
+
+  tokio::task::spawn(async move {
+    if let Err(e) = tokio::task::unconstrained(handle_client(fut)).await {
+      eprintln!("Error in websocket connection: {}", e);
+    }
+  });
+
+  Ok(response)
+}
+
+fn main() -> Result<(), WebSocketError> {
+  let rt = tokio::runtime::Builder::new_current_thread()
+    .enable_io()
+    .build()
+    .unwrap();
+
+  rt.block_on(async move {
+    let listener = TcpListener::bind("127.0.0.1:8080").await?;
+    println!("Server started, listening on {}", "127.0.0.1:8080");
+    loop {
+      let (stream, _) = listener.accept().await?;
+      println!("Client connected");
+      tokio::spawn(async move {
+        let conn_fut = Http::new()
+          .serve_connection(stream, service_fn(server_upgrade))
+          .with_upgrades();
+        if let Err(e) = conn_fut.await {
+          println!("An error occurred: {:?}", e);
+        }
+      });
+    }
+  })
+}
diff --git a/src/error.rs b/src/error.rs
@@ -41,4 +41,7 @@ pub enum WebSocketError {
   #[cfg(feature = "upgrade")]
   #[error(transparent)]
   HTTPError(#[from] hyper::Error),
+  #[cfg(feature = "unstable-split")]
+  #[error("Failed to send frame")]
+  SendError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
 }
diff --git a/src/fragment.rs b/src/fragment.rs
@@ -12,12 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#[cfg(feature = "unstable-split")]
+use std::future::Future;
+
 use crate::error::WebSocketError;
 use crate::frame::Frame;
 use crate::recv::SharedRecv;
 use crate::OpCode;
 use crate::ReadHalf;
 use crate::WebSocket;
+#[cfg(feature = "unstable-split")]
+use crate::WebSocketRead;
 use crate::WriteHalf;
 use tokio::io::AsyncReadExt;
 use tokio::io::AsyncWriteExt;
@@ -136,6 +141,60 @@ impl<'f, S> FragmentCollector<S> {
   }
 }
 
+#[cfg(feature = "unstable-split")]
+pub struct FragmentCollectorRead<S> {
+  stream: S,
+  read_half: ReadHalf,
+  fragments: Fragments,
+  // !Sync marker
+  _marker: std::marker::PhantomData<SharedRecv>,
+}
+
+#[cfg(feature = "unstable-split")]
+impl<'f, S> FragmentCollectorRead<S> {
+  /// Creates a new `FragmentCollector` with the provided `WebSocket`.
+  pub fn new(ws: WebSocketRead<S>) -> FragmentCollectorRead<S>
+  where
+    S: AsyncReadExt + Unpin,
+  {
+    let (stream, read_half) = ws.into_parts_internal();
+    FragmentCollectorRead {
+      stream,
+      read_half,
+      fragments: Fragments::new(),
+      _marker: std::marker::PhantomData,
+    }
+  }
+
+  /// Reads a WebSocket frame, collecting fragmented messages until the final frame is received and returns the completed message.
+  ///
+  /// Text frames payload is guaranteed to be valid UTF-8.
+  pub async fn read_frame<R, E>(
+    &mut self,
+    send_fn: &mut impl FnMut(Frame<'f>) -> R,
+  ) -> Result<Frame<'f>, WebSocketError>
+  where
+    S: AsyncReadExt + Unpin,
+    E: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
+    R: Future<Output = Result<(), E>>,
+  {
+    loop {
+      let (res, obligated_send) =
+        self.read_half.read_frame_inner(&mut self.stream).await;
+      if let Some(frame) = obligated_send {
+        let res = send_fn(frame).await;
+        res.map_err(|e| WebSocketError::SendError(e.into()))?;
+      }
+      let Some(frame) = res? else {
+        continue;
+      };
+      if let Some(frame) = self.fragments.accumulate(frame)? {
+        return Ok(frame);
+      }
+    }
+  }
+}
+
 /// Accumulates potentially fragmented [`Frame`]s to defragment the incoming WebSocket stream.
 struct Fragments {
   fragments: Option<Fragment>,