Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

lightclient: Fix wasm socket closure called after being dropped #1289

Merged
merged 10 commits into from
Nov 29, 2023
38 changes: 25 additions & 13 deletions lightclient/src/platform/wasm_socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,10 @@ struct InnerWasmSocket {
///
/// These need to be kept around until the socket is dropped.
type Callbacks = (
Closure<dyn FnMut()>,
JsValue,
Closure<dyn FnMut(web_sys::MessageEvent)>,
Closure<dyn FnMut(web_sys::Event)>,
Closure<dyn FnMut(web_sys::CloseEvent)>,
JsValue,
JsValue,
);

impl WasmSocket {
Expand All @@ -89,7 +89,7 @@ impl WasmSocket {
waker: None,
}));

let open_callback = Closure::<dyn FnMut()>::new({
let open_callback = Closure::once_into_js({
let inner = inner.clone();
move || {
let mut inner = inner.lock().expect("Mutex is poised; qed");
Expand Down Expand Up @@ -120,9 +120,9 @@ impl WasmSocket {
});
socket.set_onmessage(Some(message_callback.as_ref().unchecked_ref()));

let error_callback = Closure::<dyn FnMut(_)>::new({
let error_callback = Closure::once_into_js({
let inner = inner.clone();
move |_| {
move |_event: web_sys::Event| {
// Callback does not provide useful information, signal it back to the stream.
let mut inner = inner.lock().expect("Mutex is poised; qed");
inner.state = ConnectionState::Error;
Expand All @@ -134,9 +134,9 @@ impl WasmSocket {
});
socket.set_onerror(Some(error_callback.as_ref().unchecked_ref()));

let close_callback = Closure::<dyn FnMut(_)>::new({
let close_callback = Closure::once_into_js({
Copy link
Collaborator

@jsdw jsdw Nov 24, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docs for once_into_js say:

Convert a FnOnce(A...) -> R into a JavaScript Function object.

If the JavaScript function is invoked more than once, it will throw an exception.

Unlike Closure::once, this does not return a Closure that can be dropped before the function is invoked to deallocate the closure. The only way the FnOnce is deallocated is by calling the JavaScript function. If the JavaScript function is never called then the FnOnce and everything it closes over will leak.

So, since some of these callbacks might never fire, I think we'd be gradually leaking them here for any connection we make that doesn't fire one of the callbacks.

I wonder if the errors we got happened because previously, if Self is dropped before one of the JS websocket callbacks is fired, then it would have decallocated the closures and we'd get an error in JS land instead.

But anyway, from reading the spec, it sounds like we:

  • May never fire onopen if connection can't be opened (but will fire it max once)
  • May never fire onerror or onmessage if no errors/messages sent (not certain if onerror will only fire once or not)
  • Will always eventualy fire onclose (max once).

So if I read right, I wonder whether we can do some cleanup in onclose closure to ensure that all of the other closures are dropped. (eg Closure::new() for them all, save the Closures, and drop() them all when onclose fires or something like that)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense!

I've found another way, which I believe to be more straight forward than performing clean-up on the onclose:

  • keep the Closure::new() around for each callback
  • while dropping the object:
    • close the socket if in: OPEN or CONNECTING states
    • set closures on the web_sys::WebSocket as None

This shall clean-up all memory associated with the rust closures, and deny the javascript real access to the dropped objects

let inner = inner.clone();
move |_| {
move |_event: web_sys::CloseEvent| {
let mut inner = inner.lock().expect("Mutex is poised; qed");
inner.state = ConnectionState::Closed;

Expand Down Expand Up @@ -171,6 +171,10 @@ impl AsyncRead for WasmSocket {
let mut inner = self.inner.lock().expect("Mutex is poised; qed");
inner.waker = Some(cx.waker().clone());

if self.socket.ready_state() == web_sys::WebSocket::CONNECTING {
return Poll::Pending;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So I think this is ook because any of the callbacks will wake the waker when JS clals them, ensuring we call this poll_read again.

}

match inner.state {
ConnectionState::Error => {
Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, "Socket error")))
Expand Down Expand Up @@ -221,16 +225,24 @@ impl AsyncWrite for WasmSocket {
Poll::Ready(Ok(()))
}

fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
Poll::Ready(Ok(()))
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
if self.socket.ready_state() == web_sys::WebSocket::CLOSED {
return Poll::Ready(Ok(()));
}

if self.socket.ready_state() != web_sys::WebSocket::CLOSING {
let _ = self.socket.close();
}

let mut inner = self.inner.lock().expect("Mutex is poised; qed");
inner.waker = Some(cx.waker().clone());
Poll::Pending
}
}

impl Drop for WasmSocket {
fn drop(&mut self) {
let inner = self.inner.lock().expect("Mutex is poised; qed");

if inner.state == ConnectionState::Opened {
if self.socket.ready_state() == web_sys::WebSocket::OPEN {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be safer to say if ready_state != CLOSING then close()? Or maybe just call close() regardless to ensure our onclose handler always fires and can tidy up (I think nothing happens if you try closing when it's already closing anyway but not sure!)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've extended this to do clean-up for OPEN and CONNECTING, since we might generate console errors if we attempt to close a socket that's already closed 👍

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Attempting to close the socket regardless of its state will generate on --chrome:

WebSocket connection to 'ws://polkadot-boot.dwellir.com:30334/' failed: WebSocket is closed before the connection is established.

The same error is generated if the socket is in the CONNECTING state, note that --firefox is not printing this warning.

let _ = self.socket.close();
}
}
Expand Down
2 changes: 2 additions & 0 deletions testing/wasm-lightclient-tests/tests/wasm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);

#[wasm_bindgen_test]
async fn light_client_works() {
console_error_panic_hook::set_once();

let api: LightClient<PolkadotConfig> = LightClientBuilder::new()
.build_from_url("wss://rpc.polkadot.io:443")
.await
Expand Down
Loading