Add evaluation temporary directory (#2204)

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -694,6 +694,12 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     end
   end
 
+  defp evaluator_tmp_dir(state) do
+    if tmp_dir = state.tmp_dir do
+      Path.join(tmp_dir, "tmp")
+    end
+  end
+
   defp ensure_evaluator(state, container_ref) do
     if Map.has_key?(state.evaluators, container_ref) do
       state
@@ -705,6 +711,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
           runtime_broadcast_to: state.runtime_broadcast_to,
           object_tracker: state.object_tracker,
           ebin_path: state.ebin_path,
+          tmp_dir: evaluator_tmp_dir(state),
           io_proxy_registry: state.io_proxy_registry
         )
 

lib/livebook/runtime/evaluator.ex

@@ -83,6 +83,9 @@ defmodule Livebook.Runtime.Evaluator do
     * `:ebin_path` - a directory to write modules bytecode into. When
       not specified, modules are not written to disk
 
+    * `:tmp_dir` - a temporary directory for arbitrary use during
+      evaluation
+
     * `:io_proxy_registry` - the registry to register IO proxy
       processes in
 
@@ -266,6 +269,7 @@ defmodule Livebook.Runtime.Evaluator do
     runtime_broadcast_to = Keyword.get(opts, :runtime_broadcast_to, send_to)
     object_tracker = Keyword.fetch!(opts, :object_tracker)
     ebin_path = Keyword.get(opts, :ebin_path)
+    tmp_dir = Keyword.get(opts, :tmp_dir)
     io_proxy_registry = Keyword.get(opts, :io_proxy_registry)
 
     {:ok, io_proxy} =
@@ -275,6 +279,7 @@ defmodule Livebook.Runtime.Evaluator do
         runtime_broadcast_to,
         object_tracker,
         ebin_path,
+        tmp_dir,
         io_proxy_registry
       )
 

lib/livebook/runtime/evaluator/io_proxy.ex

@@ -26,23 +26,54 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
   For all supported requests a message is sent to the configured
   `:send_to` process, so this device serves as a proxy.
   """
-  @spec start(pid(), pid(), pid(), pid(), String.t() | nil, atom() | nil) :: GenServer.on_start()
-  def start(evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry) do
+  @spec start(
+          pid(),
+          pid(),
+          pid(),
+          pid(),
+          String.t() | nil,
+          String.t() | nil,
+          atom() | nil
+        ) :: GenServer.on_start()
+  def start(
+        evaluator,
+        send_to,
+        runtime_broadcast_to,
+        object_tracker,
+        ebin_path,
+        tmp_dir,
+        registry
+      ) do
     GenServer.start(
       __MODULE__,
-      {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry}
+      {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, tmp_dir, registry}
     )
   end
 
   @doc """
   Linking version of `start/4`.
   """
-  @spec start_link(pid(), pid(), pid(), pid(), String.t() | nil, atom() | nil) ::
-          GenServer.on_start()
-  def start_link(evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry) do
+  @spec start_link(
+          pid(),
+          pid(),
+          pid(),
+          pid(),
+          String.t() | nil,
+          String.t() | nil,
+          atom() | nil
+        ) :: GenServer.on_start()
+  def start_link(
+        evaluator,
+        send_to,
+        runtime_broadcast_to,
+        object_tracker,
+        ebin_path,
+        tmp_dir,
+        registry
+      ) do
     GenServer.start_link(
       __MODULE__,
-      {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry}
+      {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, tmp_dir, registry}
     )
   end
 
@@ -73,7 +104,9 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
   end
 
   @impl true
-  def init({evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, registry}) do
+  def init(
+        {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path, tmp_dir, registry}
+      ) do
     evaluator_monitor = Process.monitor(evaluator)
 
     if registry do
@@ -94,6 +127,7 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
        runtime_broadcast_to: runtime_broadcast_to,
        object_tracker: object_tracker,
        ebin_path: ebin_path,
+       tmp_dir: tmp_dir,
        tracer_info: %Evaluator.Tracer{},
        modules_defined: MapSet.new()
      }}
@@ -348,6 +382,18 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
     {result, state}
   end
 
+  defp io_request(:livebook_get_tmp_dir, state) do
+    result =
+      with tmp_dir when is_binary(tmp_dir) <- state.tmp_dir,
+           :ok <- File.mkdir_p(tmp_dir) do
+        {:ok, state.tmp_dir}
+      else
+        _ -> {:error, :not_available}
+      end
+
+    {result, state}
+  end
+
   defp io_request(_, state) do
     {{:error, :request}, state}
   end