Use WebSocket instead of Long Polling (#1142)

tests/rojo_test/internable.rs

@@ -3,7 +3,9 @@ use std::collections::HashMap;
 use rbx_dom_weak::types::Ref;
 use serde::Serialize;
 
-use librojo::web_api::{Instance, InstanceUpdate, ReadResponse, SubscribeResponse};
+use librojo::web_api::{
+    Instance, InstanceUpdate, MessagesPacket, ReadResponse, SocketPacket, SocketPacketBody,
+};
 use rojo_insta_ext::RedactionMap;
 
 /// A convenience method to store all of the redactable data from a piece of
@@ -54,7 +56,16 @@ impl<'a> Internable<&'a HashMap<Ref, Instance<'_>>> for Instance<'a> {
     }
 }
 
-impl Internable<()> for SubscribeResponse<'_> {
+impl Internable<()> for SocketPacket<'_> {
+    fn intern(&self, redactions: &mut RedactionMap, extra: ()) {
+        redactions.intern(&self.session_id);
+        match &self.body {
+            SocketPacketBody::Messages(packet) => packet.intern(redactions, extra),
+        }
+    }
+}
+
+impl Internable<()> for MessagesPacket<'_> {
     fn intern(&self, redactions: &mut RedactionMap, _extra: ()) {
         for message in &self.messages {
             intern_instance_updates(redactions, &message.updated);

tests/rojo_test/serve_util.rs

@@ -8,11 +8,14 @@ use std::{
     time::Duration,
 };
 
+use hyper_tungstenite::tungstenite::{connect, Message};
 use rbx_dom_weak::types::Ref;
 
 use tempfile::{tempdir, TempDir};
 
-use librojo::web_api::{ReadResponse, SerializeResponse, ServerInfoResponse, SubscribeResponse};
+use librojo::web_api::{
+    ReadResponse, SerializeResponse, ServerInfoResponse, SocketPacket, SocketPacketType,
+};
 use rojo_insta_ext::RedactionMap;
 
 use crate::rojo_test::io_util::{
@@ -173,13 +176,54 @@ impl TestServeSession {
         Ok(serde_json::from_value(value).expect("Server returned malformed response"))
     }
 
-    pub fn get_api_subscribe(
+    pub fn get_api_socket_packet(
         &self,
+        packet_type: SocketPacketType,
         cursor: u32,
-    ) -> Result<SubscribeResponse<'static>, reqwest::Error> {
-        let url = format!("http://localhost:{}/api/subscribe/{}", self.port, cursor);
+    ) -> Result<SocketPacket<'static>, Box<dyn std::error::Error>> {
+        let url = format!("ws://localhost:{}/api/socket/{}", self.port, cursor);
 
-        reqwest::blocking::get(url)?.json()
+        let (mut socket, _response) = connect(url)?;
+
+        // Wait for messages with a timeout
+        let timeout = Duration::from_secs(10);
+        let start = std::time::Instant::now();
+
+        loop {
+            if start.elapsed() > timeout {
+                return Err("Timeout waiting for packet from WebSocket".into());
+            }
+
+            match socket.read() {
+                Ok(Message::Text(text)) => {
+                    let packet: SocketPacket = serde_json::from_str(&text)?;
+                    if packet.packet_type != packet_type {
+                        continue;
+                    }
+
+                    // Close the WebSocket connection now that we got what we were waiting for
+                    let _ = socket.close(None);
+                    return Ok(packet);
+                }
+                Ok(Message::Close(_)) => {
+                    return Err("WebSocket closed before receiving messages".into());
+                }
+                Ok(_) => {
+                    // Ignore other message types (ping, pong, binary)
+                    continue;
+                }
+                Err(hyper_tungstenite::tungstenite::Error::Io(e))
+                    if e.kind() == std::io::ErrorKind::WouldBlock =>
+                {
+                    // No data available yet, sleep a bit and try again
+                    thread::sleep(Duration::from_millis(100));
+                    continue;
+                }
+                Err(e) => {
+                    return Err(e.into());
+                }
+            }
+        }
     }
 
     pub fn get_api_serialize(&self, ids: &[Ref]) -> Result<SerializeResponse, reqwest::Error> {

tests/tests/serve.rs

@@ -8,6 +8,8 @@ use crate::rojo_test::{
     serve_util::{run_serve_test, serialize_to_xml_model},
 };
 
+use librojo::web_api::SocketPacketType;
+
 #[test]
 fn empty() {
     run_serve_test("empty", |session, mut redactions| {
@@ -42,10 +44,12 @@ fn scripts() {
 
         fs::write(session.path().join("src/foo.lua"), "Updated foo!").unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "scripts_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -74,10 +78,12 @@ fn add_folder() {
 
         fs::create_dir(session.path().join("src/my-new-folder")).unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "add_folder_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -104,10 +110,12 @@ fn remove_file() {
 
         fs::remove_file(session.path().join("src/hello.txt")).unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "remove_file_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -134,10 +142,12 @@ fn edit_init() {
 
         fs::write(session.path().join("src/init.lua"), b"-- Edited contents").unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "edit_init_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -180,10 +190,12 @@ fn move_folder_of_stuff() {
         // will fail otherwise.
         fs::rename(stuff_path, session.path().join("src/new-stuff")).unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "move_folder_of_stuff_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -214,10 +226,12 @@ fn empty_json_model() {
         )
         .unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "empty_json_model_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -245,10 +259,12 @@ fn add_optional_folder() {
 
         fs::create_dir(session.path().join("create-later")).unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "add_optional_folder_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -437,10 +453,12 @@ fn ref_properties() {
         )
         .unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "ref_properties_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -467,10 +485,12 @@ fn ref_properties_remove() {
 
         fs::remove_file(session.path().join("src/target.model.json")).unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "ref_properties_remove_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -529,10 +549,12 @@ fn ref_properties_patch_update() {
         )
         .unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "ref_properties_patch_update_subscribe",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();
@@ -581,10 +603,12 @@ fn model_pivot_migration() {
         )
         .unwrap();
 
-        let subscribe_response = session.get_api_subscribe(0).unwrap();
+        let socket_packet = session
+            .get_api_socket_packet(SocketPacketType::Messages, 0)
+            .unwrap();
         assert_yaml_snapshot!(
             "model_pivot_migration_all",
-            subscribe_response.intern_and_redact(&mut redactions, ())
+            socket_packet.intern_and_redact(&mut redactions, ())
         );
 
         let read_response = session.get_api_read(root_id).unwrap();