Simplify the forwarding IO device (#1024)

lib/livebook/runtime/erl_dist/io_forward_gl.ex

@@ -1,20 +1,17 @@
 defmodule Livebook.Runtime.ErlDist.IOForwardGL do
   @moduledoc false
 
-  # An IO device process forwarding all requests to sender's group leader.
+  # An IO device process forwarding all requests to sender's group
+  # leader.
   #
-  # We use this device as `:standard_error` on connected runtime node,
-  # so that all evaluation warnings are treated as stdout.
+  # We register this device as the `:standard_error` in the runtime
+  # node, so that all evaluation warnings are treated as stdout.
   #
   # The process implements [The Erlang I/O Protocol](https://erlang.org/doc/apps/stdlib/io_protocol.html)
-  # and can be thought of as a *virtual* IO device.
+  # and can be thought of as a virtual IO device.
 
   use GenServer
 
-  @type state :: %{(reply_as :: term()) => from :: pid()}
-
-  ## API
-
   @doc """
   Starts the IO device.
 
@@ -25,7 +22,7 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
       starting the process and registered back when the server
       terminates.
   """
-  @spec start_link() :: GenServer.on_start()
+  @spec start_link(keyword()) :: GenServer.on_start()
   def start_link(opts \\ []) do
     name = opts[:name]
 
@@ -36,12 +33,10 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
     GenServer.start_link(__MODULE__, {name, previous}, opts)
   end
 
-  ## Callbacks
-
   @impl true
   def init({name, previous}) do
     Process.flag(:trap_exit, true)
-    {:ok, %{previous: {name, previous}, replies: %{}}}
+    {:ok, %{name: name, previous: previous}}
   end
 
   @impl true
@@ -50,26 +45,17 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
       {:group_leader, group_leader} ->
         # 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 = put_in(state.replies[reply_as], from)
-
-        {:noreply, state}
+        send(group_leader, {:io_request, from, reply_as, req})
 
       _ ->
-        {:noreply, state}
+        send(from, {:io_reply, reply_as, {:error, :terminated}})
     end
-  end
-
-  def handle_info({:io_reply, reply_as, reply}, state) do
-    # Forward the reply from group leader to the original client.
-    {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
+  def terminate(_, %{name: name, previous: previous}) do
     if name && previous do
       Process.unregister(name)
       Process.register(previous, name)