Add recipient check to support multiple clients on the same WS stream

fonsp · Nov 8, 2021 · 58fc26e · 58fc26e
1 parent 5bf0145
commit 58fc26e
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 8 deletions.
diff --git a/frontend/common/PlutoConnection.js b/frontend/common/PlutoConnection.js
@@ -193,7 +193,6 @@ const create_vscode_connection = (address, { on_message, on_socket_close }, time
             last_task = last_task.then(async () => {
                 try {
                     const raw = event.data // The json-encoded data that the extension sent
-                    console.log("raw", raw)
                     if (raw.type === "ws_proxy") {
                         const buffer = await decode_base64_to_arraybuffer(raw.base64_encoded)
                         const message = unpack(new Uint8Array(buffer))
@@ -400,9 +399,14 @@ export const create_pluto_connection = async ({
         try {
             ws_connection = await (vscode_available ? create_vscode_connection : create_ws_connection)(String(ws_address), {
                 on_message: (update) => {
-                    const by_me = update.initiator_id == client_id
+                    const for_me = update.recipient_id === client_id
+                    const by_me = update.initiator_id === client_id
                     const request_id = update.request_id
 
+                    if (!for_me) {
+                        return
+                    }
+
                     if (by_me && request_id) {
                         const request = sent_requests.get(request_id)
                         if (request) {

diff --git a/frontend/components/Editor.js b/frontend/components/Editor.js
@@ -435,7 +435,6 @@ export class Editor extends Component {
                 return await this.actions.add_remote_cell_at(index + delta, code)
             },
             confirm_delete_multiple: async (verb, cell_ids) => {
-                console.log(confirm)
                 if (cell_ids.length <= 1 || (await confirm(`${verb} ${cell_ids.length} cells?`))) {
                     if (cell_ids.some((cell_id) => this.state.notebook.cell_results[cell_id].running || this.state.notebook.cell_results[cell_id].queued)) {
                         if (await confirm('This cell is still running - would you like to interrupt the notebook?')) {

diff --git a/src/webserver/Dynamic.jl b/src/webserver/Dynamic.jl
@@ -177,6 +177,7 @@ Update the local state of all clients connected to this notebook.
 """
 function send_notebook_changes!(🙋::ClientRequest; commentary::Any=nothing)
     notebook_dict = notebook_to_js(🙋.notebook)
+    @info "Connected clients" length(🙋.session.connected_clients) Set(c -> c.stream for c in 🙋.session.connected_clients)
     for (_, client) in 🙋.session.connected_clients
         if client.connected_notebook !== nothing && client.connected_notebook.notebook_id == 🙋.notebook.notebook_id
             current_dict = get(current_state_for_clients, client, :empty)
@@ -186,6 +187,7 @@ function send_notebook_changes!(🙋::ClientRequest; commentary::Any=nothing)
 
             # Make sure we do send a confirmation to the client who made the request, even without changes
             is_response = 🙋.initiator !== nothing && client == 🙋.initiator.client
+            @info "Responding" is_response (🙋.initiator !== nothing) (commentary === nothing)
 
             if !isempty(patches) || is_response
                 response = Dict(
@@ -279,6 +281,7 @@ responses[:update_notebook] = function response_update_notebook(🙋::ClientRequ
         patches = (Base.convert(Firebasey.JSONPatch, update) for update in 🙋.body["updates"])
 
         if length(patches) == 0
+            @info "Empty patches"
             send_notebook_changes!(🙋)
             return nothing
         end

diff --git a/src/webserver/PutUpdates.jl b/src/webserver/PutUpdates.jl
@@ -1,6 +1,10 @@
 
-function serialize_message_to_stream(io::IO, message::UpdateMessage)
-    to_send = Dict(:type => message.type, :message => message.message)
+function serialize_message_to_stream(io::IO, message::UpdateMessage, recipient::ClientSession)
+    to_send = Dict{Symbol,Any}(
+        :type => message.type, 
+        :message => message.message,
+        :recipient_id => recipient.id,
+    )
     if message.notebook !== nothing
         to_send[:notebook_id] = message.notebook.notebook_id
     end
@@ -15,8 +19,8 @@ function serialize_message_to_stream(io::IO, message::UpdateMessage)
     pack(io, to_send)
 end
 
-function serialize_message(message::UpdateMessage)
-    sprint(serialize_message_to_stream, message)
+function serialize_message(message::UpdateMessage, recipient::ClientSession)
+    sprint(serialize_message_to_stream, message, recipient)
 end
 
 "Send `messages` to all clients connected to the `notebook`."
@@ -76,7 +80,7 @@ function flushclient(client::ClientSession)
                     if client.stream isa HTTP.WebSockets.WebSocket
                         client.stream.frame_type = HTTP.WebSockets.WS_BINARY
                     end
-                    write(client.stream, serialize_message(next_to_send))
+                    write(client.stream, serialize_message(next_to_send, client))
                 else
                     put!(flushtoken)
                     return false