ETS config file persistence (#1002)

* Added Ets config storage file persistence

* Adjusted to review

* Adjusted to review

* Removed redundant code

* Update lib/livebook/storage/ets.ex

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

config/test.exs

@@ -14,7 +14,14 @@ config :logger, level: :warn
 # Disable authentication mode during test
 config :livebook, :authentication_mode, :disabled
 
-config :livebook, :data_path, Path.expand("tmp/livebook_data/test")
+data_path = Path.expand("tmp/livebook_data/test")
+
+# Clear data path for tests
+if File.exists?(data_path) do
+  File.rm_rf!(data_path)
+end
+
+config :livebook, :data_path, data_path
 
 # Use the embedded runtime in tests by default, so they
 # are cheaper to run. Other runtimes can be tested by starting

lib/livebook/storage/ets.ex

@@ -10,13 +10,11 @@ defmodule Livebook.Storage.Ets do
   """
   @behaviour Livebook.Storage
 
-  @table_name __MODULE__
-
   use GenServer
 
   @impl Livebook.Storage
   def all(namespace) do
-    @table_name
+    table_name()
     |> :ets.match({{namespace, :"$1"}, :"$2", :"$3", :_})
     |> Enum.group_by(
       fn [entity_id, _attr, _val] -> entity_id end,
@@ -31,7 +29,7 @@ defmodule Livebook.Storage.Ets do
 
   @impl Livebook.Storage
   def fetch(namespace, entity_id) do
-    @table_name
+    table_name()
     |> :ets.lookup({namespace, entity_id})
     |> case do
       [] ->
@@ -48,7 +46,7 @@ defmodule Livebook.Storage.Ets do
 
   @impl Livebook.Storage
   def fetch_key(namespace, entity_id, key) do
-    @table_name
+    table_name()
     |> :ets.match({{namespace, entity_id}, key, :"$1", :_})
     |> case do
       [[value]] -> {:ok, value}
@@ -56,6 +54,11 @@ defmodule Livebook.Storage.Ets do
     end
   end
 
+  @spec config_file_path() :: Path.t()
+  def config_file_path() do
+    Path.join([Livebook.Config.data_path(), "storage.ets"])
+  end
+
   @impl Livebook.Storage
   def insert(namespace, entity_id, attributes) do
     GenServer.call(__MODULE__, {:insert, namespace, entity_id, attributes})
@@ -73,7 +76,12 @@ defmodule Livebook.Storage.Ets do
 
   @impl GenServer
   def init(_opts) do
-    table = :ets.new(@table_name, [:named_table, :protected, :duplicate_bag])
+    # Make sure that this process does not terminate abruptly
+    # in case it is persisting to disk. terminate/2 is still a no-op.
+    Process.flag(:trap_exit, true)
+
+    table = load_or_create_table()
+    :persistent_term.put(__MODULE__, table)
 
     {:ok, %{table: table}}
   end
@@ -98,6 +106,8 @@ defmodule Livebook.Storage.Ets do
 
     :ets.insert(table, attributes)
 
+    :ok = save_to_file(state)
+
     {:reply, :ok, state}
   end
 
@@ -105,6 +115,29 @@ defmodule Livebook.Storage.Ets do
   def handle_call({:delete, namespace, entity_id}, _from, %{table: table} = state) do
     :ets.delete(table, {namespace, entity_id})
 
+    :ok = save_to_file(state)
+
     {:reply, :ok, state}
   end
+
+  defp table_name(), do: :persistent_term.get(__MODULE__)
+
+  defp load_or_create_table() do
+    config_file_path()
+    |> String.to_charlist()
+    |> :ets.file2tab()
+    |> case do
+      {:ok, tab} ->
+        tab
+
+      {:error, _reason} ->
+        :ets.new(__MODULE__, [:protected, :duplicate_bag])
+    end
+  end
+
+  defp save_to_file(%{table: table}) do
+    file_path = String.to_charlist(config_file_path())
+
+    :ets.tab2file(table, file_path)
+  end
 end

test/livebook/storage/ets_test.exs

@@ -87,4 +87,35 @@ defmodule Livebook.Storage.EtsTest do
       assert [] = Ets.all(:unknown_namespace)
     end
   end
+
+  describe "persistence" do
+    defp read_table_and_lookup(path, entity) do
+      {:ok, tab} =
+        path
+        |> String.to_charlist()
+        |> :ets.file2tab()
+
+      :ets.lookup(tab, {:persistence, entity})
+    end
+
+    test "insert triggers saving to file" do
+      :ok = Ets.insert(:persistence, "insert", key: "val")
+
+      path = Ets.config_file_path()
+      assert File.exists?(path)
+
+      assert [_test] = read_table_and_lookup(path, "insert")
+    end
+
+    test "delete triggers saving to file" do
+      :ok = Ets.insert(:persistence, "delete", key: "val")
+
+      path = Ets.config_file_path()
+      assert File.exists?(path)
+
+      :ok = Ets.delete(:persistence, "delete")
+
+      assert [] = read_table_and_lookup(path, "delete")
+    end
+  end
 end