From 0386a032a149cabfb3a25565425744a24012da4c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Mon, 21 Feb 2022 20:39:29 +0100
Subject: [PATCH] Reinitialize hooks on reconnection (#1018)

* Reinitialize cells on reconnection

* Send initial session data without additional roundtrip
---
 assets/js/cell/index.js               |  8 ++++++++
 assets/js/session/index.js            | 19 +++++++++++--------
 lib/livebook_web/live/session_live.ex | 27 ++++++++++++++-------------
 3 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/assets/js/cell/index.js b/assets/js/cell/index.js
index c81be85d6..ec854141b 100644
--- a/assets/js/cell/index.js
+++ b/assets/js/cell/index.js
@@ -164,6 +164,14 @@ const Cell = {
     );
   },
 
+  disconnected() {
+    // When disconnected, this client is no longer seen by the server
+    // and misses all collaborative changes. On reconnection we want
+    // to clean up and mount a fresh hook, which we force by ensuring
+    // the DOM id doesn't match
+    this.el.removeAttribute("id");
+  },
+
   destroyed() {
     this._unsubscribeFromNavigationEvents();
     this._unsubscribeFromCellsEvents();
diff --git a/assets/js/session/index.js b/assets/js/session/index.js
index 28fc4862e..e676836c3 100644
--- a/assets/js/session/index.js
+++ b/assets/js/session/index.js
@@ -75,14 +75,6 @@ const Session = {
 
     setFavicon(faviconForEvaluationStatus(this.props.globalStatus));
 
-    // Load initial data
-
-    this.pushEvent("session_init", {}, ({ clients }) => {
-      clients.forEach((client) => {
-        this.state.clientsMap[client.pid] = client;
-      });
-    });
-
     // DOM events
 
     this.handleDocumentKeyDown = (event) => {
@@ -159,6 +151,12 @@ const Session = {
 
     // Server events
 
+    this.handleEvent("session_init", ({ clients }) => {
+      clients.forEach((client) => {
+        this.state.clientsMap[client.pid] = client;
+      });
+    });
+
     this.handleEvent("cell_inserted", ({ cell_id: cellId }) => {
       handleCellInserted(this, cellId);
     });
@@ -235,6 +233,11 @@ const Session = {
     }
   },
 
+  disconnected() {
+    // Reinitialize on reconnection
+    this.el.removeAttribute("id");
+  },
+
   destroyed() {
     this._unsubscribeFromSessionEvents();
 
diff --git a/lib/livebook_web/live/session_live.ex b/lib/livebook_web/live/session_live.ex
index ea89a48ce..a923b8203 100644
--- a/lib/livebook_web/live/session_live.ex
+++ b/lib/livebook_web/live/session_live.ex
@@ -26,6 +26,20 @@ defmodule LivebookWeb.SessionLive do
             Session.get_data(session_pid)
           end
 
+        socket =
+          if connected?(socket) do
+            payload = %{
+              clients:
+                Enum.map(data.clients_map, fn {client_pid, user_id} ->
+                  client_info(client_pid, data.users_map[user_id])
+                end)
+            }
+
+            push_event(socket, "session_init", payload)
+          else
+            socket
+          end
+
         session = Session.get_by_pid(session_pid)
 
         platform = platform_from_socket(socket)
@@ -565,19 +579,6 @@ defmodule LivebookWeb.SessionLive do
   end
 
   @impl true
-  def handle_event("session_init", _params, socket) do
-    data = socket.private.data
-
-    payload = %{
-      clients:
-        Enum.map(data.clients_map, fn {client_pid, user_id} ->
-          client_info(client_pid, data.users_map[user_id])
-        end)
-    }
-
-    {:reply, payload, socket}
-  end
-
   def handle_event("append_section", %{}, socket) do
     idx = length(socket.private.data.notebook.sections)
     Session.insert_section(socket.assigns.session.pid, idx)