Improve notebook file locking to work across nodes (#675)

* Improve notebook file locking to work across nodes

* Add node check for local file system opreations

* Replace node with host id

* Refactor process down cleanup

* Scope local file system with node

* local? -> type

lib/livebook/file_system.ex

@@ -28,6 +28,26 @@ defprotocol Livebook.FileSystem do
 
   @type access :: :read | :write | :read_write | :none
 
+  @doc """
+  Returns a term uniquely identifying the resource used as a file
+  system.
+  """
+  @spec resource_identifier(t()) :: term()
+  def resource_identifier(file_system)
+
+  @doc """
+  Returns the file system type.
+
+  Based on the underlying resource, the type can be either:
+
+    * `:local` - if the resource is local to its node
+
+    * `:global` - if the resource is external and available
+      accessible from any node
+  """
+  @spec type(t()) :: :local | :global
+  def type(file_system)
+
   @doc """
   Returns the default directory path.
 

lib/livebook/file_system/file.ex

@@ -52,6 +52,23 @@ defmodule Livebook.FileSystem.File do
     new(FileSystem.Local.new(), path)
   end
 
+  @doc """
+  Returns a term uniquely identifying the file together
+  with its file system.
+  """
+  @spec resource_identifier(t()) :: term()
+  def resource_identifier(file) do
+    {FileSystem.resource_identifier(file.file_system), file.path}
+  end
+
+  @doc """
+  Checks if the given file is within a file system local to its node.
+  """
+  @spec local?(t()) :: term()
+  def local?(file) do
+    FileSystem.type(file.file_system) == :local
+  end
+
   @doc """
   Returns a new file resulting from resolving `subject`
   against `file`.

lib/livebook/file_system/local.ex

@@ -3,11 +3,16 @@ defmodule Livebook.FileSystem.Local do
 
   # File system backed by local disk.
 
-  defstruct [:default_path]
+  defstruct [:origin_pid, :default_path]
 
   alias Livebook.FileSystem
 
   @type t :: %__MODULE__{
+          # We cannot just store the node, because when the struct is
+          # built, we may not yet be in distributed mode. Instead, we
+          # keep the pid of whatever process created this file system
+          # and we call node/1 on it whenever needed
+          origin_pid: pid(),
           default_path: FileSystem.path()
         }
 
@@ -28,13 +33,21 @@ defmodule Livebook.FileSystem.Local do
 
     FileSystem.Utils.assert_dir_path!(default_path)
 
-    %__MODULE__{default_path: default_path}
+    %__MODULE__{origin_pid: self(), default_path: default_path}
   end
 end
 
 defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
   alias Livebook.FileSystem
 
+  def resource_identifier(file_system) do
+    {:local_file_system, node(file_system.origin_pid)}
+  end
+
+  def type(_file_system) do
+    :local
+  end
+
   def default_path(file_system) do
     file_system.default_path
   end
@@ -42,6 +55,7 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
   def list(file_system, path, recursive) do
     FileSystem.Utils.assert_dir_path!(path)
 
+    with :ok <- ensure_local(file_system) do
       case File.ls(path) do
         {:ok, filenames} ->
           paths =
@@ -68,21 +82,25 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
           FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def read(_file_system, path) do
+  def read(file_system, path) do
     FileSystem.Utils.assert_regular_path!(path)
 
+    with :ok <- ensure_local(file_system) do
       case File.read(path) do
         {:ok, binary} -> {:ok, binary}
         {:error, error} -> FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def write(_file_system, path, content) do
+  def write(file_system, path, content) do
     FileSystem.Utils.assert_regular_path!(path)
 
     dir = Path.dirname(path)
 
+    with :ok <- ensure_local(file_system) do
       with :ok <- File.mkdir_p(dir),
            :ok <- File.write(path, content) do
         :ok
@@ -90,35 +108,43 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
         {:error, error} -> FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def access(_file_system, path) do
+  def access(file_system, path) do
+    with :ok <- ensure_local(file_system) do
       case File.stat(path) do
         {:ok, stat} -> {:ok, stat.access}
         {:error, error} -> FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def create_dir(_file_system, path) do
+  def create_dir(file_system, path) do
     FileSystem.Utils.assert_dir_path!(path)
 
+    with :ok <- ensure_local(file_system) do
       case File.mkdir_p(path) do
         :ok -> :ok
         {:error, error} -> FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def remove(_file_system, path) do
+  def remove(file_system, path) do
+    with :ok <- ensure_local(file_system) do
       case File.rm_rf(path) do
         {:ok, _paths} -> :ok
         {:error, error, _paths} -> FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def copy(_file_system, source_path, destination_path) do
+  def copy(file_system, source_path, destination_path) do
     FileSystem.Utils.assert_same_type!(source_path, destination_path)
 
     containing_dir = Path.dirname(destination_path)
 
+    with :ok <- ensure_local(file_system) do
       case File.mkdir_p(containing_dir) do
         :ok ->
           case File.cp_r(source_path, destination_path) do
@@ -130,10 +156,12 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
           FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def rename(_file_system, source_path, destination_path) do
+  def rename(file_system, source_path, destination_path) do
     FileSystem.Utils.assert_same_type!(source_path, destination_path)
 
+    with :ok <- ensure_local(file_system) do
       if File.exists?(destination_path) do
         FileSystem.Utils.posix_error(:eexist)
       else
@@ -148,8 +176,10 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
         end
       end
     end
+  end
 
-  def etag_for(_file_system, path) do
+  def etag_for(file_system, path) do
+    with :ok <- ensure_local(file_system) do
       case File.stat(path) do
         {:ok, stat} ->
           %{size: size, mtime: mtime} = stat
@@ -161,18 +191,22 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
           FileSystem.Utils.posix_error(error)
       end
     end
+  end
 
-  def exists?(_file_system, path) do
+  def exists?(file_system, path) do
+    with :ok <- ensure_local(file_system) do
       if FileSystem.Utils.dir_path?(path) do
         {:ok, File.dir?(path)}
       else
         {:ok, File.exists?(path)}
       end
     end
+  end
 
-  def resolve_path(_file_system, dir_path, subject) do
+  def resolve_path(file_system, dir_path, subject) do
     FileSystem.Utils.assert_dir_path!(dir_path)
 
+    with :ok <- ensure_local(file_system) do
       if subject == "" do
         dir_path
       else
@@ -187,3 +221,12 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.Local do
       end
     end
   end
+
+  defp ensure_local(file_system) do
+    if node(file_system.origin_pid) == node() do
+      :ok
+    else
+      {:error, "this local file system belongs to a different host"}
+    end
+  end
+end

lib/livebook/file_system/s3.ex

@@ -59,6 +59,14 @@ defimpl Livebook.FileSystem, for: Livebook.FileSystem.S3 do
   alias Livebook.Utils.HTTP
   alias Livebook.FileSystem.S3.XML
 
+  def resource_identifier(file_system) do
+    {:s3, file_system.bucket_url}
+  end
+
+  def type(_file_system) do
+    :global
+  end
+
   def default_path(_file_system) do
     "/"
   end

lib/livebook/session/file_guard.ex

@@ -12,7 +12,7 @@ defmodule Livebook.Session.FileGuard do
   alias Livebook.FileSystem
 
   @type state :: %{
-          file_with_owner_ref: %{FileSystem.File.t() => reference()}
+          files: %{term() => {FileSystem.File.t(), owner_pid :: pid(), reference()}}
         }
 
   @name __MODULE__
@@ -47,32 +47,52 @@ defmodule Livebook.Session.FileGuard do
 
   @impl true
   def init(_opts) do
-    {:ok, %{file_with_owner_ref: %{}}}
+    {:ok, %{files: %{}}}
   end
 
   @impl true
   def handle_call({:lock, file, owner_pid}, _from, state) do
-    if Map.has_key?(state.file_with_owner_ref, file) do
+    file_id = FileSystem.File.resource_identifier(file)
+
+    if Map.has_key?(state.files, file_id) or lock_globally(file, file_id, owner_pid) == false do
       {:reply, {:error, :already_in_use}, state}
     else
       monitor_ref = Process.monitor(owner_pid)
-      state = put_in(state.file_with_owner_ref[file], monitor_ref)
+      state = put_in(state.files[file_id], {file, owner_pid, monitor_ref})
       {:reply, :ok, state}
     end
   end
 
   @impl true
   def handle_cast({:unlock, file}, state) do
-    {maybe_ref, state} = pop_in(state.file_with_owner_ref[file])
+    file_id = FileSystem.File.resource_identifier(file)
-    maybe_ref && Process.demonitor(maybe_ref, [:flush])
+
+    {maybe_file, state} = pop_in(state.files[file_id])
+
+    with {file, owner_pid, monitor_ref} <- maybe_file do
+      unlock_globally(file, file_id, owner_pid)
+      Process.demonitor(monitor_ref, [:flush])
+    end
 
     {:noreply, state}
   end
 
   @impl true
   def handle_info({:DOWN, ref, :process, _, _}, state) do
-    {file, ^ref} = Enum.find(state.file_with_owner_ref, &(elem(&1, 1) == ref))
+    {file_id, {file, owner_pid, ^ref}} =
-    {_, state} = pop_in(state.file_with_owner_ref[file])
+      Enum.find(state.files, &match?({_file_id, {_file, _owner_pid, ^ref}}, &1))
+
+    unlock_globally(file, file_id, owner_pid)
+    {_, state} = pop_in(state.files[file_id])
+
     {:noreply, state}
   end
+
+  defp lock_globally(file, file_id, owner_pid) do
+    FileSystem.File.local?(file) or :global.set_lock({file_id, owner_pid})
+  end
+
+  defp unlock_globally(file, file_id, owner_pid) do
+    FileSystem.File.local?(file) or :global.del_lock({file_id, owner_pid})
+  end
 end

test/livebook/session/file_guard_test.exs

@@ -2,21 +2,40 @@ defmodule Livebook.Session.FileGuardTest do
   use ExUnit.Case, async: false
 
   alias Livebook.Session.FileGuard
+  alias Livebook.FileSystem
 
-  test "lock/2 returns an error if the given path is already locked" do
+  test "lock/2 returns an error if the given file is already locked" do
-    assert :ok = FileGuard.lock("/some/path", self())
+    file = FileSystem.File.local("/some/path")
-    assert {:error, :already_in_use} = FileGuard.lock("/some/path", self())
+
+    assert :ok = FileGuard.lock(file, self())
+    assert {:error, :already_in_use} = FileGuard.lock(file, self())
   end
 
-  test "unlock/1 unlocks the given path" do
+  test "lock/2 is agnostic to irrelevant file system configuration" do
-    assert :ok = FileGuard.lock("/some/path", self())
+    fs1 = FileSystem.Local.new(default_path: "/path/1/")
-    :ok = FileGuard.unlock("/some/path")
+    fs2 = FileSystem.Local.new(default_path: "/path/2/")
-    assert :ok = FileGuard.lock("/some/path", self())
+
+    # The file system has different configuration, but it's the same resource
+    file1 = FileSystem.File.new(fs1, "/some/path")
+    file2 = FileSystem.File.new(fs2, "/some/path")
+
+    assert :ok = FileGuard.lock(file1, self())
+    assert {:error, :already_in_use} = FileGuard.lock(file2, self())
   end
 
-  test "path is automatically unloacked when the owner process termiantes" do
+  test "unlock/1 unlocks the given file" do
+    file = FileSystem.File.local("/some/path")
+
+    assert :ok = FileGuard.lock(file, self())
+    :ok = FileGuard.unlock(file)
+    assert :ok = FileGuard.lock(file, self())
+  end
+
+  test "file is automatically unloacked when the owner process termiantes" do
+    file = FileSystem.File.local("/some/path")
+
     owner = spawn(fn -> :ok end)
-    :ok = FileGuard.lock("/some/path", owner)
+    :ok = FileGuard.lock(file, owner)
-    assert :ok = FileGuard.lock("/some/path", self())
+    assert :ok = FileGuard.lock(file, self())
   end
 end