From 1b7b3080e1d47d4d0c80aaf39d4cbc6d849f6527 Mon Sep 17 00:00:00 2001
From: gpopides <46693847+gpopides@users.noreply.github.com>
Date: Mon, 24 Jan 2022 23:23:31 +0100
Subject: [PATCH] disconnect node from remote machine (#915)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* disconnect node from remote machine

* return ok if disconnect is successful

* formatting

* Update lib/livebook/runtime/attached.ex

* move disconnection of node to NodeManager

* add doc about parent node option

* Update lib/livebook/runtime/erl_dist/node_manager.ex

Co-authored-by: José Valim <jose.valim@gmail.com>
Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>
---
 lib/livebook/runtime/attached.ex              |  3 ++-
 lib/livebook/runtime/erl_dist/node_manager.ex | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/lib/livebook/runtime/attached.ex b/lib/livebook/runtime/attached.ex
index 6c1126cd8..f51e2f474 100644
--- a/lib/livebook/runtime/attached.ex
+++ b/lib/livebook/runtime/attached.ex
@@ -28,7 +28,8 @@ defmodule Livebook.Runtime.Attached do
 
     case Node.ping(node) do
       :pong ->
-        server_pid = Livebook.Runtime.ErlDist.initialize(node)
+        opts = [parent_node: node()]
+        server_pid = Livebook.Runtime.ErlDist.initialize(node, opts)
         {:ok, %__MODULE__{node: node, cookie: cookie, server_pid: server_pid}}
 
       :pang ->
diff --git a/lib/livebook/runtime/erl_dist/node_manager.ex b/lib/livebook/runtime/erl_dist/node_manager.ex
index 80883e219..804317ecb 100644
--- a/lib/livebook/runtime/erl_dist/node_manager.ex
+++ b/lib/livebook/runtime/erl_dist/node_manager.ex
@@ -37,6 +37,11 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do
 
     * `:auto_termination` - whether to terminate the manager
       when the last runtime server terminates. Defaults to `true`.
+
+    * `:parent_node` - indicates which node spawned the node manager.  
+       It is used to disconnect the node when the server terminates, 
+       which happens when the last session using the node disconnects.  
+       Defaults to `nil`
   """
   def start(opts \\ []) do
     {opts, gen_opts} = split_opts(opts)
@@ -78,6 +83,7 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do
   def init(opts) do
     unload_modules_on_termination = Keyword.get(opts, :unload_modules_on_termination, true)
     auto_termination = Keyword.get(opts, :auto_termination, true)
+    parent_node = Keyword.get(opts, :parent_node)
 
     ## Initialize the node
 
@@ -105,7 +111,8 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do
        server_supevisor: server_supevisor,
        runtime_servers: [],
        initial_ignore_module_conflict: initial_ignore_module_conflict,
-       original_standard_error: original_standard_error
+       original_standard_error: original_standard_error,
+       parent_node: parent_node
      }}
   end
 
@@ -122,6 +129,10 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do
       ErlDist.unload_required_modules()
     end
 
+    if state.parent_node do
+      Node.disconnect(state.parent_node)
+    end
+
     :ok
   end