How should I structure this websocket-based WASM app?

I'm trying to build a frontend application using rust wasm and a websocket to connect to the server. Here's the desired core logic:

Server (axum) needs to send a command to the app
Server creates an instance of an Enum called ServerMessage that contains a variant for each valid command that the server can send
Server serializes that instance using serde_cbor, and sends it as a binary message over the websocket
App receives the binary data and deserializes it back into a ServerMessage
App does whatever it was commanded to do by the server

I have a similar flow working in the other direction: if there's data that needs to be sent back to the server (e.g. user input), the app creates a ClientMessage, serializes it, and sends it over the websocket.

On the server side, each active websocket connection runs in its own tokio task. When a message comes in, it is deserialized into a ClientMessage, which is then passed to the main thread over an mpsc channel. The main thread processes messages one at a time as they come in. This setup works well for what I'm looking to do.

However, I'm lost on how to structure the WASM app to handle the desired logic above, largely because (as far as I'm aware) there's no way to use an async runtime or channels.

Below is a minimal example of what I'm trying to do. In this case, there are only two commands the server can send: ConsoleLog which instructs the app to log some text to console, or Increment which instructs the app to increment its value and update the display.

use serde::{Deserialize, Serialize};
use serde_cbor::from_slice;
use wasm_bindgen::prelude::*;
use web_sys::{
    console, js_sys, window, MessageEvent, WebSocket
};

#[wasm_bindgen]
pub struct App {
    websocket: WebSocket,
    value: i32
}

#[wasm_bindgen]
impl App {
    pub fn new() -> App {
        let ws = WebSocket::new("http://127.0.0.1:8080/ws").unwrap();
        ws.set_binary_type(web_sys::BinaryType::Arraybuffer);

        return App {
            websocket: ws,
            value: 0
        }
    }

    pub fn start(&mut self) {
        // OnMessage callback
        let onmessage_callback = Closure::<dyn FnMut(_)>::new(move |e: MessageEvent| {
            if let Ok(abuf) = e.data().dyn_into::<js_sys::ArrayBuffer>() {
                // Deserialize data
                let bin_vec = js_sys::Uint8Array::new(&abuf).to_vec();
                let server_message: ServerMessage = from_slice(&bin_vec).unwrap();

                // Handle message
                self.handle_message(server_message);
            }
        });

        self.websocket.set_onmessage(Some(onmessage_callback.as_ref().unchecked_ref()));
        onmessage_callback.forget();
    }
}

/// A message sent from the server (axum)
#[derive(Debug, Deserialize, Serialize)]
pub enum ServerMessage {
    ConsoleLog(String),
    Increment(i32)
}

impl App {
    fn handle_message(&mut self, message: ServerMessage) {
        match message {
            ServerMessage::ConsoleLog(text) => {
                console::log_1(&text.into());
            }
            ServerMessage::Increment(n) => {
                // Increment value
                self.value += n;
            
                // Get counter element
                let document = window().unwrap().document().unwrap();
                let element = document.get_element_by_id("counter").unwrap();
                
                // Update display
                element.set_text_content(Some(&self.value.to_string()));
            }
        }
    }
}

This code won't compile because borrowed data (the reference to self) escapes the method body. I've tried a few other variations of the code that run into similar errors with lifetimes and ownership. I have a clear vision of how I want the code to work, but I'm stuck on how to write it within the constraints of a WASM app.

How can I adjust this code to successfully use the core logic above?

Solution

Ended up getting this working by tweaking my structure and putting persistent data behind Rc and Mutex.

use serde::{Deserialize, Serialize};
use serde_cbor::from_slice;
use std::{rc::Rc, sync::Mutex};
use wasm_bindgen::prelude::*;
use web_sys::{console, js_sys, window, MessageEvent, WebSocket};


/// A message sent from the server (axum)
#[derive(Debug, Deserialize, Serialize)]
pub enum ServerMessage {
    ConsoleLog(String),
    Increment(i32)
}

/// Shared data for WASM client app
struct AppData {
    value: Mutex<i32>
}

/// Set up event listeners
#[wasm_bindgen]
pub fn start() {
    let app_data = Rc::new(AppData {
        value: Mutex::new(0)
    });

    let ws = WebSocket::new("http://127.0.0.1:8080/ws").unwrap();
    ws.set_binary_type(web_sys::BinaryType::Arraybuffer);

    { // OnMessage callback
        let ws_clone = ws.clone();
        let onmessage_callback = Closure::<dyn FnMut(_)>::new(move |e: MessageEvent| {
            if let Ok(abuf) = e.data().dyn_into::<js_sys::ArrayBuffer>() {
                // Deserialize data
                let bin_vec = js_sys::Uint8Array::new(&abuf).to_vec();
                let server_message: ServerMessage = from_slice(&bin_vec).unwrap();

                // Handle message
                handle_message(app_data.clone(), server_message, ws.clone());
            } else {
                console::log_1(&format!("Message received with unknown type: {:?}", e.data()).into());
            }
        });

        ws_clone.set_onmessage(Some(onmessage_callback.as_ref().unchecked_ref()));
        onmessage_callback.forget();
    }
}

/// Handle message based on message type
fn handle_message(app_data: Rc<AppData>, message: ServerMessage, ws: WebSocket) {
    match message {
        ServerMessage::ConsoleLog(text) => {
            // Log text to console
            console::log_1(&text.into());
        }
        ServerMessage::Increment(n) => {
            // Increment value
            let mut value = app_data.value.lock().unwrap();
            *value += 5;
            
            // Get counter element
            let document = window().unwrap().document().unwrap();
            let element = document.get_element_by_id("counter").unwrap();
            
            // Update display
            element.set_text_content(Some(&self.value.to_string()));
        }
    }
}