Fix session assets resolution in distributed deployments (#2611)

lib/livebook/utils.ex

@@ -1,7 +1,7 @@
 defmodule Livebook.Utils do
   require Logger
 
-  @type id :: binary()
+  @type id :: String.t()
 
   @doc """
   Generates a random binary id.
@@ -82,11 +82,42 @@ defmodule Livebook.Utils do
     case Base.decode32(id, case: :lower) do
       {:ok,
        <<boot_id::binary-size(3), node_part::binary-size(16), _random_part::binary-size(11)>>} ->
-        known_nodes = [node() | Node.list()]
+        with {:ok, node} <- fetch_node_by_hash(node_part) do
+          {:ok, node, boot_id}
+        end
 
-        Enum.find_value(known_nodes, :error, fn node ->
-          node_hash(node) == node_part && {:ok, node, boot_id}
-        end)
+      _ ->
+        :error
+    end
+  end
+
+  defp fetch_node_by_hash(node_hash) do
+    known_nodes = [node() | Node.list()]
+
+    Enum.find_value(known_nodes, :error, fn node ->
+      node_hash(node) == node_hash && {:ok, node}
+    end)
+  end
+
+  @doc """
+  Returns a determinsitic short id corresponding to the current node.
+  """
+  @spec node_id() :: String.t()
+  def node_id() do
+    node_hash = node_hash(node())
+    Base.encode32(node_hash, case: :lower, padding: false)
+  end
+
+  @doc """
+  Extracts node name from the given node id, generated with `node_id/0`.
+
+  The node in question must be connected, otherwise it won't be found.
+  """
+  @spec node_from_id(id()) :: {:ok, node()} | :error
+  def node_from_id(id) do
+    case Base.decode32(id, case: :lower, padding: false) do
+      {:ok, <<node_hash::binary-size(16)>>} ->
+        fetch_node_by_hash(node_hash)
 
       _ ->
         :error

lib/livebook_web/controllers/session_controller.ex

@@ -117,43 +117,56 @@ defmodule LivebookWeb.SessionController do
     # The request comes from a cross-origin iframe
     conn = allow_cors(conn)
 
-    # This route include session id, while we want the browser to
-    # cache assets across sessions, so we only ensure the asset
-    # is available and redirect to the corresponding route without
-    # session id
+    # This route include session id, but we want the browser to cache
+    # assets across sessions, so we only ensure the asset is available
+    # on this node and redirect to the corresponding route with node
+    # id, rather than session id
     if ensure_asset?(id, hash, asset_path) do
+      node_id = Livebook.Utils.node_id()
+
       conn
       |> cache_permanently()
       |> put_status(:moved_permanently)
-      |> redirect(to: ~p"/public/sessions/assets/#{hash}/#{file_parts}")
+      |> redirect(to: ~p"/public/sessions/node/#{node_id}/assets/#{hash}/#{file_parts}")
     else
       send_resp(conn, 404, "Not found")
     end
   end
 
-  def show_cached_asset(conn, %{"hash" => hash, "file_parts" => file_parts}) do
+  def show_cached_asset(conn, %{"node_id" => node_id, "hash" => hash, "file_parts" => file_parts}) do
     asset_path = Path.join(file_parts)
 
     # The request comes from a cross-origin iframe
     conn = allow_cors(conn)
 
-    case lookup_asset(hash, asset_path) do
-      {:ok, local_asset_path} ->
-        conn =
-          conn
-          |> put_content_type(asset_path)
-          |> cache_permanently()
+    gzip_result =
+      if accept_encoding?(conn, "gzip") do
+        with {:ok, local_asset_path} <-
+               lookup_asset_or_transfer(hash, asset_path <> ".gz", node_id) do
+          conn =
+            conn
+            |> put_resp_header("content-encoding", "gzip")
+            |> put_resp_header("vary", "Accept-Encoding")
 
-        local_asset_path_gz = local_asset_path <> ".gz"
-
-        if accept_encoding?(conn, "gzip") and File.exists?(local_asset_path_gz) do
-          conn
-          |> put_resp_header("content-encoding", "gzip")
-          |> put_resp_header("vary", "Accept-Encoding")
-          |> send_file(200, local_asset_path_gz)
-        else
-          send_file(conn, 200, local_asset_path)
+          {:ok, local_asset_path, conn}
         end
+      else
+        :error
+      end
+
+    result =
+      with :error <- gzip_result do
+        with {:ok, local_asset_path} <- lookup_asset_or_transfer(hash, asset_path, node_id) do
+          {:ok, local_asset_path, conn}
+        end
+      end
+
+    case result do
+      {:ok, local_asset_path, conn} ->
+        conn
+        |> put_content_type(asset_path)
+        |> cache_permanently()
+        |> send_file(200, local_asset_path)
 
       :error ->
         send_resp(conn, 404, "Not found")
@@ -268,7 +281,8 @@ defmodule LivebookWeb.SessionController do
     end
   end
 
-  defp lookup_asset(hash, asset_path) do
+  @doc false
+  def lookup_asset(hash, asset_path) do
     with {:ok, local_asset_path} <- Session.local_asset_path(hash, asset_path),
          true <- File.exists?(local_asset_path) do
       {:ok, local_asset_path}
@@ -277,6 +291,24 @@ defmodule LivebookWeb.SessionController do
     end
   end
 
+  defp lookup_asset_or_transfer(hash, asset_path, node_id) do
+    with :error <- lookup_asset(hash, asset_path),
+         {:ok, node} <- Livebook.Utils.node_from_id(node_id),
+         {:ok, remote_asset_path} <-
+           :erpc.call(node, __MODULE__, :lookup_asset, [hash, asset_path]),
+         {:ok, local_asset_path} <- Session.local_asset_path(hash, asset_path) do
+      transfer_file!(node, remote_asset_path, local_asset_path)
+      {:ok, local_asset_path}
+    end
+  end
+
+  defp transfer_file!(remote_node, remote_path, local_path) do
+    File.mkdir_p!(Path.dirname(local_path))
+    remote_stream = :erpc.call(remote_node, File, :stream!, [remote_path, 2048, []])
+    local_stream = File.stream!(local_path)
+    Enum.into(remote_stream, local_stream)
+  end
+
   defp allow_cors(conn) do
     put_resp_header(conn, "access-control-allow-origin", "*")
   end

lib/livebook_web/router.ex

@@ -52,8 +52,8 @@ defmodule LivebookWeb.Router do
   scope "/public", LivebookWeb do
     pipe_through [:js_view_assets]
 
-    get "/sessions/assets/:hash/*file_parts", SessionController, :show_cached_asset
     get "/sessions/:id/assets/:hash/*file_parts", SessionController, :show_asset
+    get "/sessions/node/:node_id/assets/:hash/*file_parts", SessionController, :show_cached_asset
   end
 
   live_session :default,

test/livebook_web/controllers/session_controller_test.exs

@@ -255,7 +255,10 @@ defmodule LivebookWeb.SessionControllerTest do
 
       conn = start_session_and_request_asset(conn, notebook, hash)
 
-      assert redirected_to(conn, 301) == ~p"/public/sessions/assets/#{hash}/main.js"
+      node_id = Livebook.Utils.node_id()
+
+      assert redirected_to(conn, 301) ==
+               ~p"/public/sessions/node/#{node_id}/assets/#{hash}/main.js"
 
       {:ok, asset_path} = Session.local_asset_path(hash, "main.js")
       assert File.exists?(asset_path)
@@ -268,7 +271,10 @@ defmodule LivebookWeb.SessionControllerTest do
 
       conn = start_session_and_request_asset(conn, notebook, hash)
 
-      assert redirected_to(conn, 301) == ~p"/public/sessions/assets/#{hash}/main.js"
+      node_id = Livebook.Utils.node_id()
+
+      assert redirected_to(conn, 301) ==
+               ~p"/public/sessions/node/#{node_id}/assets/#{hash}/main.js"
 
       assert File.exists?(Path.join(assets_path, "main.js"))
     end
@@ -283,7 +289,10 @@ defmodule LivebookWeb.SessionControllerTest do
 
       conn = get(conn, ~p"/public/sessions/#{random_session_id}/assets/#{hash}/main.js")
 
-      assert redirected_to(conn, 301) == ~p"/public/sessions/assets/#{hash}/main.js"
+      node_id = Livebook.Utils.node_id()
+
+      assert redirected_to(conn, 301) ==
+               ~p"/public/sessions/node/#{node_id}/assets/#{hash}/main.js"
     end
   end
 
@@ -291,7 +300,8 @@ defmodule LivebookWeb.SessionControllerTest do
     test "returns not found when no matching assets are in the cache", %{conn: conn} do
       %{notebook: _notebook, hash: hash} = notebook_with_js_output()
 
-      conn = get(conn, ~p"/public/sessions/assets/#{hash}/main.js")
+      node_id = Livebook.Utils.node_id()
+      conn = get(conn, ~p"/public/sessions/node/#{node_id}/assets/#{hash}/main.js")
 
       assert conn.status == 404
       assert conn.resp_body == "Not found"
@@ -302,7 +312,8 @@ defmodule LivebookWeb.SessionControllerTest do
       # Fetch the assets for the first time
       start_session_and_request_asset(conn, notebook, hash)
 
-      conn = get(conn, ~p"/public/sessions/assets/#{hash}/main.js")
+      node_id = Livebook.Utils.node_id()
+      conn = get(conn, ~p"/public/sessions/node/#{node_id}/assets/#{hash}/main.js")
 
       assert conn.status == 200
       assert "export function init(" <> _ = conn.resp_body
@@ -313,10 +324,12 @@ defmodule LivebookWeb.SessionControllerTest do
 
       start_session_and_request_asset(conn, notebook, hash)
 
+      node_id = Livebook.Utils.node_id()
+
       conn =
         conn
         |> put_req_header("accept-encoding", "gzip")
-        |> get(~p"/public/sessions/assets/#{hash}/main.js")
+        |> get(~p"/public/sessions/node/#{node_id}/assets/#{hash}/main.js")
 
       assert conn.status == 200
       assert "export function init(" <> _ = :zlib.gunzip(conn.resp_body)