From 0c04f659c46a453a91680454fb53984a6f5ecafd Mon Sep 17 00:00:00 2001
From: aar2dee2 <85004512+aar2dee2@users.noreply.github.com>
Date: Thu, 23 Jun 2022 17:09:27 +0530
Subject: [PATCH] Capture log from processes outside Livebook supervision
 (#1235)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>
---
 lib/livebook/runtime/erl_dist/logger_gl_backend.ex |  7 +++++--
 lib/livebook/runtime/erl_dist/node_manager.ex      |  5 +++++
 lib/livebook/runtime/erl_dist/runtime_server.ex    | 14 ++++++++++++--
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/lib/livebook/runtime/erl_dist/logger_gl_backend.ex b/lib/livebook/runtime/erl_dist/logger_gl_backend.ex
index 4ad8cd611..2888f7032 100644
--- a/lib/livebook/runtime/erl_dist/logger_gl_backend.ex
+++ b/lib/livebook/runtime/erl_dist/logger_gl_backend.ex
@@ -94,9 +94,12 @@ defmodule Livebook.Runtime.ErlDist.LoggerGLBackend do
   end
 
   defp log_event(level, msg, ts, md, gl, state) do
+    output = format_event(level, msg, ts, md, state)
+
     if io_proxy?(gl) do
-      output = format_event(level, msg, ts, md, state)
       async_io(gl, output)
+    else
+      send(Livebook.Runtime.ErlDist.NodeManager, {:orphan_log, output})
     end
   end
 
@@ -105,7 +108,7 @@ defmodule Livebook.Runtime.ErlDist.LoggerGLBackend do
     info[:dictionary][:"$initial_call"] == {Livebook.Runtime.Evaluator.IOProxy, :init, 1}
   end
 
-  defp async_io(device, output) when is_pid(device) do
+  def async_io(device, output) when is_pid(device) do
     send(device, {:io_request, self(), make_ref(), {:put_chars, :unicode, output}})
   end
 
diff --git a/lib/livebook/runtime/erl_dist/node_manager.ex b/lib/livebook/runtime/erl_dist/node_manager.ex
index 38685725d..f35dbaf01 100644
--- a/lib/livebook/runtime/erl_dist/node_manager.ex
+++ b/lib/livebook/runtime/erl_dist/node_manager.ex
@@ -151,6 +151,11 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do
     end
   end
 
+  def handle_info({:orphan_log, _output} = message, state) do
+    for pid <- state.runtime_servers, do: send(pid, message)
+    {:noreply, state}
+  end
+
   def handle_info(_message, state), do: {:noreply, state}
 
   @impl true
diff --git a/lib/livebook/runtime/erl_dist/runtime_server.ex b/lib/livebook/runtime/erl_dist/runtime_server.ex
index 1b9bdd102..f6ff8aa36 100644
--- a/lib/livebook/runtime/erl_dist/runtime_server.ex
+++ b/lib/livebook/runtime/erl_dist/runtime_server.ex
@@ -201,7 +201,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
        smart_cell_definitions_module:
          Keyword.get(opts, :smart_cell_definitions_module, Kino.SmartCell),
        extra_smart_cell_definitions: Keyword.get(opts, :extra_smart_cell_definitions, []),
-       memory_timer_ref: nil
+       memory_timer_ref: nil,
+       last_evaluator: nil
      }}
   end
 
@@ -244,6 +245,15 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:noreply, finish_scan_binding(ref, state)}
   end
 
+  def handle_info({:orphan_log, output}, state) do
+    with %{} = evaluator <- state.last_evaluator,
+         {:group_leader, io_proxy} <- Process.info(evaluator.pid, :group_leader) do
+      ErlDist.LoggerGLBackend.async_io(io_proxy, output)
+    end
+
+    {:noreply, state}
+  end
+
   def handle_info(_message, state), do: {:noreply, state}
 
   defp handle_down_evaluator(state, {:DOWN, _, :process, pid, reason}) do
@@ -338,7 +348,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
       opts
     )
 
-    {:noreply, state}
+    {:noreply, %{state | last_evaluator: state.evaluators[container_ref]}}
   end
 
   def handle_cast({:forget_evaluation, {container_ref, evaluation_ref}}, state) do