Skip to content

Commit

Permalink
fix: ws on rt_worker: wrap ws in Rc instead of clone ws itself (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
kanarus authored Nov 14, 2024
1 parent b9c49b4 commit b9267df
Showing 1 changed file with 15 additions and 12 deletions.
27 changes: 15 additions & 12 deletions ohkami/src/ws/worker.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
pub use mews::{Message, CloseFrame, CloseCode};

pub(crate) use worker::WebSocket as Session;

use worker::{WebSocketPair, EventStream, wasm_bindgen_futures};
use std::rc::Rc;

impl<'req> super::WebSocketContext<'req> {
pub fn upgrade<H, F>(
Expand All @@ -17,13 +20,15 @@ impl<'req> super::WebSocketContext<'req> {
} = WebSocketPair::new().expect("failed to create WebSocketPair");

ws.accept().ok();
wasm_bindgen_futures::spawn_local(async move {
handler(Connection::new(ws.clone())).await;

// https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close
//
// > If the connection is already CLOSED, this method does nothing.
ws.close::<&str>(None, None).ok();
wasm_bindgen_futures::spawn_local({
let ws = Rc::new(ws);
async move {
handler(Connection::new(ws.clone())).await;
// https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close
//
// > If the connection is already CLOSED, this method does nothing.
ws.close::<&str>(None, None).ok();
}
});

WebSocket(session)
Expand Down Expand Up @@ -52,11 +57,11 @@ impl<'req> super::WebSocketContext<'req> {
}

pub struct Connection {
ws: worker::WebSocket,
ws: Rc<worker::WebSocket>,
events: Option<EventStream<'static>>,
}
impl Connection {
fn new(ws: worker::WebSocket) -> Self {
fn new(ws: Rc<worker::WebSocket>) -> Self {
Self { ws, events:None }
}

Expand Down Expand Up @@ -122,7 +127,7 @@ pub mod split {
use super::*;

pub struct ReadHalf(Connection);
pub struct WriteHalf(worker::WebSocket);
pub struct WriteHalf(Rc<worker::WebSocket>);

impl super::Connection {
pub fn split(self) -> (ReadHalf, WriteHalf) {
Expand Down Expand Up @@ -155,8 +160,6 @@ pub mod split {
}
}

pub type Session = ::worker::WebSocket;

pub struct WebSocket(Session);
impl crate::IntoResponse for WebSocket {
fn into_response(self) -> crate::Response {
Expand Down

0 comments on commit b9267df

Please sign in to comment.