From 69890cf43ecdeb8712bf1e56c0ba8ee45c08e398 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Wed, 2 Jun 2021 19:01:45 +0200
Subject: [PATCH] Allow errors on boot to be reported (#311)

Because we unregistered standard_error, if there
was an error booting livebook, we would fail to
print the error. This commit makes sure the original
standard error is added back whenever the proxy
process terminates.
---
 lib/livebook/application.ex                   |  3 --
 .../runtime/erl_dist/io_forward_gl.ex         | 30 +++++++++++++++----
 .../runtime/erl_dist/io_forward_gl_test.exs   |  2 +-
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex
index 3a68508a8..deebe0bfa 100644
--- a/lib/livebook/application.ex
+++ b/lib/livebook/application.ex
@@ -10,9 +10,6 @@ defmodule Livebook.Application do
     validate_hostname_resolution!()
     set_cookie()
 
-    # We register our own :standard_error below
-    Process.unregister(:standard_error)
-
     children = [
       # Start the Telemetry supervisor
       LivebookWeb.Telemetry,
diff --git a/lib/livebook/runtime/erl_dist/io_forward_gl.ex b/lib/livebook/runtime/erl_dist/io_forward_gl.ex
index a1c0a0f80..4696f943a 100644
--- a/lib/livebook/runtime/erl_dist/io_forward_gl.ex
+++ b/lib/livebook/runtime/erl_dist/io_forward_gl.ex
@@ -21,19 +21,27 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
   ## Options
 
     * `:name` - the name to regsiter the process under. Optional.
+      If the name is already used, it will be unregistered before
+      starting the process and registered back when the server
+      terminates.
   """
   @spec start_link() :: GenServer.on_start()
   def start_link(opts \\ []) do
-    {gen_opts, opts} = Keyword.split(opts, [:name])
+    name = opts[:name]
 
-    GenServer.start_link(__MODULE__, opts, gen_opts)
+    if previous = name && Process.whereis(name) do
+      Process.unregister(name)
+    end
+
+    GenServer.start_link(__MODULE__, {name, previous}, opts)
   end
 
   ## Callbacks
 
   @impl true
-  def init(_opts) do
-    {:ok, %{}}
+  def init({name, previous}) do
+    Process.flag(:trap_exit, true)
+    {:ok, %{previous: {name, previous}, replies: %{}}}
   end
 
   @impl true
@@ -43,7 +51,7 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
         # Forward the request to sender's group leader
         # and instruct it to get back to us.
         send(group_leader, {:io_request, self(), reply_as, req})
-        state = Map.put(state, reply_as, from)
+        state = put_in(state.replies[reply_as], from)
 
         {:noreply, state}
 
@@ -54,9 +62,19 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
 
   def handle_info({:io_reply, reply_as, reply}, state) do
     # Forward the reply from group leader to the original client.
-    {initially_from, state} = Map.pop(state, reply_as)
+    {initially_from, state} = pop_in(state.replies[reply_as])
     send(initially_from, {:io_reply, reply_as, reply})
 
     {:noreply, state}
   end
+
+  @impl true
+  def terminate(_, %{previous: {name, previous}}) do
+    if name && previous do
+      Process.unregister(name)
+      Process.register(previous, name)
+    end
+
+    :ok
+  end
 end
diff --git a/test/livebook/runtime/erl_dist/io_forward_gl_test.exs b/test/livebook/runtime/erl_dist/io_forward_gl_test.exs
index f77970aa9..cbce629d7 100644
--- a/test/livebook/runtime/erl_dist/io_forward_gl_test.exs
+++ b/test/livebook/runtime/erl_dist/io_forward_gl_test.exs
@@ -4,7 +4,7 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGLTest do
   alias Livebook.Runtime.ErlDist.IOForwardGL
 
   test "forwards requests to sender's group leader" do
-    {:ok, pid} = IOForwardGL.start_link()
+    pid = start_supervised!(IOForwardGL)
 
     group_leader_io =
       ExUnit.CaptureIO.capture_io(:stdio, fn ->