Merge kino_proxy proof of concept into Livebook (#2615)

lib/livebook/proxy/adapter.ex

@@ -0,0 +1,99 @@
+defmodule Livebook.Proxy.Adapter do
+  @moduledoc false
+  @behaviour Plug.Conn.Adapter
+
+  def send_resp({pid, ref}, status, headers, body) do
+    send(pid, {:send_resp, self(), ref, status, headers, body})
+
+    receive do
+      {^ref, :ok} -> {:ok, body, {pid, ref}}
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:send_resp, 4, reason)
+    end
+  end
+
+  def get_peer_data({pid, ref}) do
+    send(pid, {:get_peer_data, self(), ref})
+
+    receive do
+      {^ref, peer_data} -> peer_data
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:get_peer_data, 1, reason)
+    end
+  end
+
+  def get_http_protocol({pid, ref}) do
+    send(pid, {:get_http_protocol, self(), ref})
+
+    receive do
+      {^ref, http_protocol} -> http_protocol
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:get_http_protocol, 1, reason)
+    end
+  end
+
+  def read_req_body({pid, ref}, opts) do
+    send(pid, {:read_req_body, self(), ref, opts})
+
+    receive do
+      {^ref, {:ok, data}} -> {:ok, data, {pid, ref}}
+      {^ref, {:more, data}} -> {:more, data, {pid, ref}}
+      {^ref, {:error, _} = error} -> error
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:read_req_body, 2, reason)
+    end
+  end
+
+  def send_chunked({pid, ref}, status, headers) do
+    send(pid, {:send_chunked, self(), ref, status, headers})
+
+    receive do
+      {^ref, :ok} -> {:ok, nil, {pid, ref}}
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:send_chunked, 3, reason)
+    end
+  end
+
+  def chunk({pid, ref}, chunk) do
+    send(pid, {:chunk, self(), ref, chunk})
+
+    receive do
+      {^ref, :ok} -> :ok
+      {^ref, {:error, _} = error} -> error
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:chunk, 2, reason)
+    end
+  end
+
+  def inform({pid, ref}, status, headers) do
+    send(pid, {:inform, self(), ref, status, headers})
+
+    receive do
+      {^ref, :ok} -> {:ok, {pid, ref}}
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:inform, 3, reason)
+    end
+  end
+
+  def send_file({pid, ref}, status, headers, path, offset, length) do
+    %File.Stat{type: :regular, size: size} = File.stat!(path)
+
+    length =
+      cond do
+        length == :all -> size
+        is_integer(length) -> length
+      end
+
+    {:ok, body} =
+      File.open!(path, [:read, :raw, :binary], fn device ->
+        :file.pread(device, offset, length)
+      end)
+
+    send(pid, {:send_resp, self(), ref, status, headers, body})
+
+    receive do
+      {^ref, :ok} -> {:ok, body, {pid, ref}}
+      {:DOWN, ^ref, _, _, reason} -> exit_fun(:send_file, 6, reason)
+    end
+  end
+
+  def upgrade(_payload, _protocol, _opts), do: {:error, :not_supported}
+  def push(_payload, _path, _headers), do: {:error, :not_supported}
+
+  defp exit_fun(fun, arity, reason) do
+    exit({{__MODULE__, fun, arity}, reason})
+  end
+end

lib/livebook/proxy/handler.ex

@@ -0,0 +1,21 @@
+defmodule Livebook.Proxy.Handler do
+  @moduledoc false
+
+  def child_spec(opts) do
+    name = Keyword.fetch!(opts, :name)
+    listen = Keyword.fetch!(opts, :listen)
+    :persistent_term.put({__MODULE__, name}, listen)
+    PartitionSupervisor.child_spec(child_spec: Task.Supervisor, name: name)
+  end
+
+  def serve(parent, name, data) when is_pid(parent) and is_atom(name) do
+    Process.link(parent)
+    ref = Process.monitor(parent)
+    conn = struct!(Plug.Conn, %{data | adapter: {Livebook.Proxy.Adapter, {parent, ref}}})
+    :persistent_term.get({__MODULE__, name}).(conn)
+  end
+
+  def get_pid(name, key) do
+    GenServer.whereis({:via, PartitionSupervisor, {name, key}})
+  end
+end

lib/livebook/proxy/server.ex

@@ -0,0 +1,78 @@
+defmodule Livebook.Proxy.Server do
+  @moduledoc false
+  import Plug.Conn
+
+  def serve(pid, name, %Plug.Conn{} = conn) when is_pid(pid) and is_atom(name) do
+    args = [self(), name, build_client_conn(conn)]
+    {:ok, spawn_pid} = Task.Supervisor.start_child(pid, Livebook.Proxy.Handler, :serve, args)
+    monitor_ref = Process.monitor(spawn_pid)
+    loop(monitor_ref, conn)
+  end
+
+  defp build_client_conn(conn) do
+    %{
+      adapter: nil,
+      host: conn.host,
+      method: conn.method,
+      owner: conn.owner,
+      port: conn.port,
+      remote_ip: conn.remote_ip,
+      query_string: conn.query_string,
+      path_info: conn.path_info,
+      scheme: conn.scheme,
+      script_name: conn.script_name,
+      req_headers: conn.req_headers
+    }
+  end
+
+  defp loop(monitor_ref, conn) do
+    receive do
+      {:send_resp, pid, ref, status, headers, body} ->
+        conn = send_resp(%{conn | resp_headers: headers}, status, body)
+        send(pid, {ref, :ok})
+        loop(monitor_ref, conn)
+
+      {:get_peer_data, pid, ref} ->
+        send(pid, {ref, get_peer_data(conn)})
+        loop(monitor_ref, conn)
+
+      {:get_http_protocol, pid, ref} ->
+        send(pid, {ref, get_http_protocol(conn)})
+        loop(monitor_ref, conn)
+
+      {:read_req_body, pid, ref, opts} ->
+        {message, conn} =
+          case read_body(conn, opts) do
+            {:ok, data, conn} -> {{:ok, data}, conn}
+            {:more, data, conn} -> {{:more, data}, conn}
+            {:error, _} = error -> {error, conn}
+          end
+
+        send(pid, {ref, message})
+        loop(monitor_ref, conn)
+
+      {:send_chunked, pid, ref, status, headers} ->
+        conn = send_chunked(%{conn | resp_headers: headers}, status)
+        send(pid, {ref, :ok})
+        loop(monitor_ref, conn)
+
+      {:chunk, pid, ref, chunk} ->
+        {message, conn} =
+          case chunk(conn, chunk) do
+            {:error, _} = error -> {error, conn}
+            result -> result
+          end
+
+        send(pid, {ref, message})
+        loop(monitor_ref, conn)
+
+      {:inform, pid, ref, status, headers} ->
+        conn = inform(conn, status, headers)
+        send(pid, {ref, :ok})
+        loop(monitor_ref, conn)
+
+      {:DOWN, ^monitor_ref, :process, _pid, reason} ->
+        {conn, reason}
+    end
+  end
+end

lib/livebook/runtime.ex

@@ -1107,6 +1107,6 @@ defprotocol Livebook.Runtime do
 
   TODO: document the communication here.
   """
-  @spec fetch_proxy_handler(t()) :: {:ok, pid()} | {:error, :not_found}
-  def fetch_proxy_handler(runtime)
+  @spec fetch_proxy_handler(t(), pid()) :: {:ok, pid()} | {:error, :not_found}
+  def fetch_proxy_handler(runtime, client_pid)
 end

lib/livebook/runtime/attached.ex

@@ -205,7 +205,7 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
     RuntimeServer.unregister_clients(runtime.server_pid, client_ids)
   end
 
-  def fetch_proxy_handler(runtime) do
-    RuntimeServer.fetch_proxy_handler(runtime.server_pid)
+  def fetch_proxy_handler(runtime, client_pid) do
+    RuntimeServer.fetch_proxy_handler(runtime.server_pid, client_pid)
   end
 end

lib/livebook/runtime/elixir_standalone.ex

@@ -324,7 +324,7 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do
     RuntimeServer.unregister_clients(runtime.server_pid, client_ids)
   end
 
-  def fetch_proxy_handler(runtime) do
-    RuntimeServer.fetch_proxy_handler(runtime.server_pid)
+  def fetch_proxy_handler(runtime, client_pid) do
+    RuntimeServer.fetch_proxy_handler(runtime.server_pid, client_pid)
   end
 end

lib/livebook/runtime/embedded.ex

@@ -171,8 +171,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do
     RuntimeServer.unregister_clients(runtime.server_pid, client_ids)
   end
 
-  def fetch_proxy_handler(runtime) do
-    RuntimeServer.fetch_proxy_handler(runtime.server_pid)
+  def fetch_proxy_handler(runtime, client_pid) do
+    RuntimeServer.fetch_proxy_handler(runtime.server_pid, client_pid)
   end
 
   defp config() do

lib/livebook/runtime/erl_dist.ex

@@ -41,7 +41,9 @@ defmodule Livebook.Runtime.ErlDist do
       Livebook.Runtime.ErlDist.IOForwardGL,
       Livebook.Runtime.ErlDist.LoggerGLHandler,
       Livebook.Runtime.ErlDist.Sink,
-      Livebook.Runtime.ErlDist.SmartCellGL
+      Livebook.Runtime.ErlDist.SmartCellGL,
+      Livebook.Proxy.Adapter,
+      Livebook.Proxy.Handler
     ]
   end
 

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -320,9 +320,9 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   @doc """
   Fetches the running Proxy Handler's pid from runtime.
   """
-  @spec fetch_proxy_handler(pid()) :: {:ok, pid()} | {:error, :not_found}
-  def fetch_proxy_handler(pid) do
-    GenServer.call(pid, :fetch_proxy_handler)
+  @spec fetch_proxy_handler(pid(), pid()) :: {:ok, pid()} | {:error, :not_found}
+  def fetch_proxy_handler(pid, client_pid) do
+    GenServer.call(pid, {:fetch_proxy_handler, client_pid})
   end
 
   @doc """
@@ -752,8 +752,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:reply, has_dependencies?, state}
   end
 
-  def handle_call(:fetch_proxy_handler, _from, state) do
-    if pid = GenServer.whereis(Kino.Proxy) do
+  def handle_call({:fetch_proxy_handler, client_pid}, _from, state) do
+    if pid = Livebook.Proxy.Handler.get_pid(Kino.Proxy, client_pid) do
       {:reply, {:ok, pid}, state}
     else
       {:reply, {:error, :not_found}, state}

lib/livebook/runtime/evaluator/io_proxy.ex

@@ -382,6 +382,11 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
     {result, state}
   end
 
+  defp io_request({:livebook_get_proxy_handler_child_spec, fun}, state) do
+    result = {Livebook.Proxy.Handler, name: Kino.Proxy, listen: fun}
+    {result, state}
+  end
+
   defp io_request(_, state) do
     {{:error, :request}, state}
   end

lib/livebook/session.ex

@@ -627,9 +627,9 @@ defmodule Livebook.Session do
   @doc """
   Fetches the running Proxy Handler's pid from runtime.
   """
-  @spec fetch_proxy_handler(pid()) :: {:ok, pid()} | {:error, :not_found | :disconnected}
-  def fetch_proxy_handler(pid) do
-    GenServer.call(pid, :fetch_proxy_handler)
+  @spec fetch_proxy_handler(pid(), pid()) :: {:ok, pid()} | {:error, :not_found | :disconnected}
+  def fetch_proxy_handler(pid, client_pid) do
+    GenServer.call(pid, {:fetch_proxy_handler, client_pid})
   end
 
   @doc """
@@ -1081,9 +1081,9 @@ defmodule Livebook.Session do
     {:noreply, state}
   end
 
-  def handle_call(:fetch_proxy_handler, _from, state) do
+  def handle_call({:fetch_proxy_handler, client_pid}, _from, state) do
     if Runtime.connected?(state.data.runtime) do
-      {:reply, Runtime.fetch_proxy_handler(state.data.runtime), state}
+      {:reply, Runtime.fetch_proxy_handler(state.data.runtime, client_pid), state}
     else
       {:reply, {:error, :disconnected}, state}
     end

lib/livebook_web/plugs/proxy_plug.ex

@@ -12,7 +12,7 @@ defmodule LivebookWeb.ProxyPlug do
     session = fetch_session!(id)
     pid = fetch_proxy_handler!(session)
     conn = prepare_conn(conn, path_info, ["sessions", id, "proxy"])
-    {conn, _} = Kino.Proxy.serve(pid, conn)
+    {conn, _} = Livebook.Proxy.Server.serve(pid, Kino.Proxy, conn)
 
     halt(conn)
   end
@@ -27,7 +27,7 @@ defmodule LivebookWeb.ProxyPlug do
     session = fetch_session!(id)
     pid = fetch_proxy_handler!(session)
     conn = prepare_conn(conn, path_info, ["apps", slug, "proxy"])
-    {conn, _} = Kino.Proxy.serve(pid, conn)
+    {conn, _} = Livebook.Proxy.Server.serve(pid, Kino.Proxy, conn)
 
     halt(conn)
   end
@@ -51,7 +51,7 @@ defmodule LivebookWeb.ProxyPlug do
   end
 
   defp fetch_proxy_handler!(session) do
-    case Livebook.Session.fetch_proxy_handler(session.pid) do
+    case Livebook.Session.fetch_proxy_handler(session.pid, self()) do
       {:ok, pid} -> pid
       {:error, _} -> raise NotFoundError, "could not find a kino proxy running"
     end

mix.exs

@@ -116,7 +116,6 @@ defmodule Livebook.MixProject do
       {:mint_web_socket, "~> 1.0.0"},
       {:protobuf, "~> 0.12.0"},
       {:dns_cluster, "~> 0.1.2"},
-      {:kino_proxy, kino_proxy_opts()},
       {:phoenix_live_reload, "~> 1.2", only: :dev},
       {:floki, ">= 0.27.0", only: :test},
       {:bypass, "~> 2.1", only: :test},
@@ -128,14 +127,6 @@ defmodule Livebook.MixProject do
     ]
   end
 
-  defp kino_proxy_opts do
-    if path = System.get_env("KINO_PROXY_PATH") do
-      [path: path]
-    else
-      [github: "livebook-dev/kino_proxy"]
-    end
-  end
-
   defp target_deps(:app), do: [{:elixirkit, path: "elixirkit"}]
   defp target_deps(_), do: []
 

mix.lock

@@ -21,7 +21,6 @@
   "jason": {:hex, :jason, "1.4.1", "af1504e35f629ddcdd6addb3513c3853991f694921b1b9368b0bd32beb9f1b63", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "fbb01ecdfd565b56261302f7e1fcc27c4fb8f32d56eab74db621fc154604a7a1"},
   "jose": {:hex, :jose, "1.11.10", "a903f5227417bd2a08c8a00a0cbcc458118be84480955e8d251297a425723f83", [:mix, :rebar3], [], "hexpm", "0d6cd36ff8ba174db29148fc112b5842186b68a90ce9fc2b3ec3afe76593e614"},
   "jsx": {:hex, :jsx, "3.0.0", "20a170abd4335fc6db24d5fad1e5d677c55dadf83d1b20a8a33b5fe159892a39", [:rebar3], [], "hexpm", "37beca0435f5ca8a2f45f76a46211e76418fbef80c36f0361c249fc75059dc6d"},
-  "kino_proxy": {:git, "https://github.com/livebook-dev/kino_proxy.git", "42c434450d97bf2d2035c42296317c4955873bac", []},
   "makeup": {:hex, :makeup, "1.1.1", "fa0bc768698053b2b3869fa8a62616501ff9d11a562f3ce39580d60860c3a55e", [:mix], [{:nimble_parsec, "~> 1.2.2 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "5dc62fbdd0de44de194898b6710692490be74baa02d9d108bc29f007783b0b48"},
   "makeup_elixir": {:hex, :makeup_elixir, "0.16.2", "627e84b8e8bf22e60a2579dad15067c755531fea049ae26ef1020cad58fe9578", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "41193978704763f6bbe6cc2758b84909e62984c7752b3784bd3c218bb341706b"},
   "makeup_erlang": {:hex, :makeup_erlang, "1.0.0", "6f0eff9c9c489f26b69b61440bf1b238d95badae49adac77973cbacae87e3c2e", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "ea7a9307de9d1548d2a72d299058d1fd2339e3d398560a0e46c27dab4891e4d2"},

test/livebook/proxy_test.exs

@@ -1,4 +1,4 @@
-defmodule LivebookWeb.ProxyPlugTest do
+defmodule Livebook.ProxyTest do
   use LivebookWeb.ConnCase, async: true
 
   require Phoenix.LiveViewTest
@@ -24,30 +24,17 @@ defmodule LivebookWeb.ProxyPlugTest do
     end
 
     test "returns the proxied response defined in notebook", %{conn: conn} do
-      cell = %{
-        Notebook.Cell.new(:code)
-        | source: """
-          Kino.Proxy.listen(fn conn ->
-            conn
-            |> Plug.Conn.put_resp_header("content-type", "application/text;charset=utf-8")
-            |> Plug.Conn.send_resp(200, "used " <> conn.method <> " method")
-          end)
-          """
-      }
-
-      cell_id = cell.id
-
-      section = %{Notebook.Section.new() | cells: [cell]}
-      notebook = %{Notebook.new() | sections: [section]}
-
+      %{sections: [%{cells: [%{id: cell_id}]}]} = notebook = proxy_notebook()
       {:ok, session} = Sessions.create_session(notebook: notebook)
       {:ok, runtime} = Runtime.Embedded.new() |> Runtime.connect()
 
       Session.set_runtime(session.pid, runtime)
       Session.subscribe(session.id)
-
       Session.queue_cell_evaluation(session.pid, cell_id)
-      assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _, _}}
+
+      assert_receive {:operation,
+                      {:add_cell_evaluation_response, _, ^cell_id, _, %{errored: false}}},
+                     4_000
 
       url = "/sessions/#{session.id}/proxy/"
 
@@ -73,35 +60,16 @@ defmodule LivebookWeb.ProxyPlugTest do
 
     test "returns the proxied response defined in notebook", %{conn: conn} do
       slug = Livebook.Utils.random_short_id()
-
-      cell = %{
-        Notebook.Cell.new(:code)
-        | source: """
-          Kino.Proxy.listen(fn conn ->
-            conn
-            |> Plug.Conn.put_resp_header("content-type", "application/text;charset=utf-8")
-            |> Plug.Conn.send_resp(200, "used " <> conn.method <> " method")
-          end)
-          """
-      }
-
       app_settings = %{Notebook.AppSettings.new() | slug: slug, access_type: :public}
-      section = %{Notebook.Section.new() | cells: [cell]}
-      notebook = %{Notebook.new() | app_settings: app_settings, sections: [section]}
+      notebook = %{proxy_notebook() | app_settings: app_settings}
 
       Livebook.Apps.subscribe()
       pid = deploy_notebook_sync(notebook)
 
-      assert_receive {:app_created, %{pid: ^pid, slug: ^slug}}
+      assert_receive {:app_created, %{pid: ^pid, slug: ^slug, sessions: []}}
 
       assert_receive {:app_updated,
-                      %{
-                        pid: ^pid,
-                        slug: ^slug,
-                        sessions: [
-                          %{id: id, app_status: %{lifecycle: :active, execution: :executed}}
-                        ]
-                      }}
+                      %{slug: ^slug, sessions: [%{id: id, app_status: %{execution: :executed}}]}}
 
       url = "/apps/#{slug}/#{id}/proxy/"
 
@@ -110,8 +78,34 @@ defmodule LivebookWeb.ProxyPlugTest do
       assert text_response(put(conn, url), 200) == "used PUT method"
       assert text_response(patch(conn, url), 200) == "used PATCH method"
       assert text_response(delete(conn, url), 200) == "used DELETE method"
-
-      Livebook.App.close(pid)
     end
   end
+
+  defp proxy_notebook() do
+    cell =
+      %{
+        Notebook.Cell.new(:code)
+        | source: """
+          fun = fn conn ->
+            conn
+            |> Plug.Conn.put_resp_header("content-type", "application/text;charset=utf-8")
+            |> Plug.Conn.send_resp(200, "used " <> conn.method <> " method")
+          end
+
+          ref = make_ref()
+          request = {:livebook_get_proxy_handler_child_spec, fun}
+          send(Process.group_leader(), {:io_request, self(), ref, request})
+
+          child_spec =
+            receive do
+              {:io_reply, ^ref, child_spec} -> child_spec
+            end
+
+          Supervisor.start_link([child_spec], strategy: :one_for_one)\
+          """
+      }
+
+    section = %{Notebook.Section.new() | cells: [cell]}
+    %{Notebook.new() | sections: [section]}
+  end
 end

test/support/noop_runtime.ex

@@ -73,7 +73,7 @@ defmodule Livebook.Runtime.NoopRuntime do
 
     def register_clients(_, _), do: :ok
     def unregister_clients(_, _), do: :ok
-    def fetch_proxy_handler(_), do: {:error, :not_found}
+    def fetch_proxy_handler(_, _), do: {:error, :not_found}
 
     defp trace(runtime, fun, args) do
       if runtime.trace_to do